Skip to content

Package root

anytensor

get_backend 

get_backend(tensor)

Return the backend for tensor (e.g. NumPy for numpy.ndarray).

Optional extras (JAX, Torch, TensorFlow) are used only if already imported; this never imports them. NumPy is a required dependency and is the fallback for ndarrays.

Source code in anytensor/backends.py 
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
def get_backend(tensor: Any) -> "AbstractBackend":
    """Return the backend for ``tensor`` (e.g. NumPy for ``numpy.ndarray``).

    Optional extras (JAX, Torch, TensorFlow) are used only if already imported;
    this never imports them. NumPy is a required dependency and is the fallback
    for ndarrays.
    """
    _type = type(tensor)
    _result = _type2backend.get(_type, None)
    if _result is not None:
        return _result

    for framework_name, backend in list(_loaded_backends.items()):
        if backend.is_appropriate_type(tensor):
            _type2backend[_type] = backend
            return backend

    # Find backend subclasses recursively
    backend_subclasses = []
    backends = AbstractBackend.__subclasses__()
    while backends:
        backend = backends.pop()
        backends += backend.__subclasses__()
        backend_subclasses.append(backend)

    for BackendSubclass in backend_subclasses:
        if _debug_importing:
            print("Testing for subclass of ", BackendSubclass)
        if BackendSubclass.framework_name not in _loaded_backends:
            # Construct only if the extra is already imported; never import it here.
            if module_if_loaded(BackendSubclass.framework_name) is not None:
                if _debug_importing:
                    print("Imported backend for ", BackendSubclass.framework_name)
                backend = BackendSubclass()
                _loaded_backends[backend.framework_name] = backend
                if backend.is_appropriate_type(tensor):
                    _type2backend[_type] = backend
                    return backend

    raise RuntimeError(f"Tensor type unknown to anytensor: {type(tensor)}")

module_if_loaded 

module_if_loaded(name, callback=None, *, raises=False)

Return name if it is already imported, else None.

Never imports name. If callback is given and the module is already loaded, it is invoked immediately with the module. If not, callback is registered first and invoked later when that module is imported in this process (including when a submodule import loads the parent).

If raises is true and the module is still absent after that registration, raise :class:RuntimeError. The pending callback is kept, so a later import still runs it.

Parameters:

Name Type Description Default
name str

Absolute module name ("torch", "jax", "tensorflow").

 required
callback Optional[_Callback]

Optional callback(module) run now or on a future import.

 None
raises bool

If true, raise when name is not already imported.

 False

Returns:

Type Description
ModuleType | None

The loaded module, or None if it is not in sys.modules
ModuleType | None

(only when raises is false).

Examples:

Check without importing::

torch = module_if_loaded("torch")
if torch is None:
    return

Require an extra that must already be imported::

jax = module_if_loaded("jax", raises=True)

Register a side effect for now-or-later (JAX pytree, TorchScript)::

def _register_jax_pytree(jax):
    jax.tree_util.register_pytree_node(Ragged, flatten, unflatten)

module_if_loaded("jax", _register_jax_pytree)
# or module_if_loaded("jax", _register_jax_pytree, raises=True)
# to fail now while still running the callback on a later import.
Source code in anytensor/optional.py 
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
def module_if_loaded(
    name: str, callback: Optional[_Callback] = None, *, raises: bool = False
) -> ModuleType | None:
    """Return ``name`` if it is already imported, else ``None``.

    Never imports ``name``. If ``callback`` is given and the module is
    already loaded, it is invoked immediately with the module. If not,
    ``callback`` is registered first and invoked later when that module is
    imported in this process (including when a submodule import loads the
    parent).

    If ``raises`` is true and the module is still absent after that
    registration, raise :class:`RuntimeError`. The pending callback is kept,
    so a later import still runs it.

    Args:
        name: Absolute module name (``"torch"``, ``"jax"``, ``"tensorflow"``).
        callback: Optional ``callback(module)`` run now or on a future import.
        raises: If true, raise when ``name`` is not already imported.

    Returns:
        The loaded module, or ``None`` if it is not in ``sys.modules``
        (only when ``raises`` is false).

    Examples:
        Check without importing::

            torch = module_if_loaded("torch")
            if torch is None:
                return

        Require an extra that must already be imported::

            jax = module_if_loaded("jax", raises=True)

        Register a side effect for now-or-later (JAX pytree, TorchScript)::

            def _register_jax_pytree(jax):
                jax.tree_util.register_pytree_node(Ragged, flatten, unflatten)

            module_if_loaded("jax", _register_jax_pytree)
            # or module_if_loaded("jax", _register_jax_pytree, raises=True)
            # to fail now while still running the callback on a later import.
    """
    mod = _module(name)
    if callback is not None:
        if mod is not None:
            callback(mod)
        else:
            with _lock:
                _pending.setdefault(name, []).append(callback)
            _install_hook()
            # A concurrent import may have finished while we registered.
            _fire_ready()
            mod = _module(name)
    if mod is None and raises:
        raise RuntimeError(f"{name} is not imported")
    return mod

segment_sum 

segment_sum(x, segment_ids, num_segments, sorted=False)

Computes the sum within segments of an array.

Similar to :func:jax.ops.segment_sum and TF unsorted_segment_sum. Reduces x along axis 0, summing rows that share the same segment_ids entry.

Parameters:

Name Type Description Default
x SegmentValues

Values to sum. The leading axis is the segment axis.

 required
segment_ids SegmentIds

Integer ids with segment_ids.shape[0] == x.shape[0]. Values may be repeated and need not be sorted. Ids outside [0, num_segments) (including negatives) are dropped and do not contribute; they are not wrapped like Python negative indexing.

 required
num_segments ShapeSize

Required shape-size (unlike JAX, where omitting it defaults to max(segment_ids) + 1). Accepts a Python int, a jit/compile symbolic constant, or a 0-d integral tensor scalar — never inferred from segment_ids. Sets the output length; empty slots are 0.

 required
