Handle corner case when steps = 0

Author: jessegrabowski

pytensor/tensor/extra_ops.py

@@ -27,8 +27,19 @@
 from pytensor.tensor.exceptions import NotScalarConstantError
 from pytensor.tensor.math import abs as pt_abs
 from pytensor.tensor.math import all as pt_all
+from pytensor.tensor.math import (
+    bitwise_and,
+    ge,
+    gt,
+    log,
+    lt,
+    maximum,
+    minimum,
+    prod,
+    sign,
+    switch,
+)
 from pytensor.tensor.math import eq as pt_eq
-from pytensor.tensor.math import ge, log, lt, maximum, minimum, prod, sign, switch
 from pytensor.tensor.math import max as pt_max
 from pytensor.tensor.math import sum as pt_sum
 from pytensor.tensor.shape import specify_broadcastable
@@ -1602,17 +1613,25 @@ def _check_deprecated_inputs(stop, end, num, steps):
 
 
 def _linspace_core(
-    start, stop, num, endpoint=True, retstep=False, dtype=None, axis=0
+    start: TensorVariable,
+    stop: TensorVariable,
+    num: int,
+    dtype: str,
+    endpoint=True,
+    retstep=False,
+    axis=0,
 ) -> TensorVariable | tuple[TensorVariable, TensorVariable]:
     div = (num - 1) if endpoint else num
     delta = (stop - start).astype(dtype)
-    step = delta / div
-
     samples = ptb.arange(0, num, dtype=dtype).reshape((-1,) + (1,) * delta.ndim)
-    samples = samples * step + start
 
-    if endpoint and num > 1:
-        set_subtensor(samples[-1, ...], stop, inplace=True)
+    step = switch(gt(div, 0), delta / div, np.nan)
+    samples = switch(gt(div, 0), samples * delta / div + start, samples * delta + start)
+    samples = switch(
+        bitwise_and(gt(num, 1), np.asarray(endpoint)),
+        set_subtensor(samples[-1, ...], stop),
+        samples,
+    )
 
     if axis != 0:
         samples = ptb.moveaxis(samples, 0, axis)
@@ -1729,7 +1748,15 @@ def linspace(
     end, num = _check_deprecated_inputs(stop, end, num, steps)
     (start, stop), type = _broadcast_inputs_and_dtypes(start, stop, dtype=dtype)
 
-    return _linspace_core(start, stop, num, endpoint, retstep, dtype, axis)
+    return _linspace_core(
+        start=start,
+        stop=stop,
+        num=num,
+        dtype=dtype,
+        endpoint=endpoint,
+        retstep=retstep,
+        axis=axis,
+    )
 
 
 def geomspace(

tests/tensor/test_extra_ops.py

@@ -1286,6 +1286,7 @@ def test_broadcast_arrays():
     "start, stop, num_samples, endpoint, axis",
     [
         (1, 10, 50, True, 0),
+        (1, 10, 1, True, 0),
         (np.array([5, 6]), np.array([[10, 10], [10, 10]]), 25, True, 0),
         (np.array([5, 6]), np.array([[10, 10], [10, 10]]), 25, True, 1),
         (np.array([5, 6]), np.array([[10, 10], [10, 10]]), 25, False, -1),