Core ops

anytensor.core

Ordinary tensor ops via array-api-compat (stable, input-adaptive).

promote_options

promote_options(*, copy=False, fallback='copy')

Temporarily set NumPy-upcast defaults (copy / non-ref fallback).

Example::

with at.promote_options(copy=True):
    y = at.maximum(torch_x, numpy_y)  # safe if numpy_y will be mutated

Source code in anytensor/core.py

@contextmanager
def promote_options(*, copy: bool = False, fallback: Fallback = "copy") -> Iterator[None]:
    """Temporarily set NumPy-upcast defaults (``copy`` / non-ref ``fallback``).

    Example::

        with at.promote_options(copy=True):
            y = at.maximum(torch_x, numpy_y)  # safe if numpy_y will be mutated
    """
    t_copy = _promote_copy.set(copy)
    t_fb = _promote_fallback.set(fallback)
    try:
        yield
    finally:
        _promote_copy.reset(t_copy)
        _promote_fallback.reset(t_fb)

align_arrays

align_arrays(*arrays, copy=None, fallback=None)

Align operands on one namespace; upcast scalars/NumPy to non-NumPy peers.

None entries are preserved (useful when optional operands share a call).

Source code in anytensor/core.py

def align_arrays(
    *arrays: Any,
    copy: Optional[bool] = None,
    fallback: Optional[Fallback] = None,
) -> tuple[Any, ...]:
    """Align operands on one namespace; upcast scalars/NumPy to non-NumPy peers.

    ``None`` entries are preserved (useful when optional operands share a call).
    """
    xp = _xp(*arrays)
    return tuple(
        None if a is None else _asarray(xp, a, copy=copy, fallback=fallback)
        for a in arrays
    )

as_array_result

as_array_result(fn)

Decorator: promote scalar returns to 0-d arrays on the call's namespace.

Reductions (and any op) may hand back np.float64 / Python scalars; callers need a real array (.shape, .dtype, methods).

Source code in anytensor/core.py

def as_array_result(fn: Callable[..., ShapedArray]) -> Callable[..., ShapedArray]:
    """Decorator: promote scalar returns to 0-d arrays on the call's namespace.

    Reductions (and any op) may hand back ``np.float64`` / Python scalars;
    callers need a real array (``.shape``, ``.dtype``, methods).
    """

    @functools.wraps(fn)
    def wrapper(*args, **kwargs):
        out = fn(*args, **kwargs)
        if not _is_scalar(out):
            return out
        xp = _xp(*args, *kwargs.values())
        return xp.asarray(out)

    return wrapper

promote

promote(*, copy=None, fallback=None, **roles)

Decorator: namespace upcast + per-operand dtype policy.

Pass keyword roles for each parameter::

@promote(x="data", y="data")
def maximum(x: ArrayT, y: ArrayT) -> ArrayT: ...

@promote(x="data", indices="index")
def take(x, indices, axis=0): ...

@promote(condition="mask", x="data", y="data")
def where(condition: ArrayT, x: ArrayT, y: ArrayT) -> ArrayT: ...

@promote(x="data", segment_ids="index", num_segments="shape")
def segment_sum(x, segment_ids, num_segments): ...

data operands share Array API result_type (so a NumPy int beside a float tensor becomes float). index stays integral (width is backend-local: Torch may cast to int64; JAX/TF often keep int32). mask becomes bool. shape is a size dim (Python int / symbolic / 0-d integral tensor) and is not promoted to a 0-d array.

copy / fallback control NumPy→framework buffer sharing (see :func:promote_options).

Source code in anytensor/core.py