sorted bool

When True, JAX/TF may use a sorted-ids fast path. No-op on NumPy and Torch (unsorted-safe scatter).

 False

Returns:

Type Description
SegmentOut

Array of shape (num_segments,) + x.shape[1:] on the same backend
SegmentOut

as x.
Notes

Empty-segment identity is 0 for all dtypes (:mod:anytensor.semantics).

Index width is backend-local: Torch casts ids to int64 at scatter; JAX without jax_enable_x64 often keeps int32 and may warn on int64 ids.

Under torch.jit.script, import torch in either order relative to AnyTensor — the divert auto-enables via :func:anytensor.module_if_loaded. Eager calls still dispatch by tensor type; only the scripted path uses :mod:anytensor.torchscript.

Examples:

>>> import numpy as np
>>> import anytensor as at
>>> x = np.arange(5.0)
>>> ids = np.array([0, 0, 1, 1, 2])
>>> at.segment_sum(x, ids, num_segments=3)
array([1., 5., 4.])
Source code in anytensor/segment.py 
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
def segment_sum(x: SegmentValues, segment_ids: SegmentIds, num_segments: ShapeSize, sorted: bool = False) -> SegmentOut:
    """Computes the sum within segments of an array.

    Similar to :func:`jax.ops.segment_sum` and TF ``unsorted_segment_sum``.
    Reduces ``x`` along axis 0, summing rows that share the same
    ``segment_ids`` entry.

    Args:
        x: Values to sum. The leading axis is the segment axis.
        segment_ids: Integer ids with ``segment_ids.shape[0] == x.shape[0]``.
            Values may be repeated and need not be sorted. Ids outside
            ``[0, num_segments)`` (including negatives) are dropped and do not
            contribute; they are **not** wrapped like Python negative indexing.
        num_segments: **Required** shape-size (unlike JAX, where omitting it
            defaults to ``max(segment_ids) + 1``). Accepts a Python ``int``, a
            jit/compile symbolic constant, or a 0-d integral tensor scalar —
            never inferred from ``segment_ids``. Sets the output length; empty
            slots are ``0``.
        sorted: When True, JAX/TF may use a sorted-ids fast path. **No-op on
            NumPy and Torch** (unsorted-safe scatter).

    Returns:
        Array of shape ``(num_segments,) + x.shape[1:]`` on the same backend
        as ``x``.

    Notes:
        Empty-segment identity is ``0`` for all dtypes
        (:mod:`anytensor.semantics`).

        Index **width** is backend-local: Torch casts ids to int64 at scatter;
        JAX without ``jax_enable_x64`` often keeps int32 and may warn on int64
        ids.

        Under ``torch.jit.script``, import ``torch`` in either order relative to
        AnyTensor — the divert auto-enables via :func:`anytensor.module_if_loaded`.
        Eager calls still dispatch by tensor type; only the scripted path uses
        :mod:`anytensor.torchscript`.

    Examples:
        >>> import numpy as np
        >>> import anytensor as at
        >>> x = np.arange(5.0)
        >>> ids = np.array([0, 0, 1, 1, 2])
        >>> at.segment_sum(x, ids, num_segments=3)
        array([1., 5., 4.])
    """
    return _segment_reduce(x, segment_ids, num_segments, "sum", sorted)

segment_max 

segment_max(x, segment_ids, num_segments, sorted=False)

Computes the maximum within segments of an array.

Similar to :func:jax.ops.segment_max. Reduces x along axis 0 by segment_ids.

Parameters:

Name Type Description Default
x SegmentValues

Values to reduce. Leading axis is the segment axis.

 required
segment_ids SegmentIds

Integer ids with segment_ids.shape[0] == x.shape[0]. May be unsorted / non-unique. Out-of-range and negative ids are dropped (not wrapped).

 required
num_segments ShapeSize

Required shape-size (Python int, symbolic constant, or 0-d integral tensor). Empty slots get the max identity (-inf for floats, iinfo(dtype).min for ints).

 required
sorted bool

Honored on JAX/TF; no-op on NumPy/Torch.

 False

Returns:

Type Description
SegmentOut

Array of shape (num_segments,) + x.shape[1:].
Notes

Empty-slot identity is -inf / dtype min — not TF's native unsorted_segment_max finfo fill. Occupied ±inf stays ±inf.

On TensorFlow, scatter ignores NaN updates; AnyTensor ORs NaN back in so a segment that saw any NaN is NaN (matches NumPy/JAX/Torch).

Prefer :func:segment_max_or_constant when empty slots should be a finite fill instead of -inf.

Under TF XLA (tf.function(jit_compile=True)), NaN in min/max-like ops may become ±inf instead of NaN — avoid relying on NaN under XLA.

TorchScript: :func:enable_torchscript (eager path stays multi-backend).

Source code in anytensor/segment.py 
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
def segment_max(x: SegmentValues, segment_ids: SegmentIds, num_segments: ShapeSize, sorted: bool = False) -> SegmentOut:
    """Computes the maximum within segments of an array.

    Similar to :func:`jax.ops.segment_max`. Reduces ``x`` along axis 0 by
    ``segment_ids``.

    Args:
        x: Values to reduce. Leading axis is the segment axis.
        segment_ids: Integer ids with ``segment_ids.shape[0] == x.shape[0]``.
            May be unsorted / non-unique. Out-of-range and negative ids are
            dropped (not wrapped).
        num_segments: **Required** shape-size (Python ``int``, symbolic
            constant, or 0-d integral tensor). Empty slots get the max
            identity (``-inf`` for floats, ``iinfo(dtype).min`` for ints).
        sorted: Honored on JAX/TF; **no-op on NumPy/Torch**.

    Returns:
        Array of shape ``(num_segments,) + x.shape[1:]``.

    Notes:
        Empty-slot identity is ``-inf`` / dtype min — **not** TF's native
        ``unsorted_segment_max`` finfo fill. Occupied ±inf stays ±inf.

        On TensorFlow, scatter ignores NaN updates; AnyTensor ORs NaN back in
        so a segment that saw any NaN is NaN (matches NumPy/JAX/Torch).

        Prefer :func:`segment_max_or_constant` when empty slots should be a
        finite fill instead of ``-inf``.

        Under TF XLA (``tf.function(jit_compile=True)``), NaN in min/max-like
        ops may become ±inf instead of NaN — avoid relying on NaN under XLA.

        TorchScript: :func:`enable_torchscript` (eager path stays multi-backend).
    """
    return _segment_reduce(x, segment_ids, num_segments, "max", sorted)

