Refactor linspace, logspace, and geomspace to match numpy implementation

Author: jessegrabowski

pytensor/tensor/extra_ops.py

@@ -1,3 +1,4 @@
+import warnings
 from collections.abc import Collection, Iterable
 
 import numpy as np
@@ -20,14 +21,14 @@
 from pytensor.raise_op import Assert
 from pytensor.scalar import int32 as int_t
 from pytensor.scalar import upcast
-from pytensor.tensor import as_tensor_variable
+from pytensor.tensor import TensorLike, as_tensor_variable
 from pytensor.tensor import basic as ptb
 from pytensor.tensor.basic import alloc, second
 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 eq as pt_eq
-from pytensor.tensor.math import ge, lt, maximum, minimum, prod, switch
+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
@@ -1583,27 +1584,278 @@ def broadcast_shape_iter(
     return tuple(result_dims)
 
 
-def geomspace(start, end, steps, base=10.0):
-    from pytensor.tensor.math import log
+def _check_deprecated_inputs(stop, end, num, steps):
+    if end is not None:
+        warnings.warn(
+            "The 'end' parameter is deprecated and will be removed in a future version. Use 'stop' instead.",
+            DeprecationWarning,
+        )
+        stop = end
+    if steps is not None:
+        warnings.warn(
+            "The 'steps' parameter is deprecated and will be removed in a future version. Use 'num' instead.",
+            DeprecationWarning,
+        )
+        num = steps
+
+    return stop, num
+
+
+def _linspace_core(
+    start, stop, num, endpoint=True, retstep=False, dtype=None, 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)
+
+    if axis != 0:
+        samples = ptb.moveaxis(samples, 0, axis)
+
+    if retstep:
+        return samples, step
+
+    return samples
+
+
+def _broadcast_inputs_and_dtypes(*args, dtype=None):
+    args = map(ptb.as_tensor_variable, args)
+    args = broadcast_arrays(*args)
+
+    if dtype is None:
+        dtype = pytensor.config.floatX
+
+    return args, dtype
+
+
+def _broadcast_base_with_inputs(start, stop, base, dtype, axis):
+    """
+    Broadcast the base tensor with the start and stop tensors if base is not a scalar. This is important because it
+    may change how the axis argument is interpreted in the final output.
+
+    Parameters
+    ----------
+    start
+    stop
+    base
+    dtype
+    axis
+
+    Returns
+    -------
+
+    """
+    base = ptb.as_tensor_variable(base, dtype=dtype)
+    if base.ndim > 0:
+        ndmax = len(broadcast_shape(start, stop, base))
+        start, stop, base = (
+            ptb.shape_padleft(a, ndmax - a.ndim) for a in (start, stop, base)
+        )
+        base = ptb.expand_dims(base, axis=(axis,))
+
+    return start, stop, base
+
+
+def linspace(
+    start: TensorLike,
+    stop: TensorLike,
+    num: TensorLike = 50,
+    endpoint: bool = True,
+    retstep: bool = False,
+    dtype: str | None = None,
+    axis: int = 0,
+    end: TensorLike | None = None,
+    steps: TensorLike | None = None,
+) -> TensorVariable | tuple[TensorVariable, TensorVariable]:
+    """
+    Return evenly spaced numbers over a specified interval.
+
+    Returns `num` evenly spaced samples, calculated over the interval [`start`, `stop`].
+
+    The endpoint of the interval can optionally be excluded.
+
+    Parameters
+    ----------
+    start: int, float, or TensorVariable
+        The starting value of the sequence.
+
+    stop: int, float or TensorVariable
+        The end value of the sequence, unless `endpoint` is set to False.
+        In that case, the sequence consists of all but the last of `num + 1` evenly spaced samples, such that `stop` is excluded.
+
+    num: int
+        Number of samples to generate. Must be non-negative.
+
+    endpoint: bool
+        Whether to include the endpoint in the range.
+
+    retstep: bool
+        If true, returns both the samples and an array of steps between samples.
+
+    dtype: str, optional
+        dtype of the output tensor(s). If None, the dtype is inferred from that of the values provided to the `start`
+        and `end` arguments.
 
-    start = ptb.as_tensor_variable(start)
-    end = ptb.as_tensor_variable(end)
-    return base ** linspace(log(start) / log(base), log(end) / log(base), steps)
+    axis: int
+        Axis along which to generate samples. Ignored if both `start` and `end` have dimension 0. By default, axis=0
+        will insert the samples on a new left-most dimension. To insert samples on a right-most dimension, use axis=-1.
 
+    end:  int, float or TensorVariable
+        .. warning::
+            The "end" parameter is deprecated and will be removed in a future version. Use "stop" instead.
+        The end value of the sequence, unless `endpoint` is set to False.
+        In that case, the sequence consists of all but the last of `num + 1` evenly spaced samples, such that `end` is
+        excluded.
 
-def logspace(start, end, steps, base=10.0):
-    start = ptb.as_tensor_variable(start)
-    end = ptb.as_tensor_variable(end)
-    return base ** linspace(start, end, steps)
+    steps: float, int, or TensorVariable
+        .. warning::
+            The "steps" parameter is deprecated and will be removed in a future version. Use "num" instead.
 
+        Number of samples to generate. Must be non-negative
+
+    Returns
+    -------
+    samples: TensorVariable
+        Tensor containing `num` evenly-spaced values between [start, stop]. The range is inclusive if `endpoint` is True.
+
+    step: TensorVariable
+        Tensor containing the spacing between samples. Only returned if `retstep` is True.
+    """
+    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)
+
+
+def geomspace(
+    start: TensorLike,
+    stop: TensorLike,
+    num: int = 50,
+    base: float = 10.0,
+    endpoint: bool = True,
+    dtype: str | None = None,
+    axis: int = 0,
+    end: TensorLike | None = None,
+    steps: TensorLike | None = None,
+) -> TensorVariable:
+    """
+    Return numbers spaced evenly on a log scale (a geometric progression).
+
+    Parameters
+    ----------
+    Returns `num` evenly spaced samples, calculated over the interval [`start`, `stop`].
+
+    The endpoint of the interval can optionally be excluded.
+
+    Parameters
+    ----------
+    start: int, float, or TensorVariable
+        The starting value of the sequence.
+
+    stop: int, float or TensorVariable
+        The end value of the sequence, unless `endpoint` is set to False.
+        In that case, the sequence consists of all but the last of `num + 1` evenly spaced samples, such that `stop` is excluded.
+
+    num: int
+        Number of samples to generate. Must be non-negative.
+
+    base: float
+        The base of the log space. The step size between the elements in ln(samples) / ln(base)
+        (or log_base(samples)) is uniform.
+
+    endpoint: bool
+        Whether to include the endpoint in the range.
+
+    dtype: str, optional
+        dtype of the output tensor(s). If None, the dtype is inferred from that of the values provided to the `start`
+        and `end` arguments.
+
+    axis: int
+        Axis along which to generate samples. Ignored if both `start` and `end` have dimension 0. By default, axis=0
+        will insert the samples on a new left-most dimension. To insert samples on a right-most dimension, use axis=-1.
+
+    end:  int, float or TensorVariable
+        .. warning::
+            The "end" parameter is deprecated and will be removed in a future version. Use "stop" instead.
+        The end value of the sequence, unless `endpoint` is set to False.
+        In that case, the sequence consists of all but the last of `num + 1` evenly spaced samples, such that `end` is
+        excluded.
+
+    steps: float, int, or TensorVariable
+        .. warning::
+            The "steps" parameter is deprecated and will be removed in a future version. Use "num" instead.
+
+        Number of samples to generate. Must be non-negative
+
+    Returns
+    -------
+    samples: TensorVariable
+        Tensor containing `num` evenly-spaced values between [start, stop]. The range is inclusive if `endpoint` is True.
+    """
+    stop, num = _check_deprecated_inputs(stop, end, num, steps)
+    (start, stop), dtype = _broadcast_inputs_and_dtypes(start, stop, dtype=dtype)
+    start, stop, base = _broadcast_base_with_inputs(start, stop, base, dtype, axis)
+
+    out_sign = sign(start)
+    log_start, log_stop = (
+        log(start * out_sign) / log(base),
+        log(stop * out_sign) / log(base),
+    )
+    result = _linspace_core(
+        start=log_start,
+        stop=log_stop,
+        num=num,
+        endpoint=endpoint,
+        dtype=dtype,
+        axis=0,
+        retstep=False,
+    )
+    result = base**result
+
+    if num > 0:
+        set_subtensor(result[0, ...], start, inplace=True)
+        if num > 1 and endpoint:
+            set_subtensor(result[-1, ...], stop, inplace=True)
+
+    result = result * out_sign
+
+    if axis != 0:
+        result = ptb.moveaxis(result, 0, axis)
+
+    return result
+
+
+def logspace(
+    start: TensorLike,
+    stop: TensorLike,
+    num: int = 50,
+    base: float = 10.0,
+    endpoint: bool = True,
+    dtype: str | None = None,
+    axis: int = 0,
+    end: TensorLike | None = None,
+    steps: TensorLike | None = None,
+) -> TensorVariable:
+    stop, num = _check_deprecated_inputs(stop, end, num, steps)
+    (start, stop), type = _broadcast_inputs_and_dtypes(start, stop, dtype=dtype)
+    start, stop, base = _broadcast_base_with_inputs(start, stop, base, dtype, axis)
+
+    ls = _linspace_core(
+        start=start,
+        stop=stop,
+        num=num,
+        endpoint=endpoint,
+        dtype=dtype,
+        axis=axis,
+        retstep=False,
+    )
 