def promote(
    *,
    copy: Optional[bool] = None,
    fallback: Optional[Fallback] = None,
    **roles: OperandKind,
) -> Callable:
    """Decorator: namespace upcast + per-operand dtype policy.

    Pass keyword roles for each parameter::

        @promote(x="data", y="data")
        def maximum(x: ArrayT, y: ArrayT) -> ArrayT: ...

        @promote(x="data", indices="index")
        def take(x, indices, axis=0): ...

        @promote(condition="mask", x="data", y="data")
        def where(condition: ArrayT, x: ArrayT, y: ArrayT) -> ArrayT: ...

        @promote(x="data", segment_ids="index", num_segments="shape")
        def segment_sum(x, segment_ids, num_segments): ...

    ``data`` operands share Array API ``result_type`` (so a NumPy int beside a
    float tensor becomes float). ``index`` stays integral (width is
    backend-local: Torch may cast to int64; JAX/TF often keep int32). ``mask``
    becomes bool. ``shape`` is a size dim (Python int / symbolic / 0-d integral
    tensor) and is **not** promoted to a 0-d array.

    ``copy`` / ``fallback`` control NumPy→framework buffer sharing (see
    :func:`promote_options`).
    """
    if roles:  # pragma: no branch
        def decorator(fn: Callable) -> Callable:
            sig = inspect.signature(fn)

            @functools.wraps(fn)
            def wrapper(*args, **kwargs):
                bound = sig.bind(*args, **kwargs)
                bound.apply_defaults()
                names = [n for n in roles if n in bound.arguments]
                # Namespace from data/index/mask only — pure Python shape ints stay host.
                ns_values = [
                    bound.arguments[n]
                    for n in names
                    if roles[n] != "shape" or not _is_scalar(bound.arguments[n])
                ]
                xp = _xp(*ns_values) if ns_values else _xp()
                converted = {}
                for n in names:
                    v = bound.arguments[n]
                    if roles[n] == "shape":
                        if v is None:
                            continue
                        converted[n] = _normalize_shape_dim(v)
                    else:
                        converted[n] = _asarray(xp, v, copy=copy, fallback=fallback)
                _apply_dtype_roles(xp, converted, roles)
                for n, v in converted.items():
                    bound.arguments[n] = v
                return fn(*bound.args, **bound.kwargs)

            return wrapper

        return decorator
    raise TypeError("promote() requires at least one name=kind role")

promote_scalars

promote_scalars(*names, copy=None, fallback=None)

Upcast named operands as data (namespace + result_type).

Prefer :func:promote when some args are indices/masks. Kept as a short form of @promote(x="data", y="data").

Source code in anytensor/core.py

def promote_scalars(
    *names: str,
    copy: Optional[bool] = None,
    fallback: Optional[Fallback] = None,
) -> Callable:
    """Upcast named operands as ``data`` (namespace + ``result_type``).

    Prefer :func:`promote` when some args are indices/masks. Kept as a short
    form of ``@promote(x="data", y="data")``.
    """
    if not names:
        raise TypeError("promote_scalars() requires at least one parameter name")
    return promote(**{n: "data" for n in names}, copy=copy, fallback=fallback)

exp

exp(x)

Element-wise exponential.

Source code in anytensor/core.py

@as_array_result
def exp(x: ShapedArray) -> ShapedArray:
    """Element-wise exponential."""
    return array_namespace(x).exp(x)

log

log(x)

Element-wise natural logarithm.

Source code in anytensor/core.py

@as_array_result
def log(x: ShapedArray) -> ShapedArray:
    """Element-wise natural logarithm."""
    return array_namespace(x).log(x)

sum

sum(x, axes=None)

Sum over axes; full reduce returns a 0-d array (not a scalar).

Source code in anytensor/core.py

@as_array_result
def sum(x: ShapedArray, axes: Axes = None) -> ShapedArray:
    """Sum over axes; full reduce returns a **0-d array** (not a scalar)."""
    xp = array_namespace(x)
    with _ignore_fp_invalid(xp):
        return xp.sum(x, axis=_axis(axes))

min

min(x, axes=None)

Minimum over axes; full reduce returns a 0-d array (not a scalar).

Notes

Length-0 reductions are framework-defined (often error) — prefer nonempty. Under TF XLA (jit_compile=True), NaN inputs may yield ±inf instead of NaN; eager TF / NumPy / JAX usually keep NaN.

Source code in anytensor/core.py

@as_array_result
def min(x: ShapedArray, axes: Axes = None) -> ShapedArray:
    """Minimum over axes; full reduce returns a **0-d array** (not a scalar).

    Notes:
        Length-0 reductions are framework-defined (often error) — prefer
        nonempty. Under TF XLA (``jit_compile=True``), NaN inputs may yield
        ``±inf`` instead of NaN; eager TF / NumPy / JAX usually keep NaN.
    """
    xp = array_namespace(x)
    with _ignore_fp_invalid(xp):
        return xp.min(x, axis=_axis(axes))

max

max(x, axes=None)

Maximum over axes; full reduce returns a 0-d array (not a scalar).

Notes

Length-0 reductions are framework-defined (often error) — prefer nonempty. Under TF XLA, NaN inputs may yield ±inf instead of NaN.

Source code in anytensor/core.py