segment_min 

segment_min(x, segment_ids, num_segments, sorted=False)

Computes the minimum within segments of an array.

Similar to :func:jax.ops.segment_min. Reduces x along axis 0 by segment_ids.

Parameters:

Name Type Description Default
x SegmentValues

Values to reduce. Leading axis is the segment axis.

 required
segment_ids SegmentIds

Integer ids with segment_ids.shape[0] == x.shape[0]. May be unsorted / non-unique. Out-of-range and negative ids are dropped (not wrapped).

 required
num_segments ShapeSize

Required shape-size (Python int, symbolic constant, or 0-d integral tensor). Empty slots get the min identity (+inf for floats, iinfo(dtype).max for ints).

 required
sorted bool

Honored on JAX/TF; no-op on NumPy/Torch.

 False

Returns:

Type Description
SegmentOut

Array of shape (num_segments,) + x.shape[1:].
Notes

Empty-slot identity is +inf / dtype max — not TF's native unsorted_segment_min finfo fill. Occupied ±inf stays ±inf.

On TensorFlow, scatter ignores NaN updates; AnyTensor ORs NaN back in so a segment that saw any NaN is NaN (matches NumPy/JAX/Torch).

Prefer :func:segment_min_or_constant when empty slots should be a finite fill instead of +inf.

Under TF XLA, NaN in min/max-like ops may become ±inf — avoid relying on NaN under XLA.

TorchScript: :func:enable_torchscript (eager path stays multi-backend).

Source code in anytensor/segment.py 
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
def segment_min(x: SegmentValues, segment_ids: SegmentIds, num_segments: ShapeSize, sorted: bool = False) -> SegmentOut:
    """Computes the minimum within segments of an array.

    Similar to :func:`jax.ops.segment_min`. Reduces ``x`` along axis 0 by
    ``segment_ids``.

    Args:
        x: Values to reduce. Leading axis is the segment axis.
        segment_ids: Integer ids with ``segment_ids.shape[0] == x.shape[0]``.
            May be unsorted / non-unique. Out-of-range and negative ids are
            dropped (not wrapped).
        num_segments: **Required** shape-size (Python ``int``, symbolic
            constant, or 0-d integral tensor). Empty slots get the min
            identity (``+inf`` for floats, ``iinfo(dtype).max`` for ints).
        sorted: Honored on JAX/TF; **no-op on NumPy/Torch**.

    Returns:
        Array of shape ``(num_segments,) + x.shape[1:]``.

    Notes:
        Empty-slot identity is ``+inf`` / dtype max — **not** TF's native
        ``unsorted_segment_min`` finfo fill. Occupied ±inf stays ±inf.

        On TensorFlow, scatter ignores NaN updates; AnyTensor ORs NaN back in
        so a segment that saw any NaN is NaN (matches NumPy/JAX/Torch).

        Prefer :func:`segment_min_or_constant` when empty slots should be a
        finite fill instead of ``+inf``.

        Under TF XLA, NaN in min/max-like ops may become ±inf — avoid relying
        on NaN under XLA.

        TorchScript: :func:`enable_torchscript` (eager path stays multi-backend).
    """
    return _segment_reduce(x, segment_ids, num_segments, "min", sorted)

segment_mean 

segment_mean(x, segment_ids, num_segments, sorted=False)

Mean of values of x within each segment along axis 0.

Parameters:

Name Type Description Default
x SegmentValues

Values to average. Leading axis is the segment axis.

 required
segment_ids SegmentIds

Integer segment ids (see :func:segment_sum).

 required
num_segments ShapeSize

Required shape-size.

 required
sorted bool

Forwarded to underlying segment ops (no-op on NumPy/Torch).

 False

Returns:

Type Description
SegmentOut

Array of shape (num_segments,) + x.shape[1:].
Notes

Empty segments yield 0 (sum is already 0; the denominator is clamped away from zero only so division stays defined). This differs from a NaN-on-empty mean.

Source code in anytensor/segment.py 
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
def segment_mean(x: SegmentValues, segment_ids: SegmentIds, num_segments: ShapeSize, sorted: bool = False) -> SegmentOut:
    """Mean of values of ``x`` within each segment along axis 0.

    Args:
        x: Values to average. Leading axis is the segment axis.
        segment_ids: Integer segment ids (see :func:`segment_sum`).
        num_segments: Required shape-size.
        sorted: Forwarded to underlying segment ops (no-op on NumPy/Torch).

    Returns:
        Array of shape ``(num_segments,) + x.shape[1:]``.

    Notes:
        Empty segments yield **0** (sum is already 0; the denominator is
        clamped away from zero only so division stays defined). This differs
        from a NaN-on-empty mean.
    """
    total = segment_sum(x, segment_ids, num_segments, sorted=sorted)
    counts = segment_count(segment_ids, num_segments, sorted=sorted)
    xp = array_namespace(total, counts)
    # Broadcast counts over trailing dims of x.
    while counts.ndim < total.ndim:
        counts = xp.expand_dims(counts, axis=-1)
    denom = maximum(counts, xp.asarray(1.0, dtype=counts.dtype))
    denom = xp.astype(denom, total.dtype)
    return total / denom

