Fix linspace retstep dtype handling

[Sanjays2402]

Description

linspace(..., retstep=True) now casts only the generated samples instead of calling astype on the returned tuple. The step keeps NumPy-compatible floating-point precision.

Related Issue

• Closes linspace and retstep=True fails #1297
• Related to #

Checklist

• Checked that the pre-commit linting/style checks pass
• Included tests that prove the fix is effective or that the new feature works
• Added necessary documentation (docstrings and/or example notebooks)
• If you are a pro: each commit corresponds to a relevant logical change

Type of change

• New feature / enhancement
• Bug fix
• Documentation
• Maintenance
• Other (please specify):

[ricardoV94]

the other places that make use of this helper don't share the same problem?

[Sanjays2402]

Good question — I checked, and no, the other callers are unaffected.

_linspace_core has three callers, and the two besides linspace pass retstep=False explicitly and hardcode it (not exposed as a kwarg on their own signatures):

• geomspace (extra_ops.py:1783) — retstep=False, then result = base**result and .astype(dtype)
• logspace (extra_ops.py:1858) — retstep=False, then return (base**ls).astype(dtype)

Since _linspace_core only returns the (samples, step) tuple when retstep=True, both of them always get a bare tensor back, so the .astype(dtype) on the return value can never hit a tuple. linspace was the only one that forwarded a user-controllable retstep through and then unconditionally called .astype on the result.

Confirmed on the branch:

geomspace dtype=int64 -> type=TensorVariable out_dtype=int64
geomspace dtype=float64 -> type=TensorVariable out_dtype=float64
logspace dtype=int64 -> type=TensorVariable out_dtype=int64
logspace dtype=float64 -> type=TensorVariable out_dtype=float64
linspace retstep int64: samples.dtype= int64  step.dtype= float64

One thing worth flagging from that last line: I left step un-cast, matching numpy — np.linspace(0, 1, num=3, retstep=True, dtype='int64') also returns an int64 samples array with a float64 step, since the step is the true spacing rather than a sample value. That's what the added test asserts against.

$ python -m pytest tests/tensor/test_extra_ops.py -k "linspace or space_ops or geomspace or logspace" -q
66 passed, 257 deselected, 14 warnings in 17.03s