@as_array_result
def max(x: ShapedArray, axes: Axes = None) -> ShapedArray:
    """Maximum over axes; full reduce returns a **0-d array** (not a scalar).

    Notes:
        Length-0 reductions are framework-defined (often error) — prefer
        nonempty. Under TF XLA, NaN inputs may yield ``±inf`` instead of NaN.
    """
    xp = array_namespace(x)
    with _ignore_fp_invalid(xp):
        return xp.max(x, axis=_axis(axes))

mean

mean(x, axes=None)

Mean over axes; full reduce returns a 0-d array (not a scalar).

Notes

Empty full-reduce (x.size == 0) returns a 0-d NaN on the backend dtype to avoid NumPy's Mean of empty slice warning path.

Source code in anytensor/core.py

@as_array_result
def mean(x: ShapedArray, axes: Axes = None) -> ShapedArray:
    """Mean over axes; full reduce returns a **0-d array** (not a scalar).

    Notes:
        Empty full-reduce (``x.size == 0``) returns a 0-d NaN on the backend
        dtype to avoid NumPy's ``Mean of empty slice`` warning path.
    """
    xp = array_namespace(x)
    # NumPy emits RuntimeWarning "Mean of empty slice" (not via errstate).
    size = getattr(x, "size", None)
    if axes is None and size == 0:
        dtype = getattr(x, "dtype", None)
        return xp.asarray(float("nan"), dtype=dtype) if dtype is not None else xp.asarray(float("nan"))
    with _ignore_fp_invalid(xp):
        return xp.mean(x, axis=_axis(axes))

prod

prod(x, axes=None)

Product over axes; full reduce returns a 0-d array (not a scalar).

Notes

inf * a subnormal or float32-min value may be inf (NumPy / eager TF) or nan (JAX / TF XLA) when the tiny flushes to 0. Not standardized — keep finite samples away from the underflow edge if you need portable results.

Source code in anytensor/core.py

@as_array_result
def prod(x: ShapedArray, axes: Axes = None) -> ShapedArray:
    """Product over axes; full reduce returns a **0-d array** (not a scalar).

    Notes:
        ``inf *`` a subnormal or float32-min value may be ``inf`` (NumPy /
        eager TF) or ``nan`` (JAX / TF XLA) when the tiny flushes to 0. Not
        standardized — keep finite samples away from the underflow edge if
        you need portable results.
    """
    xp = array_namespace(x)
    with _ignore_fp_invalid(xp):
        return xp.prod(x, axis=_axis(axes))

cumsum

cumsum(x, axis=0)

Cumulative sum along axis (default 0, never flatten).

Source code in anytensor/core.py

@as_array_result
def cumsum(x: ShapedArray, axis: int = 0) -> ShapedArray:
    """Cumulative sum along ``axis`` (default ``0``, never flatten)."""
    xp = array_namespace(x)
    with _ignore_fp_invalid(xp):
        return xp.cumulative_sum(x, axis=axis)

shape

shape(x)

Return the shape of x as a tuple.

Under tracing (tf.function, jax.jit, torch.compile), unknown dims are backend size tensors / symbols rather than Python None, so callers can build matching outputs under jit/compile.

Notes

Prefer this over raw x.shape when feeding zeros / full / *_like under polymorphic TF graphs — tnp.zeros errors on TensorShape(None,) after retracing.

Source code in anytensor/core.py

def shape(x: ShapedArray) -> ShapeLike:
    """Return the shape of ``x`` as a tuple.

    Under tracing (``tf.function``, ``jax.jit``, ``torch.compile``), unknown
    dims are backend size tensors / symbols rather than Python ``None``, so
    callers can build matching outputs under jit/compile.

    Notes:
        Prefer this over raw ``x.shape`` when feeding ``zeros`` / ``full`` /
        ``*_like`` under polymorphic TF graphs — ``tnp.zeros`` errors on
        ``TensorShape(None,)`` after retracing.
    """
    from .backends import HashableTuple, get_backend

    s = get_backend(x).shape(x)
    if isinstance(s, HashableTuple):  # pragma: no cover - TF graph shape wrapper
        return tuple(s)
    return tuple(s)

take

take(x, indices, axis=0)

Take elements from x along axis (default 0).

Source code in anytensor/core.py

@as_array_result
@promote(x="data", indices="index")
def take(x: ShapedArray, indices: IntArray, axis: int = 0) -> ShapedArray:
    """Take elements from ``x`` along ``axis`` (default ``0``)."""
    return array_namespace(x, indices).take(x, indices, axis=axis)

reshape

reshape(x, shape)

Reshape x to shape.

Source code in anytensor/core.py

