From bd7e35922b1832afd834d6f1a279d94ce07b71bd Mon Sep 17 00:00:00 2001
From: Sanjay Santhanam <51058514+Sanjays2402@users.noreply.github.com>
Date: Tue, 21 Jul 2026 04:34:50 -0700
Subject: [PATCH] Fix linspace retstep dtype handling

Cast only the generated samples when retstep is enabled, preserving the step value returned by the core implementation. Add regression coverage across inferred, integer, and floatX dtypes.
---
 pytensor/tensor/extra_ops.py   |  4 ++++
 tests/tensor/test_extra_ops.py | 16 ++++++++++++++++
 2 files changed, 20 insertions(+)

diff --git a/pytensor/tensor/extra_ops.py b/pytensor/tensor/extra_ops.py
index 21166ecc5b..543926920a 100644
--- a/pytensor/tensor/extra_ops.py
+++ b/pytensor/tensor/extra_ops.py
@@ -1710,6 +1710,10 @@ def linspace(
         axis=axis,
     )
 
+    if retstep:
+        samples, step = ls
+        return samples.astype(dtype), step
+
     return ls.astype(dtype)
 
 
diff --git a/tests/tensor/test_extra_ops.py b/tests/tensor/test_extra_ops.py
index 88de5cbb1a..0ab310c918 100644
--- a/tests/tensor/test_extra_ops.py
+++ b/tests/tensor/test_extra_ops.py
@@ -1356,6 +1356,22 @@ def test_space_ops(op, dtype, start, stop, num_samples, endpoint, axis):
     )
 
 
+@pytest.mark.parametrize("dtype", [None, "int64", "floatX"])
+def test_linspace_retstep(dtype):
+    dtype = dtype if dtype != "floatX" else config.floatX
+    samples, step = pt.linspace(0, 1, num=3, retstep=True, dtype=dtype)
+
+    result = function(inputs=[], outputs=[samples, step], mode="FAST_COMPILE")()
+    assert result is not None
+    actual_samples, actual_step = result
+    expected_samples, expected_step = np.linspace(
+        0, 1, num=3, retstep=True, dtype=dtype
+    )
+
+    np.testing.assert_allclose(actual_samples, expected_samples)
+    np.testing.assert_allclose(actual_step, expected_step)
+
+
 def test_concat_with_broadcast():
     rng = np.random.default_rng()
     a = pt.tensor("a", shape=(1, 3, 5))