-def linspace(start, end, steps):
-    start = ptb.as_tensor_variable(start)
-    end = ptb.as_tensor_variable(end)
-    arr = ptb.arange(steps)
-    arr = ptb.shape_padright(arr, max(start.ndim, end.ndim))
-    multiplier = (end - start) / (steps - 1)
-    return start + arr * multiplier
+    return base**ls
 
 
 def broadcast_to(

tests/tensor/test_extra_ops.py

@@ -35,9 +35,6 @@
     diff,
     fill_diagonal,
     fill_diagonal_offset,
-    geomspace,
-    linspace,
-    logspace,
     ravel_multi_index,
     repeat,
     searchsorted,
@@ -1281,25 +1278,27 @@ def test_broadcast_arrays():
 
 
 @pytest.mark.parametrize(
-    "start, stop, num_samples",
+    "op",
+    ["linspace", "logspace", "geomspace"],
+    ids=["linspace", "logspace", "geomspace"],
+)
+@pytest.mark.parametrize(
+    "start, stop, num_samples, endpoint, axis",
     [
-        (1, 10, 50),
-        (np.array([5, 6]), np.array([[10, 10], [10, 10]]), 25),
-        (1, np.array([5, 6]), 30),
+        (1, 10, 50, 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),
+        (1, np.array([5, 6]), 30, True, 0),
+        (1, np.array([5, 6]), 30, False, -1),
     ],
 )
-def test_space_ops(start, stop, num_samples):
-    z = linspace(start, stop, num_samples)
-    pytensor_res = function(inputs=[], outputs=z)()
-    numpy_res = np.linspace(start, stop, num=num_samples)
-    assert np.allclose(pytensor_res, numpy_res)
-
-    z = logspace(start, stop, num_samples)
-    pytensor_res = function(inputs=[], outputs=z)()
-    numpy_res = np.logspace(start, stop, num=num_samples)
-    assert np.allclose(pytensor_res, numpy_res)
-
-    z = geomspace(start, stop, num_samples)
-    pytensor_res = function(inputs=[], outputs=z)()
-    numpy_res = np.geomspace(start, stop, num=num_samples)
-    assert np.allclose(pytensor_res, numpy_res)
+def test_space_ops(op, start, stop, num_samples, endpoint, axis):
+    pt_func = getattr(pt, op)
+    np_func = getattr(np, op)
+    z = pt_func(start, stop, num_samples, endpoint=endpoint, axis=axis)
+
+    numpy_res = np_func(start, stop, num=num_samples, endpoint=endpoint, axis=axis)
+    pytensor_res = function(inputs=[], outputs=z, mode="FAST_COMPILE")()
+
+    np.testing.assert_allclose(pytensor_res, numpy_res, atol=1e-6, rtol=1e-6)