@as_array_result
def reshape(x: ShapedArray, shape: ShapeLike) -> ShapedArray:
    """Reshape ``x`` to ``shape``."""
    return array_namespace(x).reshape(x, shape)

transpose

transpose(x, axes=None)

Permute axes of x.

Source code in anytensor/core.py

@as_array_result
def transpose(x: ShapedArray, axes: Optional[Sequence[int]] = None) -> ShapedArray:
    """Permute axes of ``x``."""
    xp = array_namespace(x)
    if axes is None:
        return xp.permute_dims(x, axes=tuple(range(x.ndim - 1, -1, -1)))
    return xp.permute_dims(x, axes=tuple(axes))

concatenate

concatenate(arrays, axis=0)

Concatenate a sequence of arrays along axis.

Source code in anytensor/core.py

@as_array_result
def concatenate(arrays: Sequence[ShapedArray], axis: int = 0) -> ShapedArray:
    """Concatenate a sequence of arrays along ``axis``."""
    return array_namespace(*arrays).concat(arrays, axis=axis)

split

split(x, indices_or_sections, axis=0)

Split x along axis (NumPy split semantics).

Parameters:

Name	Type	Description	Default
x	ShapedArray	Array to split.	required
indices_or_sections	Any	Either an int (equal-sized parts) or a sequence of cut indices along axis (as in numpy.split).	required
axis	int	Axis along which to split.	0

Returns:

Type	Description
list	A list of array chunks (may include empty leading-axis pieces).

Notes

Dispatches via :func:anytensor.backends.get_backend so Torch uses tensor_split (index cuts) rather than size-based torch.split.

Source code in anytensor/core.py

def split(x: ShapedArray, indices_or_sections: Any, axis: int = 0) -> list:
    """Split ``x`` along ``axis`` (NumPy ``split`` semantics).

    Args:
        x: Array to split.
        indices_or_sections: Either an ``int`` (equal-sized parts) or a sequence
            of cut indices along ``axis`` (as in ``numpy.split``).
        axis: Axis along which to split.

    Returns:
        A ``list`` of array chunks (may include empty leading-axis pieces).

    Notes:
        Dispatches via :func:`anytensor.backends.get_backend` so Torch uses
        ``tensor_split`` (index cuts) rather than size-based ``torch.split``.
    """
    from .backends import get_backend

    return get_backend(x).split(x, indices_or_sections, axis=axis)

stack

stack(arrays, axis=0)

Stack a sequence of arrays along a new axis.

Source code in anytensor/core.py

@as_array_result
def stack(arrays: Sequence[ShapedArray], axis: int = 0) -> ShapedArray:
    """Stack a sequence of arrays along a new ``axis``."""
    return array_namespace(*arrays).stack(arrays, axis=axis)

maximum

maximum(x, y)

Element-wise maximum. Scalars/NumPy upcast; dtypes via result_type.

Notes

Under TF XLA (jit_compile=True), NaN inputs may yield ±inf instead of NaN; eager TF / NumPy / JAX usually keep NaN. Not standardized across backends.

Source code in anytensor/core.py

@as_array_result
@promote(x="data", y="data")
def maximum(x: ShapedArray, y: ShapedArray) -> ShapedArray:
    """Element-wise maximum. Scalars/NumPy upcast; dtypes via ``result_type``.

    Notes:
        Under TF XLA (``jit_compile=True``), NaN inputs may yield ``±inf``
        instead of NaN; eager TF / NumPy / JAX usually keep NaN. Not
        standardized across backends.
    """
    return array_namespace(x, y).maximum(x, y)

minimum

minimum(x, y)

Element-wise minimum. Scalars/NumPy upcast; dtypes via result_type.

Notes

Under TF XLA, NaN inputs may yield ±inf instead of NaN. Not standardized across backends.

Source code in anytensor/core.py

@as_array_result
@promote(x="data", y="data")
def minimum(x: ShapedArray, y: ShapedArray) -> ShapedArray:
    """Element-wise minimum. Scalars/NumPy upcast; dtypes via ``result_type``.

    Notes:
        Under TF XLA, NaN inputs may yield ``±inf`` instead of NaN. Not
        standardized across backends.
    """
    return array_namespace(x, y).minimum(x, y)

sqrt

sqrt(x)

Element-wise square root.

Source code in anytensor/core.py

@as_array_result
def sqrt(x: ShapedArray) -> ShapedArray:
    """Element-wise square root."""
    return array_namespace(x).sqrt(x)

rsqrt

rsqrt(x)

Element-wise reciprocal square root (1 / sqrt(x)).

Source code in anytensor/core.py

