jax.numpy: move linspace & friends into array_creation

Author: jakevdp

jax/_src/numpy/array_creation.py

@@ -13,18 +13,23 @@
 # limitations under the License.
 
 import types
-from typing import Any
+from functools import partial
+import operator
+from typing import Any, Literal, overload
 
 import numpy as np
 
 import jax
 from jax import lax
+from jax._src.api import jit
 from jax._src import core
 from jax._src import dtypes
+from jax._src.lax import lax as lax_internal
 from jax._src.lib import xla_client as xc
+from jax._src.numpy import ufuncs
 from jax._src.numpy import util
 from jax._src.typing import Array, ArrayLike, DuckTypedArray, DTypeLike
-from jax._src.util import set_module
+from jax._src.util import canonicalize_axis, set_module
 from jax.sharding import Sharding
 
 
@@ -405,3 +410,315 @@ def full_like(a: ArrayLike | DuckTypedArray,
     dtype = dtypes.result_type(a) if dtype is None else dtype
     return jax.device_put(
         util._broadcast_to(jax.numpy.asarray(fill_value, dtype=dtype), shape), device)
+
+@overload
+def linspace(start: ArrayLike, stop: ArrayLike, num: int = 50,
+             endpoint: bool = True, retstep: Literal[False] = False,
+             dtype: DTypeLike | None = None,
+             axis: int = 0,
+             *, device: xc.Device | Sharding | None = None) -> Array: ...
+@overload
+def linspace(start: ArrayLike, stop: ArrayLike, num: int,
+             endpoint: bool, retstep: Literal[True],
+             dtype: DTypeLike | None = None,
+             axis: int = 0,
+             *, device: xc.Device | Sharding | None = None) -> tuple[Array, Array]: ...
+@overload
+def linspace(start: ArrayLike, stop: ArrayLike, num: int = 50,
+             endpoint: bool = True, *, retstep: Literal[True],
+             dtype: DTypeLike | None = None,
+             axis: int = 0,
+             device: xc.Device | Sharding | None = None) -> tuple[Array, Array]: ...
+@overload
+def linspace(start: ArrayLike, stop: ArrayLike, num: int = 50,
+             endpoint: bool = True, retstep: bool = False,
+             dtype: DTypeLike | None = None,
+             axis: int = 0,
+             *, device: xc.Device | Sharding | None = None) -> Array | tuple[Array, Array]: ...
+@export
+def linspace(start: ArrayLike, stop: ArrayLike, num: int = 50,
+             endpoint: bool = True, retstep: bool = False,
+             dtype: DTypeLike | None = None,
+             axis: int = 0,
+             *, device: xc.Device | Sharding | None = None) -> Array | tuple[Array, Array]:
+  """Return evenly-spaced numbers within an interval.
+
+  JAX implementation of :func:`numpy.linspace`.
+
+  Args:
+    start: scalar or array of starting values.
+    stop: scalar or array of stop values.
+    num: number of values to generate. Default: 50.
+    endpoint: if True (default) then include the ``stop`` value in the result.
+      If False, then exclude the ``stop`` value.
+    retstep: If True, then return a ``(result, step)`` tuple, where ``step`` is the
+      interval between adjacent values in ``result``.
+    axis: integer axis along which to generate the linspace. Defaults to zero.
+    device: optional :class:`~jax.Device` or :class:`~jax.sharding.Sharding`
+      to which the created array will be committed.
+
+  Returns:
+    An array ``values``, or a tuple ``(values, step)`` if ``retstep`` is True, where:
+
+    - ``values`` is an array of evenly-spaced values from ``start`` to ``stop``
+    - ``step`` is the interval between adjacent values.
+
+  See also:
+    - :func:`jax.numpy.arange`: Generate ``N`` evenly-spaced values given a starting
+      point and a step
+    - :func:`jax.numpy.logspace`: Generate logarithmically-spaced values.
+    - :func:`jax.numpy.geomspace`: Generate geometrically-spaced values.
+
+  Examples:
+    List of 5 values between 0 and 10:
+
+    >>> jnp.linspace(0, 10, 5)
+    Array([ 0. ,  2.5,  5. ,  7.5, 10. ], dtype=float32)
+
+    List of 8 values between 0 and 10, excluding the endpoint:
+
+    >>> jnp.linspace(0, 10, 8, endpoint=False)
+    Array([0.  , 1.25, 2.5 , 3.75, 5.  , 6.25, 7.5 , 8.75], dtype=float32)
+
+    List of values and the step size between them
+
+    >>> vals, step = jnp.linspace(0, 10, 9, retstep=True)
+    >>> vals
+    Array([ 0.  ,  1.25,  2.5 ,  3.75,  5.  ,  6.25,  7.5 ,  8.75, 10.  ],      dtype=float32)
+    >>> step
+    Array(1.25, dtype=float32)
+
+    Multi-dimensional linspace:
+
+    >>> start = jnp.array([0, 5])
+    >>> stop = jnp.array([5, 10])
+    >>> jnp.linspace(start, stop, 5)
+    Array([[ 0.  ,  5.  ],
+           [ 1.25,  6.25],
+           [ 2.5 ,  7.5 ],
+           [ 3.75,  8.75],
+           [ 5.  , 10.  ]], dtype=float32)
+  """
+  num = core.concrete_dim_or_error(num, "'num' argument of jnp.linspace")
+  axis = core.concrete_or_error(operator.index, axis, "'axis' argument of jnp.linspace")
+  return _linspace(start, stop, num, endpoint, retstep, dtype, axis, device=device)
+
+@partial(jit, static_argnames=('num', 'endpoint', 'retstep', 'dtype', 'axis', 'device'))
+def _linspace(start: ArrayLike, stop: ArrayLike, num: int = 50,
+              endpoint: bool = True, retstep: bool = False,
+              dtype: DTypeLike | None = None,
+              axis: int = 0,
+              *, device: xc.Device | Sharding | None = None) -> Array | tuple[Array, Array]:
+  """Implementation of linspace differentiable in start and stop args."""
+  dtypes.check_user_dtype_supported(dtype, "linspace")
+  if num < 0:
+    raise ValueError(f"Number of samples, {num}, must be non-negative.")
+  start, stop = util.ensure_arraylike("linspace", start, stop)
+
+  if dtype is None:
+    dtype = dtypes.to_inexact_dtype(dtypes.result_type(start, stop))
+  dtype = dtypes.jax_dtype(dtype)
+  computation_dtype = dtypes.to_inexact_dtype(dtype)
+  start = start.astype(computation_dtype)
+  stop = stop.astype(computation_dtype)
+
+  bounds_shape = list(lax.broadcast_shapes(np.shape(start), np.shape(stop)))
+  broadcast_start = util._broadcast_to(start, bounds_shape)
+  broadcast_stop = util._broadcast_to(stop, bounds_shape)
+  axis = len(bounds_shape) + axis + 1 if axis < 0 else axis
+  bounds_shape.insert(axis, 1)
+  div = (num - 1) if endpoint else num
+  if num > 1:
+    delta: Array = lax.convert_element_type(stop - start, computation_dtype) / jax.numpy.array(div, dtype=computation_dtype)
+    iota_shape = [1,] * len(bounds_shape)
+    iota_shape[axis] = div
+    # This approach recovers the endpoints with float32 arithmetic,
+    # but can lead to rounding errors for integer outputs.
+    real_dtype = dtypes.finfo(computation_dtype).dtype
+    step = lax.iota(real_dtype, div).reshape(iota_shape) / jax.numpy.array(div, real_dtype)
+    step = step.astype(computation_dtype)
+    out = (broadcast_start.reshape(bounds_shape) * (1 - step) +
+      broadcast_stop.reshape(bounds_shape) * step)
+
+    if endpoint:
+      out = lax.concatenate([out, lax.expand_dims(broadcast_stop, (axis,))],
+                            canonicalize_axis(axis, out.ndim))
+
+  elif num == 1:
+    delta = jax.numpy.asarray(np.nan if endpoint else stop - start, dtype=computation_dtype)
+    out = broadcast_start.reshape(bounds_shape)
+  else:  # num == 0 degenerate case, match numpy behavior
+    empty_shape = list(lax.broadcast_shapes(np.shape(start), np.shape(stop)))
+    empty_shape.insert(axis, 0)
+    delta = full((), np.nan, computation_dtype)
+    out = empty(empty_shape, dtype)
+
+  if dtypes.issubdtype(dtype, np.integer) and not dtypes.issubdtype(out.dtype, np.integer):
+    out = lax.floor(out)
+
+  sharding = util.canonicalize_device_to_sharding(device)
+  result = lax_internal._convert_element_type(out, dtype, sharding=sharding)
+  return (result, delta) if retstep else result
+
+
+@export
+def logspace(start: ArrayLike, stop: ArrayLike, num: int = 50,
+             endpoint: bool = True, base: ArrayLike = 10.0,
+             dtype: DTypeLike | None = None, axis: int = 0) -> Array:
+  """Generate logarithmically-spaced values.
+
+  JAX implementation of :func:`numpy.logspace`.
+
+  Args:
+    start: scalar or array. Used to specify the start value. The start value is
+      ``base ** start``.
+    stop: scalar or array. Used to specify the stop value. The end value is
+      ``base ** stop``.
+    num: int, optional, default=50. Number of values to generate.
+    endpoint: bool, optional, default=True. If True, then include the ``stop`` value
+      in the result. If False, then exclude the ``stop`` value.
+    base: scalar or array, optional, default=10. Specifies the base of the logarithm.
+    dtype: optional. Specifies the dtype of the output.
+    axis: int, optional, default=0. Axis along which to generate the logspace.
+
+  Returns:
+    An array of logarithm.
+
+  See also:
+    - :func:`jax.numpy.arange`: Generate ``N`` evenly-spaced values given a starting
+      point and a step value.
+    - :func:`jax.numpy.linspace`: Generate evenly-spaced values.
+    - :func:`jax.numpy.geomspace`: Generate geometrically-spaced values.
+
+  Examples:
+    List 5 logarithmically spaced values between 1 (``10 ** 0``) and 100
+    (``10 ** 2``):
+
+    >>> with jnp.printoptions(precision=3, suppress=True):
+    ...   jnp.logspace(0, 2, 5)
+    Array([  1.   ,   3.162,  10.   ,  31.623, 100.   ], dtype=float32)
+
+    List 5 logarithmically-spaced values between 1(``10 ** 0``) and 100
+    (``10 ** 2``), excluding endpoint:
+
+    >>> with jnp.printoptions(precision=3, suppress=True):
+    ...   jnp.logspace(0, 2, 5, endpoint=False)
+    Array([ 1.   ,  2.512,  6.31 , 15.849, 39.811], dtype=float32)
+
+    List 7 logarithmically-spaced values between 1 (``2 ** 0``) and 4 (``2 ** 2``)
+    with base 2:
+
+    >>> with jnp.printoptions(precision=3, suppress=True):
+    ...   jnp.logspace(0, 2, 7, base=2)
+    Array([1.   , 1.26 , 1.587, 2.   , 2.52 , 3.175, 4.   ], dtype=float32)
+
+    Multi-dimensional logspace:
+
+    >>> start = jnp.array([0, 5])
+    >>> stop = jnp.array([5, 0])
+    >>> base = jnp.array([2, 3])
+    >>> with jnp.printoptions(precision=3, suppress=True):
+    ...   jnp.logspace(start, stop, 5, base=base)
+    Array([[  1.   , 243.   ],
+           [  2.378,  61.547],
+           [  5.657,  15.588],
+           [ 13.454,   3.948],
+           [ 32.   ,   1.   ]], dtype=float32)
+  """
+  num = core.concrete_or_error(operator.index, num, "'num' argument of jnp.logspace")
+  axis = core.concrete_or_error(operator.index, axis, "'axis' argument of jnp.logspace")
+  return _logspace(start, stop, num, endpoint, base, dtype, axis)
+
+@partial(jit, static_argnames=('num', 'endpoint', 'dtype', 'axis'))
+def _logspace(start: ArrayLike, stop: ArrayLike, num: int = 50,
+              endpoint: bool = True, base: ArrayLike = 10.0,
+              dtype: DTypeLike | None = None, axis: int = 0) -> Array:
+  """Implementation of logspace differentiable in start and stop args."""
+  dtypes.check_user_dtype_supported(dtype, "logspace")
+  if dtype is None:
+    dtype = dtypes.to_inexact_dtype(dtypes.result_type(start, stop))
+  dtype = dtypes.jax_dtype(dtype)
+  computation_dtype = dtypes.to_inexact_dtype(dtype)
+  start, stop = util.ensure_arraylike("logspace", start, stop)
+  start = start.astype(computation_dtype)
+  stop = stop.astype(computation_dtype)
+  lin = linspace(start, stop, num,
+                 endpoint=endpoint, retstep=False, dtype=None, axis=axis)
+  return lax.convert_element_type(ufuncs.power(base, lin), dtype)
+
+
+@export
+def geomspace(start: ArrayLike, stop: ArrayLike, num: int = 50, endpoint: bool = True,
+              dtype: DTypeLike | None = None, axis: int = 0) -> Array:
+  """Generate geometrically-spaced values.
+
+  JAX implementation of :func:`numpy.geomspace`.
+
+  Args:
+    start: scalar or array. Specifies the starting values.
+    stop: scalar or array. Specifies the stop values.
+    num: int, optional, default=50. Number of values to generate.
+    endpoint: bool, optional, default=True. If True, then include the ``stop`` value
+      in the result. If False, then exclude the ``stop`` value.
+    dtype: optional. Specifies the dtype of the output.
+    axis: int, optional, default=0. Axis along which to generate the geomspace.
+
+  Returns:
+    An array containing the geometrically-spaced values.
+
+  See also:
+    - :func:`jax.numpy.arange`: Generate ``N`` evenly-spaced values given a starting
+      point and a step value.
+    - :func:`jax.numpy.linspace`: Generate evenly-spaced values.
+    - :func:`jax.numpy.logspace`: Generate logarithmically-spaced values.
+
+  Examples:
+    List 5 geometrically-spaced values between 1 and 16:
+
+    >>> with jnp.printoptions(precision=3, suppress=True):
+    ...   jnp.geomspace(1, 16, 5)
+    Array([ 1.,  2.,  4.,  8., 16.], dtype=float32)
+
+    List 4 geomtrically-spaced values between 1 and 16, with ``endpoint=False``:
+
+    >>> with jnp.printoptions(precision=3, suppress=True):
+    ...   jnp.geomspace(1, 16, 4, endpoint=False)
+    Array([1., 2., 4., 8.], dtype=float32)
+
+    Multi-dimensional geomspace:
+
+    >>> start = jnp.array([1, 1000])
+    >>> stop = jnp.array([27, 1])
+    >>> with jnp.printoptions(precision=3, suppress=True):
+    ...   jnp.geomspace(start, stop, 4)
+    Array([[   1., 1000.],
+           [   3.,  100.],
+           [   9.,   10.],
+           [  27.,    1.]], dtype=float32)
+  """
+  num = core.concrete_or_error(operator.index, num, "'num' argument of jnp.geomspace")
+  axis = core.concrete_or_error(operator.index, axis, "'axis' argument of jnp.geomspace")
+  return _geomspace(start, stop, num, endpoint, dtype, axis)
+
+@partial(jit, static_argnames=('num', 'endpoint', 'dtype', 'axis'))
+def _geomspace(start: ArrayLike, stop: ArrayLike, num: int = 50, endpoint: bool = True,
+               dtype: DTypeLike | None = None, axis: int = 0) -> Array:
+  """Implementation of geomspace differentiable in start and stop args."""
+  dtypes.check_user_dtype_supported(dtype, "geomspace")
+  if dtype is None:
+    dtype = dtypes.to_inexact_dtype(dtypes.result_type(start, stop))
+  dtype = dtypes.jax_dtype(dtype)
+  computation_dtype = dtypes.to_inexact_dtype(dtype)
+  start, stop = util.ensure_arraylike("geomspace", start, stop)
+  start = start.astype(computation_dtype)
+  stop = stop.astype(computation_dtype)
+
+  sign = ufuncs.sign(start)
+  res = sign * logspace(ufuncs.log10(start / sign), ufuncs.log10(stop / sign),
+                        num, endpoint=endpoint, base=10.0,
+                        dtype=computation_dtype, axis=0)
+  axis = canonicalize_axis(axis, res.ndim)
+  if axis != 0:
+    # res = moveaxis(res, 0, axis)
+    res = lax.transpose(res, permutation=(*range(1, axis + 1), 0, *range(axis + 1, res.ndim)))
+  return lax.convert_element_type(res, dtype)