pytensor icon indicating copy to clipboard operation
pytensor copied to clipboard

`linspace` and `retstep=True` fails

Open williambdean opened this issue 9 months ago • 1 comments

retstep=True returns a tuple which doesn't work with the astype

https://github.com/pymc-devs/pytensor/blob/c822a8e626f8c4538e954f134fc4a509232bcf62/pytensor/tensor/extra_ops.py?plain=1#L1799-L1808

[ins] In [1]: import pytensor.tensor as pt

[ins] In [2]: pt.linspace(0, 1, retstep=True)
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Cell In[2], line 1
----> 1 pt.linspace(0, 1, retstep=True)

File ~/GitHub/pymc-eco/pytensor/pytensor/tensor/extra_ops.py:1808, in linspace(start, stop, num, endpoint, retstep, dtype, axis, end, steps)
   1797 start, stop = broadcast_arrays(start, stop)
   1799 ls = _linspace_core(
   1800     start=start,
   1801     stop=stop,
   (...)
   1805     axis=axis,
   1806 )
-> 1808 return ls.astype(dtype)

AttributeError: 'tuple' object has no attribute 'astype'

williambdean avatar Mar 17 '25 14:03 williambdean

Will this work?

if retstep:
    ls = tuple([x.astype(dtype) for x in ls])
else:
    ls = ls.astype(dtype)
return ls

Aarsh-Wankar avatar Mar 20 '25 09:03 Aarsh-Wankar

That should work. Or maybe only astype on one of the return values

williambdean avatar May 22 '25 13:05 williambdean