@as_array_result
def rsqrt(x: ShapedArray) -> ShapedArray:
    """Element-wise reciprocal square root (``1 / sqrt(x)``)."""
    xp = array_namespace(x)
    return xp.asarray(1.0, dtype=getattr(x, "dtype", None)) / xp.sqrt(x)

where

where(condition, x, y)

Choose from x or y by condition. Scalars/NumPy upcast.

Source code in anytensor/core.py

@as_array_result
@promote(condition="mask", x="data", y="data")
def where(condition: ShapedArray, x: ShapedArray, y: ShapedArray) -> ShapedArray:
    """Choose from ``x`` or ``y`` by ``condition``. Scalars/NumPy upcast."""
    return array_namespace(condition, x, y).where(condition, x, y)

clip

clip(x, min=None, max=None)

Clip values to [min, max].

Implemented with :func:maximum / :func:minimum so TF tracing does not emit BroadcastArgs (tf2onnx cannot lower that op).

Source code in anytensor/core.py

@as_array_result
def clip(x: ShapedArray, min: Any = None, max: Any = None) -> ShapedArray:
    """Clip values to ``[min, max]``.

    Implemented with :func:`maximum` / :func:`minimum` so TF tracing does not
    emit ``BroadcastArgs`` (tf2onnx cannot lower that op).
    """
    out = x
    if min is not None:
        out = maximum(out, min)
    if max is not None:
        out = minimum(out, max)
    return out

astype

astype(x, dtype)

Cast x to dtype.

Source code in anytensor/core.py

@as_array_result
def astype(x: ShapedArray, dtype: DtypeLike) -> ShapedArray:
    """Cast ``x`` to ``dtype``."""
    return array_namespace(x).astype(x, dtype)

cast

cast(x, dtype)

Alias of :func:astype.

Source code in anytensor/core.py

def cast(x: ShapedArray, dtype: DtypeLike) -> ShapedArray:
    """Alias of :func:`astype`."""
    return astype(x, dtype)

zeros_like

zeros_like(x, dtype=None)

Return zeros with the same shape (and backend) as x.

Built from :func:shape so polymorphic tf.function sees symbolic sizes instead of None dims (raw tnp.zeros_like / tnp.zeros can fail after retracing).

Source code in anytensor/core.py

@as_array_result
def zeros_like(x: ShapedArray, dtype: DtypeLike = None) -> ShapedArray:
    """Return zeros with the same shape (and backend) as ``x``.

    Built from :func:`shape` so polymorphic ``tf.function`` sees symbolic
    sizes instead of ``None`` dims (raw ``tnp.zeros_like`` / ``tnp.zeros`` can
    fail after retracing).
    """
    xp = array_namespace(x)
    return xp.zeros(shape(x), dtype=x.dtype if dtype is None else dtype)

ones_like

ones_like(x, dtype=None)

Return ones with the same shape (and backend) as x.

Uses :func:shape for graph-safe sizes under tf.function (see :func:zeros_like).

Source code in anytensor/core.py

@as_array_result
def ones_like(x: ShapedArray, dtype: DtypeLike = None) -> ShapedArray:
    """Return ones with the same shape (and backend) as ``x``.

    Uses :func:`shape` for graph-safe sizes under ``tf.function`` (see
    :func:`zeros_like`).
    """
    xp = array_namespace(x)
    return xp.ones(shape(x), dtype=x.dtype if dtype is None else dtype)

full_like

full_like(x, fill_value, dtype=None)

Return an array filled with fill_value matching x.

Uses :func:shape for graph-safe sizes under tf.function (see :func:zeros_like).

Source code in anytensor/core.py

@as_array_result
def full_like(x: ShapedArray, fill_value: Any, dtype: DtypeLike = None) -> ShapedArray:
    """Return an array filled with ``fill_value`` matching ``x``.

    Uses :func:`shape` for graph-safe sizes under ``tf.function`` (see
    :func:`zeros_like`).
    """
    xp = array_namespace(x)
    return xp.full(shape(x), fill_value, dtype=x.dtype if dtype is None else dtype)

zeros

zeros(shape, *, dtype=None, like=None)

Return zeros; pass like= to select the backend.

Source code in anytensor/core.py

@as_array_result
def zeros(shape: ShapeLike, *, dtype: DtypeLike = None, like: Optional[ShapedArray] = None) -> ShapedArray:
    """Return zeros; pass ``like=`` to select the backend."""
    if like is None:
        import array_api_compat.numpy as xp

        return xp.zeros(shape, dtype=dtype)
    xp = array_namespace(like)
    return xp.zeros(shape, dtype=dtype if dtype is not None else like.dtype)