segment_count 

segment_count(segment_ids, num_segments, sorted=False)

Count how many elements fall in each segment.

Implemented as segment_sum of ones, so empty slots are 0.0.

Parameters:

Name Type Description Default
segment_ids SegmentIds

Integer ids (same conventions as :func:segment_sum).

 required
num_segments ShapeSize

Required shape-size (Python int / symbolic / 0-d integral tensor).

 required
sorted bool

Forwarded to :func:segment_sum (no-op on NumPy/Torch).

 False

Returns:

Type Description
SegmentOut

Float vector of shape (num_segments,) on the same backend as
SegmentOut

segment_ids. Dtype follows the backend default float (often
SegmentOut

float32 on JAX without x64).
Source code in anytensor/segment.py 
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
def segment_count(segment_ids: SegmentIds, num_segments: ShapeSize, sorted: bool = False) -> SegmentOut:
    """Count how many elements fall in each segment.

    Implemented as ``segment_sum`` of ones, so empty slots are ``0.0``.

    Args:
        segment_ids: Integer ids (same conventions as :func:`segment_sum`).
        num_segments: Required shape-size (Python ``int`` / symbolic / 0-d
            integral tensor).
        sorted: Forwarded to :func:`segment_sum` (no-op on NumPy/Torch).

    Returns:
        Float vector of shape ``(num_segments,)`` on the same backend as
        ``segment_ids``. Dtype follows the backend default float (often
        float32 on JAX without x64).
    """
    ones = ones_like(segment_ids)
    # Float counts for division; use default float (float32 on JAX without x64).
    xp = array_namespace(ones)
    ones = xp.astype(ones, xp.asarray(0.0).dtype)
    return segment_sum(ones, segment_ids, num_segments, sorted=sorted)

segment_variance 

segment_variance(
    x, segment_ids, num_segments, sorted=False
)

Population variance of x within each segment along axis 0.

Computed as the segment mean of squared deviations from the segment mean (divide by n, not n-1).

Parameters:

Name Type Description Default
x SegmentValues

Values. Leading axis is the segment axis.

 required
segment_ids SegmentIds

Integer segment ids (see :func:segment_sum).

 required
num_segments ShapeSize

Required shape-size.

 required
sorted bool

Forwarded to underlying segment ops (no-op on NumPy/Torch).

 False

Returns:

Type Description
SegmentOut

Array of shape (num_segments,) + x.shape[1:]. Empty segments are
SegmentOut

0 (same empty-mean convention as :func:segment_mean).
Source code in anytensor/segment.py 
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
def segment_variance(x: SegmentValues, segment_ids: SegmentIds, num_segments: ShapeSize, sorted: bool = False) -> SegmentOut:
    """Population variance of ``x`` within each segment along axis 0.

    Computed as the segment mean of squared deviations from the segment mean
    (divide by ``n``, not ``n-1``).

    Args:
        x: Values. Leading axis is the segment axis.
        segment_ids: Integer segment ids (see :func:`segment_sum`).
        num_segments: Required shape-size.
        sorted: Forwarded to underlying segment ops (no-op on NumPy/Torch).

    Returns:
        Array of shape ``(num_segments,) + x.shape[1:]``. Empty segments are
        ``0`` (same empty-mean convention as :func:`segment_mean`).
    """
    mean = segment_mean(x, segment_ids, num_segments, sorted=sorted)
    mean_gathered = take(mean, segment_ids)
    centered = x - mean_gathered
    return segment_mean(centered * centered, segment_ids, num_segments, sorted=sorted)

segment_normalize 

segment_normalize(
    x, segment_ids, num_segments, sorted=False
)

Divide each value by its segment sum (0/0 → 0).

Parameters:

Name Type Description Default
x SegmentValues

Values. Leading axis is the segment axis.

 required
segment_ids SegmentIds

Integer segment ids (see :func:segment_sum).

 required
num_segments ShapeSize

Required shape-size.

 required
sorted bool

Forwarded to :func:segment_sum (no-op on NumPy/Torch).

 False

Returns:

Type Description
SegmentValues

Array with the same shape as x. Rows whose segment sum is 0
SegmentValues

become 0 (no NaN from 0/0).
Source code in anytensor/segment.py 
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
def segment_normalize(x: SegmentValues, segment_ids: SegmentIds, num_segments: ShapeSize, sorted: bool = False) -> SegmentValues:
    """Divide each value by its segment sum (``0/0`` → ``0``).

    Args:
        x: Values. Leading axis is the segment axis.
        segment_ids: Integer segment ids (see :func:`segment_sum`).
        num_segments: Required shape-size.
        sorted: Forwarded to :func:`segment_sum` (no-op on NumPy/Torch).

    Returns:
        Array with the same shape as ``x``. Rows whose segment sum is ``0``
        become ``0`` (no NaN from ``0/0``).
    """
    sum_x = segment_sum(x, segment_ids, num_segments, sorted=sorted)
    sum_x = take(sum_x, segment_ids)
    xp = array_namespace(x, sum_x)
    safe = where(sum_x == 0, xp.ones_like(sum_x), sum_x)
    out = x / safe
    return where(sum_x == 0, xp.zeros_like(out), out)

segment_softmax 

segment_softmax(
    logits, segment_ids, num_segments, sorted=False
)

Softmax within segments (numerically stable).

Subtracts the per-segment max before exp, then normalizes by the per-segment sum of exps — same pattern as a stable full softmax.

Parameters:

Name Type Description Default
logits SegmentValues

Scores. Leading axis is the segment axis.

 required
segment_ids SegmentIds

Integer segment ids (see :func:segment_sum).

 required
num_segments ShapeSize

Required shape-size (Python int, symbolic constant, or 0-d integral tensor) — not inferred from ids.

 required
sorted bool

Forwarded to underlying segment ops (no-op on NumPy/Torch).

 False

Returns:

