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))