ones

ones(shape, *, dtype=None, like=None)

Return ones; pass like= to select the backend.

Source code in anytensor/core.py

@as_array_result
def ones(shape: ShapeLike, *, dtype: DtypeLike = None, like: Optional[ShapedArray] = None) -> ShapedArray:
    """Return ones; pass ``like=`` to select the backend."""
    if like is None:
        import array_api_compat.numpy as xp

        return xp.ones(shape, dtype=dtype)
    xp = array_namespace(like)
    return xp.ones(shape, dtype=dtype if dtype is not None else like.dtype)

full

full(shape, fill_value, *, dtype=None, like=None)

Return a filled array; pass like= to select the backend.

Source code in anytensor/core.py

@as_array_result
def full(shape: ShapeLike, fill_value: Any, *, dtype: DtypeLike = None, like: Optional[ShapedArray] = None) -> ShapedArray:
    """Return a filled array; pass ``like=`` to select the backend."""
    if like is None:
        import array_api_compat.numpy as xp

        return xp.full(shape, fill_value, dtype=dtype)
    xp = array_namespace(like)
    return xp.full(shape, fill_value, dtype=dtype if dtype is not None else like.dtype)

arange

arange(
    start,
    /,
    stop=None,
    step=1,
    *,
    dtype=None,
    like=None,
    device=None,
)

Evenly spaced values; pass like= to select the backend.

device is forwarded when supported (e.g. Torch).

Source code in anytensor/core.py

@as_array_result
def arange(start: Any, /, stop: Any = None, step: Any = 1, *, dtype: DtypeLike = None, like: Optional[ShapedArray] = None, device: Any = None) -> ShapedArray:
    """Evenly spaced values; pass ``like=`` to select the backend.

    ``device`` is forwarded when supported (e.g. Torch).
    """
    if stop is None:
        start, stop = 0, start
    if like is None:
        import array_api_compat.numpy as xp

        return xp.arange(start, stop, step, dtype=dtype)
    xp = array_namespace(like)
    kwargs = {}
    if dtype is not None:
        kwargs["dtype"] = dtype
    if device is not None:
        kwargs["device"] = device
    try:
        return xp.arange(start, stop, step, **kwargs)
    except TypeError:
        kwargs.pop("device", None)
        return xp.arange(start, stop, step, **kwargs)

repeat

repeat(x, repeats, *, total_repeat_length=None, axis=None)

Repeat elements of x.

Similar to NumPy / JAX repeat. When total_repeat_length is set with per-element repeats, matches JAX jnp.repeat(..., total_repeat_length=) (pad or slice the leading axis to that length).

Parameters:

Name	Type	Description	Default
x	ShapedArray	Input array.	required
repeats	Any	Python int (same count for every element) or an integral array of per-element counts. Python scalar repeats stay Python (not promoted to 0-d tensors) so tf.function / jax.jit / torch.compile stay happy — promoting a 2 to a 0-d TF tensor breaks tf.experimental.numpy.repeat.	required
total_repeat_length	Optional[ShapeSize]	Optional shape-size (Python int, symbolic constant, or 0-d integral tensor) for the flattened output length. Required for a static output size under jax.jit when repeats are dynamic. Not supported together with axis is not None yet. :func:~anytensor.partition_softmax always passes this as total_length.	None
axis	Optional[int]	Axis to repeat along; None flattens (Array API / NumPy style).	None

Returns:

Type	Description
ShapedArray	Array with repeated elements on the same backend as x.

Notes

Under jax.jit, jnp.repeat needs static repeat counts or a static total_repeat_length. :func:~anytensor.partition_softmax always passes total_length as that length. Omitting total_repeat_length on :func:repeat itself is fine eagerly and on TensorFlow.

Source code in anytensor/core.py