Type Description
SegmentValues

Array with the same shape as logits. Within each segment, values
SegmentValues

sum to 1 along the segment grouping (empty segments contribute
SegmentValues

nothing useful if referenced via ids).
Notes

Inherits empty-slot segment_max identity (-inf) and TF NaN OR-in behavior from :func:segment_max / :func:segment_sum. Not TorchScript-safe today (needs Python dispatch).

Source code in anytensor/segment.py 
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
def segment_softmax(logits: SegmentValues, segment_ids: SegmentIds, num_segments: ShapeSize, sorted: bool = False) -> SegmentValues:
    """Softmax within segments (numerically stable).

    Subtracts the per-segment max before ``exp``, then normalizes by the
    per-segment sum of exps — same pattern as a stable full softmax.

    Args:
        logits: Scores. Leading axis is the segment axis.
        segment_ids: Integer segment ids (see :func:`segment_sum`).
        num_segments: **Required** shape-size (Python ``int``, symbolic
            constant, or 0-d integral tensor) — not inferred from ids.
        sorted: Forwarded to underlying segment ops (no-op on NumPy/Torch).

    Returns:
        Array with the same shape as ``logits``. Within each segment, values
        sum to ``1`` along the segment grouping (empty segments contribute
        nothing useful if referenced via ids).

    Notes:
        Inherits empty-slot ``segment_max`` identity (``-inf``) and TF NaN
        OR-in behavior from :func:`segment_max` / :func:`segment_sum`.
        Not TorchScript-safe today (needs Python dispatch).
    """
    num_segments = _normalize_shape_dim(num_segments)
    maxs = segment_max(logits, segment_ids, num_segments, sorted=sorted)
    maxs = take(maxs, segment_ids)
    centered = logits - maxs
    exps = exp(centered)
    normalizers = segment_sum(exps, segment_ids, num_segments, sorted=sorted)
    normalizers = take(normalizers, segment_ids)
    return exps / normalizers

segment_min_or_constant 

segment_min_or_constant(
    x, segment_ids, num_segments, constant=0.0, sorted=False
)

Segment min with a finite fill for empty segments.

Like :func:segment_min, but empty slots become constant instead of +inf / dtype max.

Parameters:

Name Type Description Default
x SegmentValues

Values. Leading axis is the segment axis.

 required
segment_ids SegmentIds

Integer segment ids (see :func:segment_sum).

 required
num_segments ShapeSize

Required shape-size.

 required
constant float

Fill for empty segments (default 0.0). Broadcast to the aggregated shape.

 0.0
sorted bool

Forwarded to :func:segment_min (no-op on NumPy/Torch).

 False

Returns:

Type Description
SegmentOut

Array of shape (num_segments,) + x.shape[1:].
Source code in anytensor/segment.py 
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
def segment_min_or_constant(
    x: SegmentValues,
    segment_ids: SegmentIds,
    num_segments: ShapeSize,
    constant: float = 0.0,
    sorted: bool = False,
) -> SegmentOut:
    """Segment min with a finite fill for empty segments.

    Like :func:`segment_min`, but empty slots become ``constant`` instead of
    ``+inf`` / dtype max.

    Args:
        x: Values. Leading axis is the segment axis.
        segment_ids: Integer segment ids (see :func:`segment_sum`).
        num_segments: Required shape-size.
        constant: Fill for empty segments (default ``0.0``). Broadcast to the
            aggregated shape.
        sorted: Forwarded to :func:`segment_min` (no-op on NumPy/Torch).

    Returns:
        Array of shape ``(num_segments,) + x.shape[1:]``.
    """
    num_segments = _normalize_shape_dim(num_segments)
    out = segment_min(x, segment_ids, num_segments, sorted=sorted)
    return _replace_empty_with_constant(out, segment_ids, num_segments, constant, sorted=sorted)

segment_max_or_constant 

segment_max_or_constant(
    x, segment_ids, num_segments, constant=0.0, sorted=False
)

Segment max with a finite fill for empty segments.

Like :func:segment_max, but empty slots become constant instead of -inf / dtype min.

Parameters:

Name Type Description Default
x SegmentValues

Values. Leading axis is the segment axis.

 required
segment_ids SegmentIds

Integer segment ids (see :func:segment_sum).

 required
num_segments ShapeSize

Required shape-size.

 required
constant float

Fill for empty segments (default 0.0). Broadcast to the aggregated shape.

 0.0
sorted bool

Forwarded to :func:segment_max (no-op on NumPy/Torch).

 False

Returns:

Type Description
SegmentOut

Array of shape (num_segments,) + x.shape[1:].
Source code in anytensor/segment.py 
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
def segment_max_or_constant(
    x: SegmentValues,
    segment_ids: SegmentIds,
    num_segments: ShapeSize,
    constant: float = 0.0,
    sorted: bool = False,
) -> SegmentOut:
    """Segment max with a finite fill for empty segments.

    Like :func:`segment_max`, but empty slots become ``constant`` instead of
    ``-inf`` / dtype min.

    Args:
        x: Values. Leading axis is the segment axis.
        segment_ids: Integer segment ids (see :func:`segment_sum`).
        num_segments: Required shape-size.
        constant: Fill for empty segments (default ``0.0``). Broadcast to the
            aggregated shape.
        sorted: Forwarded to :func:`segment_max` (no-op on NumPy/Torch).

    Returns:
        Array of shape ``(num_segments,) + x.shape[1:]``.
    """
    num_segments = _normalize_shape_dim(num_segments)
    out = segment_max(x, segment_ids, num_segments, sorted=sorted)
    return _replace_empty_with_constant(out, segment_ids, num_segments, constant, sorted=sorted)

partition_softmax 

partition_softmax(logits, partitions, total_length)

Softmax within contiguous partitions of lengths partitions.

Convenience: :func:partition_ids then :func:segment_softmax. num_segments is shape(partitions)[0] — not an argument. total_length is required (shape(logits)[0], not a data sum(partitions)). Does not talk to :data:cache itself — :func:partition_ids does, so a cache hit is shared with every partition helper. Ids are rebuilt on every call unless that cache is active. A compiler may CSE the rebuild; eager will not. If you already have ids, call :func:segment_softmax.

Parameters:

Name Type Description Default
logits ShapedArray

Scores aligned with the flattened partitions (length total_length).

 required
partitions IntArray

1-D integer vector of partition sizes. Length is the number of segments.

 required
total_length ShapeSize

Required shape-size for the flattened length (shape(logits)[0], not a data sum(partitions)). Passed to :func:repeat as total_repeat_length.

 required

Returns:

Type Description
ShapedArray

Softmax of logits within each partition (same shape as logits).
Source code in anytensor/segment.py 
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
def partition_softmax(
    logits: ShapedArray,
    partitions: IntArray,
    total_length: ShapeSize,
) -> ShapedArray:
    """Softmax within contiguous partitions of lengths ``partitions``.

    Convenience: :func:`partition_ids` then :func:`segment_softmax`.
    ``num_segments`` is ``shape(partitions)[0]`` — not an argument.
    ``total_length`` is **required** (``shape(logits)[0]``, not a data
    ``sum(partitions)``). Does not talk to :data:`cache` itself —
    :func:`partition_ids` does, so a cache hit is shared with every
    partition helper. **Ids are rebuilt on every call** unless that
    cache is active. A compiler may CSE the rebuild; eager will not.
    If you already have ids, call :func:`segment_softmax`.

    Args:
        logits: Scores aligned with the flattened partitions (length
            ``total_length``).
        partitions: 1-D integer vector of partition sizes. Length is the
            number of segments.
        total_length: **Required** shape-size for the flattened length
            (``shape(logits)[0]``, not a data ``sum(partitions)``). Passed
            to :func:`repeat` as ``total_repeat_length``.

    Returns:
        Softmax of ``logits`` within each partition (same shape as ``logits``).
    """
    return _partition_apply(logits, partitions, total_length, segment_softmax)

partition_sum 

partition_sum(x, partitions, total_length)

Sum within contiguous partitions of lengths partitions.

Convenience: :func:partition_ids then :func:segment_sum. num_segments is shape(partitions)[0] — not an argument. total_length is required (shape(x)[0], not a data sum(partitions)). Does not talk to :data:cache itself — :func:partition_ids does, so a cache hit is shared with every partition helper. If you already have ids, call :func:segment_sum.

Parameters:

Name Type Description Default
x ShapedArray

Values aligned with the flattened partitions (length total_length).

 required
partitions IntArray

1-D integer vector of partition sizes. Length is the number of segments.

 required
total_length ShapeSize

Required shape-size for the flattened length (shape(x)[0], not a data sum(partitions)).

 required

Returns:

Type Description
ShapedArray

Array of shape (shape(partitions)[0],) + x.shape[1:]. Empty
ShapedArray

partitions are 0 (same identity as :func:segment_sum).
Source code in anytensor/segment.py 
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
def partition_sum(
    x: ShapedArray,
    partitions: IntArray,
    total_length: ShapeSize,
) -> ShapedArray:
    """Sum within contiguous partitions of lengths ``partitions``.

    Convenience: :func:`partition_ids` then :func:`segment_sum`.
    ``num_segments`` is ``shape(partitions)[0]`` — not an argument.
    ``total_length`` is **required** (``shape(x)[0]``, not a data
    ``sum(partitions)``). Does not talk to :data:`cache` itself —
    :func:`partition_ids` does, so a cache hit is shared with every
    partition helper. If you already have ids, call :func:`segment_sum`.

    Args:
        x: Values aligned with the flattened partitions (length
            ``total_length``).
        partitions: 1-D integer vector of partition sizes. Length is the
            number of segments.
        total_length: **Required** shape-size for the flattened length
            (``shape(x)[0]``, not a data ``sum(partitions)``).

    Returns:
        Array of shape ``(shape(partitions)[0],) + x.shape[1:]``. Empty
        partitions are ``0`` (same identity as :func:`segment_sum`).
    """
    return _partition_apply(x, partitions, total_length, segment_sum)

partition_min 

partition_min(x, partitions, total_length)

Minimum within contiguous partitions of lengths partitions.

Convenience: :func:partition_ids then :func:segment_min. Same contracts as :func:partition_sum. Empty partitions keep the min identity (+inf / dtype max). Prefer :func:segment_min_or_constant after :func:partition_ids for a finite empty fill.

Source code in anytensor/segment.py 
542
543
544
545
546
547
548
549
550
551
552
553
554
555
def partition_min(
    x: ShapedArray,
    partitions: IntArray,
    total_length: ShapeSize,
) -> ShapedArray:
    """Minimum within contiguous partitions of lengths ``partitions``.

    Convenience: :func:`partition_ids` then :func:`segment_min`. Same
    contracts as :func:`partition_sum`. Empty partitions keep the min
    identity (``+inf`` / dtype max). Prefer
    :func:`segment_min_or_constant` after :func:`partition_ids` for a
    finite empty fill.
    """
    return _partition_apply(x, partitions, total_length, segment_min)

partition_max 

partition_max(x, partitions, total_length)

Maximum within contiguous partitions of lengths partitions.

Convenience: :func:partition_ids then :func:segment_max. Same contracts as :func:partition_sum. Empty partitions keep the max identity (-inf / dtype min). Prefer :func:segment_max_or_constant after :func:partition_ids for a finite empty fill.