@as_array_result
@promote(x="data")
def repeat(x: ShapedArray, repeats: Any, *, total_repeat_length: Optional[ShapeSize] = None, axis: Optional[int] = None) -> ShapedArray:
    """Repeat elements of ``x``.

    Similar to NumPy / JAX ``repeat``. When ``total_repeat_length`` is set with
    per-element ``repeats``, matches JAX ``jnp.repeat(..., total_repeat_length=)``
    (pad or slice the leading axis to that length).

    Args:
        x: Input array.
        repeats: Python ``int`` (same count for every element) or an integral
            array of per-element counts. **Python scalar repeats stay Python**
            (not promoted to 0-d tensors) so ``tf.function`` / ``jax.jit`` /
            ``torch.compile`` stay happy — promoting a ``2`` to a 0-d TF tensor
            breaks ``tf.experimental.numpy.repeat``.
        total_repeat_length: Optional shape-size (Python ``int``, symbolic
            constant, or 0-d integral tensor) for the flattened output length.
            Required for a static output size under ``jax.jit`` when repeats
            are dynamic. Not supported together with ``axis is not None`` yet.
            :func:`~anytensor.partition_softmax` always passes this as
            ``total_length``.
        axis: Axis to repeat along; ``None`` flattens (Array API / NumPy style).

    Returns:
        Array with repeated elements on the same backend as ``x``.

    Notes:
        Under ``jax.jit``, ``jnp.repeat`` needs static repeat counts or a
        static ``total_repeat_length``. :func:`~anytensor.partition_softmax`
        always passes ``total_length`` as that length. Omitting
        ``total_repeat_length`` on :func:`repeat` itself is fine eagerly and
        on TensorFlow.
    """
    from .backends import get_backend

    xp = array_namespace(x)
    if not _is_scalar(repeats):
        repeats = _asarray(xp, repeats)
        if not xp.isdtype(repeats.dtype, "integral"):
            raise TypeError(
                f"repeats must be integral, got dtype={getattr(repeats, 'dtype', type(repeats))}"
            )

    if total_repeat_length is not None:
        if axis is not None:
            raise NotImplementedError("total_repeat_length with axis!=None is not supported yet")
        backend = get_backend(x)
        try:
            return backend.repeat(x, repeats, total_repeat_length)
        except NotImplementedError:
            pass
        flat = xp.reshape(x, (-1,))
        # Python-loop path only when repeat *values* are host ints. Under
        # tf.function, shapes can look concrete while repeats[i] is symbolic.
        if (
            _leading_dim_is_concrete(flat.shape[0])
            and _repeats_are_host_concrete(repeats)
            and _leading_dim_is_concrete(total_repeat_length)
        ):
            return _repeat_total_length_concrete(
                xp, x, repeats, int(_normalize_shape_dim(total_repeat_length))
            )
        out = xp.repeat(flat, repeats)
        return _pad_or_slice_leading(xp, out, _normalize_shape_dim(total_repeat_length))

    if axis is None:
        return xp.repeat(x, repeats)
    return xp.repeat(x, repeats, axis=axis)

matmul

matmul(x, y)

Matrix product of two arrays. NumPy operands upcast onto peers.

Source code in anytensor/core.py

@as_array_result
@promote(x="data", y="data")
def matmul(x: ShapedArray, y: ShapedArray) -> ShapedArray:
    """Matrix product of two arrays. NumPy operands upcast onto peers."""
    return array_namespace(x, y).matmul(x, y)

inf

inf(like)

+inf from the backend of like.

Source code in anytensor/core.py

def inf(like: ShapedArray) -> Any:
    """``+inf`` from the backend of ``like``."""
    return _backend_attr(like, "inf")

ninf

ninf(like)

-inf from the backend of like.

Source code in anytensor/core.py

def ninf(like: ShapedArray) -> Any:
    """``-inf`` from the backend of ``like``."""
    return _backend_attr(like, "ninf")

nan

nan(like)

NaN from the backend of like.

Source code in anytensor/core.py

def nan(like: ShapedArray) -> Any:
    """``NaN`` from the backend of ``like``."""
    return _backend_attr(like, "nan")

pi

pi(like)

π from the backend of like.

Source code in anytensor/core.py

def pi(like: ShapedArray) -> Any:
    """``π`` from the backend of ``like``."""
    return _backend_attr(like, "pi")

e

e(like)

Euler's number from the backend of like.

Source code in anytensor/core.py

def e(like: ShapedArray) -> Any:
    """Euler's number from the backend of ``like``."""
    return _backend_attr(like, "e")

finfo

finfo(x)

Floating limits for x.dtype via x's backend.

Source code in anytensor/core.py

def finfo(x: ArrayT) -> Any:
    """Floating limits for ``x.dtype`` via ``x``'s backend."""
    return _backend_attr(x, "finfo")(x.dtype)

iinfo

iinfo(x)

Integral limits for x.dtype via x's backend.

Source code in anytensor/core.py

def iinfo(x: ArrayT) -> Any:
    """Integral limits for ``x.dtype`` via ``x``'s backend."""
    return _backend_attr(x, "iinfo")(x.dtype)

dtype