Source code in anytensor/segment.py 
558
559
560
561
562
563
564
565
566
567
568
569
570
571
def partition_max(
    x: ShapedArray,
    partitions: IntArray,
    total_length: ShapeSize,
) -> ShapedArray:
    """Maximum within contiguous partitions of lengths ``partitions``.

    Convenience: :func:`partition_ids` then :func:`segment_max`. Same
    contracts as :func:`partition_sum`. Empty partitions keep the max
    identity (``-inf`` / dtype min). Prefer
    :func:`segment_max_or_constant` after :func:`partition_ids` for a
    finite empty fill.
    """
    return _partition_apply(x, partitions, total_length, segment_max)

partition_ids 

partition_ids(partitions, total_length)

Expand partition lengths to segment ids ([0,0,…,1,1,…,n-1]).

This is the conversion other partition helpers call internally. num_segments is shape(partitions)[0] (not an argument; not data-dependent). total_length is the required flattened length (shape(logits)[0], not a data sum(partitions)). Passed to :func:repeat as total_repeat_length.

The only partition helper that talks to :data:cache. Uses :meth:cache.lookup / :meth:cache.store on "partition". Outside the cache, every call rebuilds ids. Inside, the same partitions tensor returns the previous ids from cache["partition"] until the tensor is collected or the block exits — one entry per partition vector. The flattened total is shape(ids)[0] (not a separate sum(partitions) cache); on ONNX export that length is a dim_param. If a cached expansion's length does not match total_length (in-place edit of a 0-d size, or a stale entry), that entry is purged, a warning is issued, and ids are recomputed. The length check uses host Python ints only; tracing skips it. Passing None for total_length is a TypeError.

Source code in anytensor/segment.py 
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
def partition_ids(
    partitions: IntArray,
    total_length: ShapeSize,
) -> IntArray:
    """Expand partition lengths to segment ids (``[0,0,…,1,1,…,n-1]``).

    This is the conversion other partition helpers call internally.
    ``num_segments`` is ``shape(partitions)[0]`` (not an argument; not
    data-dependent). ``total_length`` is the required flattened length
    (``shape(logits)[0]``, not a data ``sum(partitions)``). Passed to
    :func:`repeat` as ``total_repeat_length``.

    The only partition helper that talks to :data:`cache`. Uses
    :meth:`cache.lookup` / :meth:`cache.store` on ``"partition"``. Outside
    the cache, every call rebuilds ids. Inside, the same ``partitions``
    tensor returns the previous ids from ``cache["partition"]`` until
    the tensor is collected or the block exits — **one entry per
    partition vector**. The flattened total is ``shape(ids)[0]`` (not a
    separate ``sum(partitions)`` cache); on ONNX export that length is a
    ``dim_param``. If a cached expansion's length does not match
    ``total_length`` (in-place edit of a 0-d size, or a stale entry),
    that entry is purged, a warning is issued, and ids are recomputed.
    The length check uses host Python ints only; tracing skips it.
    Passing ``None`` for ``total_length`` is a ``TypeError``.
    """
    n_part = shape(partitions)[0]
    total = _require_shape_size("total_length", total_length)
    cached = cache.lookup("partition", partitions)
    if cached is not None:
        stale = _stale_cached_ids(cached, total)
        if stale is not None:
            cache.purge("partition", partitions)
            got, want = stale
            warnings.warn(
                f"cached partition ids length {got} != total_length {want}; "
                "purging and recomputing",
                stacklevel=2,
            )
        else:
            return cached
    ids = repeat(arange(n_part, like=partitions), partitions, total_repeat_length=total)
    cache.store("partition", partitions, ids)
    return ids

enable_torchscript 

enable_torchscript()

Enable torch.jit.script through public segment_sum / min / max.

Wraps those helpers with a torch.jit.is_scripting() divert to :mod:anytensor.torchscript. Eager behavior is unchanged: NumPy, JAX, Torch, and TF tensors still dispatch via backends. Only while scripting (or inside an already-scripted graph) do we take the pure-Torch kernels.

That lets a third-party library call anytensor.segment_sum in ordinary Python, while an end user of that library can torch.jit.script their own code that reaches those calls.

Does not import torch and does not require a particular import order. If Torch is not loaded yet, a helper is registered with :func:anytensor.module_if_loaded and the divert enables on a later import torch. Returns False until then; safe to call more than once.

Source code in anytensor/segment.py 
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
def enable_torchscript() -> bool:
    """Enable ``torch.jit.script`` through public ``segment_sum`` / ``min`` / ``max``.

    Wraps those helpers with a ``torch.jit.is_scripting()`` divert to
    :mod:`anytensor.torchscript`. **Eager behavior is unchanged**: NumPy, JAX,
    Torch, and TF tensors still dispatch via backends. Only while scripting
    (or inside an already-scripted graph) do we take the pure-Torch kernels.

    That lets a third-party library call ``anytensor.segment_sum`` in ordinary
    Python, while an end user of that library can ``torch.jit.script`` their
    own code that reaches those calls.

    Does not import ``torch`` and does not require a particular import order.
    If Torch is not loaded yet, a helper is registered with
    :func:`anytensor.module_if_loaded` and the divert enables on a later
    ``import torch``. Returns ``False`` until then; safe to call more than once.
    """
    if _TORCHSCRIPT_ENABLED:
        return True
    return module_if_loaded("torch", _enable_torchscript) is not None

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 
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
@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)

take 

take(x, indices, axis=0)

Take elements from x along axis (default 0).

Source code in anytensor/core.py 
442
443
444
445
446
@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)

exp 

exp(x)

Element-wise exponential.

Source code in anytensor/core.py 
334
335
336
337
@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 
340
341
342
343
@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 
346
347
348
349
350
351
@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 
354
355
356
357
358
359
360
361
362
363
364
365
@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 
368
369
370
371
372
373
374
375
376
377
378
@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 
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
@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 
399
400
401
402
403
404
405
406
407
408
409
410
411
@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))

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 
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
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)

cumsum 

cumsum(x, axis=0)

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

Source code in anytensor/core.py 
414
415
416
417
418
419
@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)

reshape 

reshape(x, shape)

Reshape x to shape.