dtype(name, like)

Framework dtype name from the backend of like (e.g. "bool").

Source code in anytensor/core.py

def dtype(name: str, like: ShapedArray) -> DtypeLike:
    """Framework dtype ``name`` from the backend of ``like`` (e.g. ``\"bool\"``)."""
    return _backend_attr(like, name)

is_nan

is_nan(x)

Element-wise NaN test (Array API isnan).

Source code in anytensor/core.py

@as_array_result
def is_nan(x: ShapedArray) -> ShapedArray:
    """Element-wise NaN test (Array API ``isnan``)."""
    return array_namespace(x).isnan(x)

is_finite

is_finite(x)

Element-wise finite test (Array API isfinite).

Source code in anytensor/core.py

@as_array_result
def is_finite(x: ShapedArray) -> ShapedArray:
    """Element-wise finite test (Array API ``isfinite``)."""
    return array_namespace(x).isfinite(x)

is_inf

is_inf(x)

Element-wise infinity test (Array API isinf).

Source code in anytensor/core.py

@as_array_result
def is_inf(x: ShapedArray) -> ShapedArray:
    """Element-wise infinity test (Array API ``isinf``)."""
    return array_namespace(x).isinf(x)

fill_nan

fill_nan(x, value=0.0)

Replace NaNs in x with value (broadcasts). Leaves ±inf unchanged.

Source code in anytensor/core.py

@as_array_result
@promote(x="data", value="data")
def fill_nan(x: ShapedArray, value: Any = 0.0) -> ShapedArray:
    """Replace NaNs in ``x`` with ``value`` (broadcasts). Leaves ±inf unchanged."""
    xp = array_namespace(x, value)
    return xp.where(xp.isnan(x), value, x)

fill_nan_mask

fill_nan_mask(x, value=0.0)

Return (filled, mask): NaNs replaced, plus a boolean NaN mask.

mask is True where x was NaN (same polarity as :func:is_nan / NumPy masked-array invalid). Boolean, not 0/1 — cast if you need weights. Equivalent to (fill_nan(x, value), is_nan(x)); not ~is_finite (±inf is non-NaN).

Source code in anytensor/core.py

@promote(x="data", value="data")
def fill_nan_mask(x: ShapedArray, value: Any = 0.0) -> tuple[ShapedArray, ShapedArray]:
    """Return ``(filled, mask)``: NaNs replaced, plus a boolean NaN mask.

    ``mask`` is ``True`` where ``x`` was NaN (same polarity as :func:`is_nan` /
    NumPy masked-array invalid). Boolean, not 0/1 — cast if you need weights.
    Equivalent to ``(fill_nan(x, value), is_nan(x))``; not ``~is_finite``
    (±inf is non-NaN).
    """
    xp = array_namespace(x, value)
    mask = xp.isnan(x)
    filled = xp.where(mask, value, x)
    return filled, mask

nan_to_num

nan_to_num(x, *, nan=0.0, posinf=None, neginf=None)

Replace NaN and ±inf (Array API nan_to_num).

Defaults: NaN → nan (0.0); posinf / neginf None → large finite values from the dtype's finfo (framework-dependent).

Source code in anytensor/core.py

@as_array_result
@promote(x="data")
def nan_to_num(x: ShapedArray, *, nan: Any = 0.0, posinf: Any = None, neginf: Any = None) -> ShapedArray:
    """Replace NaN and ±inf (Array API ``nan_to_num``).

    Defaults: NaN → ``nan`` (0.0); ``posinf`` / ``neginf`` ``None`` → large
    finite values from the dtype's finfo (framework-dependent).
    """
    return array_namespace(x).nan_to_num(x, nan=nan, posinf=posinf, neginf=neginf)

equal_nan

equal_nan(x, y)

Element-wise equality treating NaN as equal to NaN.

Returns a boolean array: (x == y) | (isnan(x) & isnan(y)). Non-NaN values compare with ordinary == (so +inf == +inf).

Source code in anytensor/core.py

@as_array_result
@promote(x="data", y="data")
def equal_nan(x: ShapedArray, y: ShapedArray) -> ShapedArray:
    """Element-wise equality treating NaN as equal to NaN.

    Returns a boolean array: ``(x == y) | (isnan(x) & isnan(y))``.
    Non-NaN values compare with ordinary ``==`` (so ``+inf == +inf``).
    """
    xp = array_namespace(x, y)
    both_nan = xp.logical_and(xp.isnan(x), xp.isnan(y))
    return xp.logical_or(x == y, both_nan)