Source code in anytensor/core.py 
449
450
451
452
@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 
455
456
457
458
459
460
461
@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 
464
465
466
467
@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 
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
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 
491
492
493
494
@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 
497
498
499
500
501
502
503
504
505
506
507
@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 
510
511
512
513
514
515
516
517
518
519
@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 
522
523
524
525
@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 
528
529
530
531
532
@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 
535
536
537
538
539
@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 
542
543
544
545
546
547
548
549
550
551
552
553
554
@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 
557
558
559
560
@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 
563
564
565
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 
568
569
570
571
572
573
574
575
576
577
@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 
580
581
582
583
584
585
586
587
588
@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 
591
592
593
594
595
596
597
598
599
@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 
602
603
604
605
606
607
608
609
610
@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 
613
614
615
616
617
618
619
620
621
@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 
624
625
626
627
628
629
630
631
632
@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 
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
@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)

matmul 

matmul(x, y)

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

Source code in anytensor/core.py 
821
822
823
824
825
@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 
841
842
843
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 
846
847
848
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 
851
852
853
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 
856
857
858
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 
861
862
863
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 
869
870
871
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 
874
875
876
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 
879
880
881
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 
887
888
889
890
@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 
893
894
895
896
@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 
899
900
901
902
@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 
911
912
913
914
915
916
@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 
922
923
924
925
926
927
928
929
930
931
932
933
934
@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 
940
941
942
943
944
945
946
947
948
@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 
951
952
953
954
955
956
957
958
959
960
961
@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)

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 
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
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 
303
304
305
306
307
308
309
310
311
312
313
314
315
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)

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 
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
@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 
121
122
123
124
125
126
127
128
129
130
131
132
133
134
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
    )

empty_segment_identity 

empty_segment_identity(dtype, reduction, *, xp)

Return the AnyTensor empty-segment identity for reduction on dtype.

This is the portable fill used for empty slots in segment_sum / segment_min / segment_max (and backends that implement them).

Parameters:

Name Type Description Default
dtype Any

Target dtype (NumPy / Array API / TF dtype with as_numpy_dtype).

 required
reduction str

One of "sum", "min", "max".

 required
xp Any

Array API-ish namespace with iinfo, asarray, and floating constants (inf) — e.g. numpy or an array_api_compat ns.

 required

Returns:

Type Description
Any

0 for sum; +inf / iinfo.max for min; -inf /
Any

iinfo.min for max.
Notes

Differs from TF unsorted_segment_{min,max}, which fill empties (and ±inf-only segments) with finfo limits. AnyTensor uses these identities so empties and occupied ±inf match NumPy/JAX/Torch.

Source code in anytensor/semantics.py 
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
def empty_segment_identity(dtype: Any, reduction: str, *, xp: Any) -> Any:
    """Return the AnyTensor empty-segment identity for ``reduction`` on ``dtype``.

    This is the portable fill used for empty slots in ``segment_sum`` /
    ``segment_min`` / ``segment_max`` (and backends that implement them).

    Args:
        dtype: Target dtype (NumPy / Array API / TF dtype with
            ``as_numpy_dtype``).
        reduction: One of ``"sum"``, ``"min"``, ``"max"``.
        xp: Array API-ish namespace with ``iinfo``, ``asarray``, and floating
            constants (``inf``) — e.g. ``numpy`` or an ``array_api_compat`` ns.

    Returns:
        ``0`` for ``sum``; ``+inf`` / ``iinfo.max`` for ``min``; ``-inf`` /
        ``iinfo.min`` for ``max``.

    Notes:
        Differs from TF ``unsorted_segment_{min,max}``, which fill empties
        (and ±inf-only segments) with **finfo** limits. AnyTensor uses these
        identities so empties and occupied ±inf match NumPy/JAX/Torch.
    """
    if reduction == "sum":
        return 0
    if reduction not in ("min", "max"):
        raise ValueError(f"reduction type {reduction} not supported")

    # Prefer Array API / NumPy dtype checks; accept TF dtypes via as_numpy_dtype.
    is_float = False
    try:
        import numpy as np

        np_dtype = getattr(dtype, "as_numpy_dtype", dtype)
        is_float = np.issubdtype(np_dtype, np.floating)
    except Exception:
        kind = getattr(dtype, "kind", None)
        is_float = kind == "f"

    if is_float:
        return xp.inf if reduction == "min" else -xp.inf

    info = xp.iinfo(dtype)
    return info.max if reduction == "min" else info.min

enable_typecheck 

enable_typecheck(typechecker='beartype.beartype')

Opt-in runtime jaxtyping checks for subsequent anytensor imports.

Must be called before importing the modules you want checked (or use :func:jaxtyping.install_import_hook yourself at process start). Default AnyTensor imports do not install this hook.

Hooks core / segment / namespace / semantics / backends / typing — not :mod:anytensor.torchscript, whose divert wrappers must remain torch.jit.script-compilable.

Requires the named typechecker package (default: beartype).

Source code in anytensor/typing.py 
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
def enable_typecheck(typechecker: str = "beartype.beartype") -> None:
    """Opt-in runtime jaxtyping checks for subsequent ``anytensor`` imports.

    Must be called **before** importing the modules you want checked (or use
    :func:`jaxtyping.install_import_hook` yourself at process start). Default
    AnyTensor imports do **not** install this hook.

    Hooks ``core`` / ``segment`` / ``namespace`` / ``semantics`` / ``backends`` /
    ``typing`` — not :mod:`anytensor.torchscript`, whose divert wrappers must
    remain ``torch.jit.script``-compilable.

    Requires the named typechecker package (default: ``beartype``).
    """
    from jaxtyping import install_import_hook

    install_import_hook(
        [
            "anytensor.core",
            "anytensor.segment",
            "anytensor.namespace",
            "anytensor.semantics",
            "anytensor.backends",
            "anytensor.typing",
        ],
        typechecker,
    )