mixle.engines package

Public compute-engine registry and precision front door for Mixle kernels.

This namespace exposes the NumPy, Torch, JAX, symbolic, and precision-aware engine utilities used by backend-neutral scoring, estimation, and data-transfer paths.

class ComputeEngine[source]

Bases: ABC

Small array-backend interface for numpy/torch/etc.

Engines own arithmetic policy: array library, device, dtype, and optional compilation. Distribution and kernel code should depend only on this surface when it wants backend-neutral arrays.

is_concrete(x)[source]

Return whether x holds real values now, so host-side inspection of it is meaningful.

Every eager backend answers True. A tracing backend (JAX under jit) answers False for a traced value: its contents are not available in Python, so to_numpy raises and any control flow branching on them is impossible by construction. Shared code that wants to run an optional host-side diagnostic – scanning for NaN rows to name them in an exception, say – asks first and skips the diagnostic while tracing; the engine math itself is unaffected either way.

Parameters:

x (Any)

Return type:

bool

constant(value)[source]

Return value in this engine’s scalar representation (identity for numeric engines).

Parameters:

value (Any)

Return type:

Any

property accumulator_dtype: Any

High-precision dtype for sufficient-statistic reductions, or None when not applicable.

Numeric engines override this with their float64 accumulator so a reduced-precision fit does not drift on large N (see NumpyEngine/TorchEngine). The base returns None: meaning “no separate accumulator dtype”, which is the correct policy for engines that never drive the numeric accumulate path (e.g. the symbolic engine, where reductions are exact expression trees). None is also a valid dtype= argument to sum (NumPy’s default).

property precision: str

Return the engine dtype policy as a stable user-facing name.

with_precision(precision)[source]

Return an equivalent engine with a different floating-point policy.

Parameters:

precision (Any)

Return type:

ComputeEngine

abstractmethod asarray(x, dtype=None)[source]

Convert x into this engine’s array/tensor representation.

Parameters:
Return type:

Any

abstractmethod zeros(shape, dtype=None)[source]

Allocate a zero-filled array on this engine.

Parameters:
Return type:

Any

abstractmethod empty(shape, dtype=None)[source]

Allocate an uninitialized array on this engine.

Parameters:
Return type:

Any

abstractmethod arange(*args, **kwargs)[source]

Return an evenly spaced one-dimensional array on this engine.

Parameters:
Return type:

Any

abstractmethod to_numpy(x)[source]

Move an engine array back to a NumPy/host representation.

Parameters:

x (Any)

Return type:

Any

abstractmethod stack(arrays, axis=0)[source]

Stack a sequence of arrays along axis.

Parameters:
Return type:

Any

abstractmethod concatenate(arrays, axis=0)[source]

Join a sequence of arrays along an existing axis.

Parameters:
Return type:

Any

requires_grad(x)[source]

Return whether x participates in this engine’s autograd graph.

Parameters:

x (Any)

Return type:

bool

compile(fn)[source]

Optionally compile fn; engines without a compiler return it unchanged.

Parameters:

fn (Callable)

Return type:

Callable

replicate(x)[source]

Return x in the engine’s replicated placement, when applicable.

Parameters:

x (Any)

Return type:

Any

place_component_axis(x, axis=0)[source]

Return x with a component-axis placement, when the engine supports it.

Parameters:
Return type:

Any

class NumpyEngine(dtype=None, prefer_fused=False)[source]

Bases: ComputeEngine

Host NumPy/SciPy engine used by the default local execution path.

Parameters:

  • dtype (Any)

  • prefer_fused (bool)

property accumulator_dtype: Any

High-precision dtype for sufficient-statistic reductions (always float64).

Reductions that aggregate over observations (the sum/index-add forming sufficient statistics) accumulate in this dtype even when scoring runs in reduced precision, so a float32 fit does not drift on large N (e.g. catastrophic cancellation in a variance).

with_precision(precision)[source]

Return a NumPy engine using precision for floating arrays.

Parameters:

precision (Any)

Return type:

NumpyEngine

asarray(x, dtype=None)[source]

Convert x to a NumPy ndarray under the engine dtype policy.

Parameters:
Return type:

ndarray

zeros(shape, dtype=None)[source]

Allocate a NumPy zero array using the configured float dtype.

Parameters:
Return type:

ndarray

empty(shape, dtype=None)[source]

Allocate an uninitialized NumPy array using the configured float dtype.

Parameters:
Return type:

ndarray

arange(*args, **kwargs)[source]

Return np.arange with float ranges honoring the precision policy.

Parameters:
Return type:

ndarray

to_numpy(x)[source]

Return x as a host NumPy array.

Parameters:

x (Any)

Return type:

ndarray

stack(arrays, axis=0)[source]

Stack arrays with np.stack along the requested axis.

Parameters:
Return type:

ndarray

concatenate(arrays, axis=0)[source]

Join arrays with np.concatenate along the requested axis.

Parameters:
Return type:

ndarray

log = <ufunc 'log'>
exp = <ufunc 'exp'>
sqrt = <ufunc 'sqrt'>
abs = <ufunc 'absolute'>
static where(condition, [x, y, ]/)

Return elements chosen from x or y depending on condition.

Note

When only condition is provided, this function is a shorthand for np.asarray(condition).nonzero(). Using nonzero directly should be preferred, as it behaves correctly for subclasses. The rest of this documentation covers only the case where all three arguments are provided.

Parameters:

  • condition (array_like, bool) – Where True, yield x, otherwise yield y.

  • x (array_like) – Values from which to choose. x, y and condition need to be broadcastable to some shape.

  • y (array_like) – Values from which to choose. x, y and condition need to be broadcastable to some shape.

Returns:

out – An array with elements from x where condition is True, and elements from y elsewhere.

Return type:

ndarray

See also

choose

nonzero

The function that is called when x and y are omitted

Notes

If all the arrays are 1-D, where is equivalent to:

[xv if c else yv
 for c, xv, yv in zip(condition, x, y)]

Examples

>>> import numpy as np
>>> a = np.arange(10)
>>> a
array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
>>> np.where(a < 5, a, 10*a)
array([ 0,  1,  2,  3,  4, 50, 60, 70, 80, 90])

This can be used on multidimensional arrays too:

>>> np.where([[True, False], [True, True]],
...          [[1, 2], [3, 4]],
...          [[9, 8], [7, 6]])
array([[1, 8],
       [3, 4]])

The shapes of x, y, and the condition are broadcast together:

>>> x, y = np.ogrid[:3, :4]
>>> np.where(x < y, x, 10 + y)  # both x and 10+y are broadcast
array([[10,  0,  0,  0],
       [10, 11,  1,  1],
       [10, 11, 12,  2]])

>>> a = np.array([[0, 1, 2],
...               [0, 2, 4],
...               [0, 3, 6]])
>>> np.where(a < 4, a, -1)  # -1 is broadcast
array([[ 0,  1,  2],
       [ 0,  2, -1],
       [ 0,  3, -1]])
maximum = <ufunc 'maximum'>
static clip(a, a_min=np._NoValue, a_max=np._NoValue, out=None, *, min=np._NoValue, max=np._NoValue, **kwargs)

Clip (limit) the values in an array.

Given an interval, values outside the interval are clipped to the interval edges. For example, if an interval of [0, 1] is specified, values smaller than 0 become 0, and values larger than 1 become 1.

Equivalent to but faster than np.minimum(a_max, np.maximum(a, a_min)).

No check is performed to ensure a_min < a_max.

Parameters:

  • a (array_like) – Array containing elements to clip.

  • a_min (array_like or None) – Minimum and maximum value. If None, clipping is not performed on the corresponding edge. If both a_min and a_max are None, the elements of the returned array stay the same. Both are broadcasted against a.

  • a_max (array_like or None) – Minimum and maximum value. If None, clipping is not performed on the corresponding edge. If both a_min and a_max are None, the elements of the returned array stay the same. Both are broadcasted against a.

  • out (ndarray, optional) – The results will be placed in this array. It may be the input array for in-place clipping. out must be of the right shape to hold the output. Its type is preserved.

  • min (array_like or None) –

    Array API compatible alternatives for a_min and a_max arguments. Either a_min and a_max or min and max can be passed at the same time. Default: None.

    Added in version 2.1.0.

  • max (array_like or None) –

    Array API compatible alternatives for a_min and a_max arguments. Either a_min and a_max or min and max can be passed at the same time. Default: None.

    Added in version 2.1.0.

  • **kwargs – For other keyword-only arguments, see the ufunc docs.

Returns:

clipped_array – An array with the elements of a, but where values < a_min are replaced with a_min, and those > a_max with a_max.

Return type:

ndarray

See also

Output type determination

Notes

When a_min is greater than a_max, clip returns an array in which all values are equal to a_max, as shown in the second example.

Examples

>>> import numpy as np
>>> a = np.arange(10)
>>> a
array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
>>> np.clip(a, 1, 8)
array([1, 1, 2, 3, 4, 5, 6, 7, 8, 8])
>>> np.clip(a, 8, 1)
array([1, 1, 1, 1, 1, 1, 1, 1, 1, 1])
>>> np.clip(a, 3, 6, out=a)
array([3, 3, 3, 3, 4, 5, 6, 6, 6, 6])
>>> a
array([3, 3, 3, 3, 4, 5, 6, 6, 6, 6])
>>> a = np.arange(10)
>>> a
array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
>>> np.clip(a, [3, 4, 1, 1, 1, 4, 4, 4, 4, 4], 8)
array([3, 4, 2, 3, 4, 5, 6, 7, 8, 8])
floor = <ufunc 'floor'>
isnan = <ufunc 'isnan'>
isinf = <ufunc 'isinf'>
sum(a, *args, dtype=None, **kwargs)[source]

Reduce with np.sum, accumulating floats in accumulator_dtype by default.

A float32 np.sum drifts on large N; when the caller passes no explicit dtype we promote floating inputs to the engine’s high-precision accumulator (float64) so correctness no longer relies on every caller threading dtype=accumulator_dtype. Integer inputs and explicit-dtype calls keep NumPy’s default behavior.

Parameters:
Return type:

Any

static max(a, axis=None, out=None, keepdims=np._NoValue, initial=np._NoValue, where=np._NoValue)

Return the maximum of an array or maximum along an axis.

Parameters:

  • a (array_like) – Input data.

  • axis (None or int or tuple of ints, optional) – Axis or axes along which to operate. By default, flattened input is used. If this is a tuple of ints, the maximum is selected over multiple axes, instead of a single axis or all the axes as before.

  • out (ndarray, optional) – Alternative output array in which to place the result. Must be of the same shape and buffer length as the expected output. See Output type determination for more details.

  • keepdims (bool, optional) –

    If this is set to True, the axes which are reduced are left in the result as dimensions with size one. With this option, the result will broadcast correctly against the input array.

    If the default value is passed, then keepdims will not be passed through to the max method of sub-classes of ndarray, however any non-default value will be. If the sub-class’ method does not implement keepdims any exceptions will be raised.

  • initial (scalar, optional) – The minimum value of an output element. Must be present to allow computation on empty slice. See ~numpy.ufunc.reduce for details.

  • where (array_like of bool, optional) – Elements to compare for the maximum. See ~numpy.ufunc.reduce for details.

Returns:

max – Maximum of a. If axis is None, the result is a scalar value. If axis is an int, the result is an array of dimension a.ndim - 1. If axis is a tuple, the result is an array of dimension a.ndim - len(axis).

Return type:

ndarray or scalar

See also

amin

The minimum value of an array along a given axis, propagating any NaNs.

nanmax

The maximum value of an array along a given axis, ignoring any NaNs.

maximum

Element-wise maximum of two arrays, propagating any NaNs.

fmax

Element-wise maximum of two arrays, ignoring any NaNs.

argmax

Return the indices of the maximum values.

nanmin, minimum, fmin

Notes

NaN values are propagated, that is if at least one item is NaN, the corresponding max value will be NaN as well. To ignore NaN values (MATLAB behavior), please use nanmax.

Don’t use ~numpy.max for element-wise comparison of 2 arrays; when a.shape[0] is 2, maximum(a[0], a[1]) is faster than max(a, axis=0).

Examples

>>> import numpy as np
>>> a = np.arange(4).reshape((2,2))
>>> a
array([[0, 1],
       [2, 3]])
>>> np.max(a)           # Maximum of the flattened array
3
>>> np.max(a, axis=0)   # Maxima along the first axis
array([2, 3])
>>> np.max(a, axis=1)   # Maxima along the second axis
array([1, 3])
>>> np.max(a, where=[False, True], initial=-1, axis=0)
array([-1,  3])
>>> b = np.arange(5, dtype=float)
>>> b[2] = np.nan
>>> np.max(b)
np.float64(nan)
>>> np.max(b, where=~np.isnan(b), initial=-1)
4.0
>>> np.nanmax(b)
4.0

You can use an initial value to compute the maximum of an empty slice, or to initialize it to a different value:

>>> np.max([[-50], [10]], axis=-1, initial=0)
array([ 0, 10])

Notice that the initial value is used as one of the elements for which the maximum is determined, unlike for the default argument Python’s max function, which is only used for empty iterables.

>>> np.max([5], initial=6)
6
>>> max([5], default=6)
5
static dot(a, b, out=None)

Dot product of two arrays. Specifically,

  • If both a and b are 1-D arrays, it is inner product of vectors (without complex conjugation).

  • If both a and b are 2-D arrays, it is matrix multiplication, but using matmul() or a @ b is preferred.

  • If either a or b is 0-D (scalar), it is equivalent to multiply() and using numpy.multiply(a, b) or a * b is preferred.

  • If a is an N-D array and b is a 1-D array, it is a sum product over the last axis of a and b.

  • If a is an N-D array and b is an M-D array (where M>=2), it is a sum product over the last axis of a and the second-to-last axis of b:

    dot(a, b)[i,j,k,m] = sum(a[i,j,:] * b[k,:,m])

It uses an optimized BLAS library when possible (see numpy.linalg).

Parameters:

  • a (array_like) – First argument.

  • b (array_like) – Second argument.

  • out (ndarray, optional) – Output argument. This must have the exact kind that would be returned if it was not used. In particular, it must have the right type, must be C-contiguous, and its dtype must be the dtype that would be returned for dot(a,b). This is a performance feature. Therefore, if these conditions are not met, an exception is raised, instead of attempting to be flexible.

Returns:

output – Returns the dot product of a and b. If a and b are both scalars or both 1-D arrays then a scalar is returned; otherwise an array is returned. If out is given, then it is returned.

Return type:

ndarray

Raises:

ValueError – If the last dimension of a is not the same size as the second-to-last dimension of b.

See also

vdot

Complex-conjugating dot product.

vecdot

Vector dot product of two arrays.

tensordot

Sum products over arbitrary axes.

einsum

Einstein summation convention.

matmul

‘@’ operator as method with out parameter.

linalg.multi_dot

Chained dot product.

Examples

>>> import numpy as np
>>> np.dot(3, 4)
12

Neither argument is complex-conjugated:

>>> np.dot([2j, 3j], [2j, 3j])
(-13+0j)

For 2-D arrays it is the matrix product:

>>> a = [[1, 0], [0, 1]]
>>> b = [[4, 1], [2, 2]]
>>> np.dot(a, b)
array([[4, 1],
       [2, 2]])

>>> a = np.arange(3*4*5*6).reshape((3,4,5,6))
>>> b = np.arange(3*4*5*6)[::-1].reshape((5,4,6,3))
>>> np.dot(a, b)[2,3,2,1,2,2]
499128
>>> sum(a[2,3,2,:] * b[1,2,:,2])
499128
matmul = <ufunc 'matmul'>
static cumsum(a, axis=None, dtype=None, out=None)

Return the cumulative sum of the elements along a given axis.

Parameters:

  • a (array_like) – Input array.

  • axis (int, optional) – Axis along which the cumulative sum is computed. The default (None) is to compute the cumsum over the flattened array.

  • dtype (dtype, optional) – Type of the returned array and of the accumulator in which the elements are summed. If dtype is not specified, it defaults to the dtype of a, unless a has an integer dtype with a precision less than that of the default platform integer. In that case, the default platform integer is used.

  • out (ndarray, optional) – Alternative output array in which to place the result. It must have the same shape and buffer length as the expected output but the type will be cast if necessary. See Output type determination for more details.

Returns:

cumsum_along_axis – A new array holding the result is returned unless out is specified, in which case a reference to out is returned. The result has the same size as a, and the same shape as a if axis is not None or a is a 1-d array.

Return type:

ndarray.

See also

cumulative_sum

Array API compatible alternative for cumsum.

sum

Sum array elements.

trapezoid

Integration of array values using composite trapezoidal rule.

diff

Calculate the n-th discrete difference along given axis.

Notes

Arithmetic is modular when using integer types, and no error is raised on overflow.

cumsum(a)[-1] may not be equal to sum(a) for floating-point values since sum may use a pairwise summation routine, reducing the roundoff-error. See sum for more information.

Examples

>>> import numpy as np
>>> a = np.array([[1,2,3], [4,5,6]])
>>> a
array([[1, 2, 3],
       [4, 5, 6]])
>>> np.cumsum(a)
array([ 1,  3,  6, 10, 15, 21])
>>> np.cumsum(a, dtype=float)     # specifies type of output value(s)
array([  1.,   3.,   6.,  10.,  15.,  21.])

>>> np.cumsum(a,axis=0)      # sum over rows for each of the 3 columns
array([[1, 2, 3],
       [5, 7, 9]])
>>> np.cumsum(a,axis=1)      # sum over columns for each of the 2 rows
array([[ 1,  3,  6],
       [ 4,  9, 15]])

cumsum(b)[-1] may not be equal to sum(b)

>>> b = np.array([1, 2e-9, 3e-9] * 1000000)
>>> b.cumsum()[-1]
1000000.0050045159
>>> b.sum()
1000000.0050000029
static logsumexp(a, axis=None, b=None, keepdims=False, return_sign=False)

Compute the log of the sum of exponentials of input elements.

Parameters:

  • a (array_like) – Input array.

  • axis (None or int or tuple of ints, optional) –

    Axis or axes over which the sum is taken. By default axis is None, and all elements are summed.

    Added in version 0.11.0.

  • b (array-like, optional) –

    Scaling factor for exp(a) must be of the same shape as a or broadcastable to a. These values may be negative in order to implement subtraction.

    Added in version 0.12.0.

  • keepdims (bool, optional) –

    If this is set to True, the axes which are reduced are left in the result as dimensions with size one. With this option, the result will broadcast correctly against the original array.

    Added in version 0.15.0.

  • return_sign (bool, optional) –

    If this is set to True, the result will be a pair containing sign information; if False, results that are negative will be returned as NaN. Default is False (no sign information).

    Added in version 0.16.0.

Returns:

  • res (ndarray) – The result, np.log(np.sum(np.exp(a))) calculated in a numerically more stable way. If b is given then np.log(np.sum(b*np.exp(a))) is returned. If return_sign is True, res contains the log of the absolute value of the argument.

  • sgn (ndarray) – If return_sign is True, this will be an array of floating-point numbers matching res containing +1, 0, -1 (for real-valued inputs) or a complex phase (for complex inputs). This gives the sign of the argument of the logarithm in res. If return_sign is False, only one result is returned.

See also

numpy.logaddexp
numpy.logaddexp2

Notes

NumPy has a logaddexp function which is very similar to logsumexp, but only handles two arguments. logaddexp.reduce is similar to this function, but may be less stable.

The logarithm is a multivalued function: for each \(x\) there is an infinite number of \(z\) such that \(exp(z) = x\). The convention is to return the \(z\) whose imaginary part lies in \((-pi, pi]\).

Array API Standard Support

logsumexp has experimental support for Python Array API Standard compatible backends in addition to NumPy. Please consider testing these features by setting an environment variable SCIPY_ARRAY_API=1 and providing CuPy, PyTorch, JAX, or Dask arrays as array arguments. The following combinations of backend and device (or other capability) are supported.

Library

CPU

GPU

NumPy

✅

n/a

CuPy

n/a

✅

PyTorch

✅

✅

JAX

✅

✅

Dask

✅

n/a

See Support for the array API standard for more information.

Examples

>>> import numpy as np
>>> from scipy.special import logsumexp
>>> a = np.arange(10)
>>> logsumexp(a)
9.4586297444267107
>>> np.log(np.sum(np.exp(a)))
9.4586297444267107

With weights

>>> a = np.arange(10)
>>> b = np.arange(10, 0, -1)
>>> logsumexp(a, b=b)
9.9170178533034665
>>> np.log(np.sum(b*np.exp(a)))
9.9170178533034647

Returning a sign flag

>>> logsumexp([1,2],b=[1,-1],return_sign=True)
(1.5413248546129181, -1.0)

Notice that logsumexp does not directly support masked arrays. To use it on a masked array, convert the mask into zero weights:

>>> a = np.ma.array([np.log(2), 2, np.log(3)],
...                  mask=[False, True, False])
>>> b = (~a.mask).astype(int)
>>> logsumexp(a.data, b=b), np.log(5)
1.6094379124341005, 1.6094379124341005
static bincount(x, /, weights=None, minlength=0)

Count number of occurrences of each value in array of non-negative ints.

The number of bins (of size 1) is one larger than the largest value in x. If minlength is specified, there will be at least this number of bins in the output array (though it will be longer if necessary, depending on the contents of x). Each bin gives the number of occurrences of its index value in x. If weights is specified the input array is weighted by it, i.e. if a value n is found at position i, out[n] += weight[i] instead of out[n] += 1.

Parameters:

  • x (array_like, 1 dimension, nonnegative ints) – Input array.

  • weights (array_like, optional) – Weights, array of the same shape as x.

  • minlength (int, optional) – A minimum number of bins for the output array.

Returns:

out – The result of binning the input array. The length of out is equal to np.amax(x)+1.

Return type:

ndarray of ints

Raises:

  • ValueError – If the input is not 1-dimensional, or contains elements with negative values, or if minlength is negative.

  • TypeError – If the type of the input is float or complex.

See also

histogram, digitize, unique

Examples

>>> import numpy as np
>>> np.bincount(np.arange(5))
array([1, 1, 1, 1, 1])
>>> np.bincount(np.array([0, 1, 1, 3, 2, 1, 7]))
array([1, 3, 1, 1, 0, 0, 0, 1])

>>> x = np.array([0, 1, 1, 3, 2, 1, 7, 23])
>>> np.bincount(x).size == np.amax(x)+1
True

The input array needs to be of integer dtype, otherwise a TypeError is raised:

>>> np.bincount(np.arange(5, dtype=float))
Traceback (most recent call last):
  ...
TypeError: Cannot cast array data from dtype('float64') to dtype('int64')
according to the rule 'safe'

A possible use of bincount is to perform sums over variable-size chunks of an array, using the weights keyword.

>>> w = np.array([0.3, 0.5, 0.2, 0.7, 1., -0.6]) # weights
>>> x = np.array([0, 1, 1, 2, 2, 2])
>>> np.bincount(x,  weights=w)
array([ 0.3,  0.7,  1.1])
static unique(ar, return_index=False, return_inverse=False, return_counts=False, axis=None, *, equal_nan=True, sorted=True)

Find the unique elements of an array.

Returns the sorted unique elements of an array. There are three optional outputs in addition to the unique elements:

  • the indices of the input array that give the unique values

  • the indices of the unique array that reconstruct the input array

  • the number of times each unique value comes up in the input array

Parameters:

  • ar (array_like) – Input array. Unless axis is specified, this will be flattened if it is not already 1-D.

  • return_index (bool, optional) – If True, also return the indices of ar (along the specified axis, if provided, or in the flattened array) that result in the unique array.

  • return_inverse (bool, optional) – If True, also return the indices of the unique array (for the specified axis, if provided) that can be used to reconstruct ar.

  • return_counts (bool, optional) – If True, also return the number of times each unique item appears in ar.

  • axis (int or None, optional) – The axis to operate on. If None, ar will be flattened. If an integer, the subarrays indexed by the given axis will be flattened and treated as the elements of a 1-D array with the dimension of the given axis, see the notes for more details. Object arrays or structured arrays that contain objects are not supported if the axis kwarg is used. The default is None.

  • equal_nan (bool, optional) –

    If True, collapses multiple NaN values in the return array into one.

    Added in version 1.24.

  • sorted (bool, optional) –

    If True, the unique elements are sorted. Elements may be sorted in practice even if sorted=False, but this could change without notice.

    Added in version 2.3.

Returns:

  • unique (ndarray) – The sorted unique values.

  • unique_indices (ndarray, optional) – The indices of the first occurrences of the unique values in the original array. Only provided if return_index is True.

  • unique_inverse (ndarray, optional) – The indices to reconstruct the original array from the unique array. Only provided if return_inverse is True.

  • unique_counts (ndarray, optional) – The number of times each of the unique values comes up in the original array. Only provided if return_counts is True.

See also

repeat

Repeat elements of an array.

sort

Return a sorted copy of an array.

Notes

When an axis is specified the subarrays indexed by the axis are sorted. This is done by making the specified axis the first dimension of the array (move the axis to the first dimension to keep the order of the other axes) and then flattening the subarrays in C order. The flattened subarrays are then viewed as a structured type with each element given a label, with the effect that we end up with a 1-D array of structured types that can be treated in the same way as any other 1-D array. The result is that the flattened subarrays are sorted in lexicographic order starting with the first element.

Changed in version 1.21: Like np.sort, NaN will sort to the end of the values. For complex arrays all NaN values are considered equivalent (no matter whether the NaN is in the real or imaginary part). As the representant for the returned array the smallest one in the lexicographical order is chosen - see np.sort for how the lexicographical order is defined for complex arrays.

Changed in version 2.0: For multi-dimensional inputs, unique_inverse is reshaped such that the input can be reconstructed using np.take(unique, unique_inverse, axis=axis). The result is now not 1-dimensional when axis=None.

Note that in NumPy 2.0.0 a higher dimensional array was returned also when axis was not None. This was reverted, but inverse.reshape(-1) can be used to ensure compatibility with both versions.

Examples

>>> import numpy as np
>>> np.unique([1, 1, 2, 2, 3, 3])
array([1, 2, 3])
>>> a = np.array([[1, 1], [2, 3]])
>>> np.unique(a)
array([1, 2, 3])

Return the unique rows of a 2D array

>>> a = np.array([[1, 0, 0], [1, 0, 0], [2, 3, 4]])
>>> np.unique(a, axis=0)
array([[1, 0, 0], [2, 3, 4]])

Return the indices of the original array that give the unique values:

>>> a = np.array(['a', 'b', 'b', 'c', 'a'])
>>> u, indices = np.unique(a, return_index=True)
>>> u
array(['a', 'b', 'c'], dtype='<U1')
>>> indices
array([0, 1, 3])
>>> a[indices]
array(['a', 'b', 'c'], dtype='<U1')

Reconstruct the input array from the unique values and inverse:

>>> a = np.array([1, 2, 6, 4, 2, 3, 2])
>>> u, indices = np.unique(a, return_inverse=True)
>>> u
array([1, 2, 3, 4, 6])
>>> indices
array([0, 1, 4, 3, 1, 2, 1])
>>> u[indices]
array([1, 2, 6, 4, 2, 3, 2])

Reconstruct the input values from the unique values and counts:

>>> a = np.array([1, 2, 6, 4, 2, 3, 2])
>>> values, counts = np.unique(a, return_counts=True)
>>> values
array([1, 2, 3, 4, 6])
>>> counts
array([1, 3, 1, 1, 1])
>>> np.repeat(values, counts)
array([1, 2, 2, 2, 3, 4, 6])    # original order not preserved
static searchsorted(a, v, side='left', sorter=None)

Find indices where elements should be inserted to maintain order.

Find the indices into a sorted array a such that, if the corresponding elements in v were inserted before the indices, the order of a would be preserved.

Assuming that a is sorted:

side

returned index i satisfies

left

a[i-1] < v <= a[i]

right

a[i-1] <= v < a[i]
Parameters:

  • a (1-D array_like) – Input array. If sorter is None, then it must be sorted in ascending order, otherwise sorter must be an array of indices that sort it.

  • v (array_like) – Values to insert into a.

  • side ({'left', 'right'}, optional) – If ‘left’, the index of the first suitable location found is given. If ‘right’, return the last such index. If there is no suitable index, return either 0 or N (where N is the length of a).

  • sorter (1-D array_like, optional) – Optional array of integer indices that sort array a into ascending order. They are typically the result of argsort.

Returns:

indices – Array of insertion points with the same shape as v, or an integer if v is a scalar.

Return type:

int or array of ints

See also

sort

Return a sorted copy of an array.

histogram

Produce histogram from 1-D data.

Notes

Binary search is used to find the required insertion points.

As of NumPy 1.4.0 searchsorted works with real/complex arrays containing nan values. The enhanced sort order is documented in sort.

This function uses the same algorithm as the builtin python bisect.bisect_left (side='left') and bisect.bisect_right (side='right') functions, which is also vectorized in the v argument.

Examples

>>> import numpy as np
>>> np.searchsorted([11,12,13,14,15], 13)
2
>>> np.searchsorted([11,12,13,14,15], 13, side='right')
3
>>> np.searchsorted([11,12,13,14,15], [-10, 20, 12, 13])
array([0, 5, 1, 2])

When sorter is used, the returned indices refer to the sorted array of a and not a itself:

>>> a = np.array([40, 10, 20, 30])
>>> sorter = np.argsort(a)
>>> sorter
array([1, 2, 3, 0])  # Indices that would sort the array 'a'
>>> result = np.searchsorted(a, 25, sorter=sorter)
>>> result
2
>>> a[sorter[result]]
30  # The element at index 2 of the sorted array is 30.
gammaln = <ufunc 'gammaln'>
digamma = <ufunc 'psi'>
betaln = <ufunc 'betaln'>
erf = <ufunc 'erf'>
cos = <ufunc 'cos'>
sin = <ufunc 'sin'>
arctan2 = <ufunc 'arctan2'>
i0e = <ufunc 'i0e'>
erfcx = <ufunc 'erfcx'>
index_add(out, index, values)[source]

Add values into out at index using NumPy’s add.at semantics.

Parameters:
Return type:

ndarray

class SymbolicEngine[source]

Bases: ComputeEngine

Small symbolic expression engine over scalar nodes and object arrays.

with_precision(precision)[source]

Return this engine unchanged: symbolic expressions carry no float precision policy.

The numeric engines swap their float dtype here, but symbolic nodes are exact expression trees with no reduced-precision representation, so precision adjustment is a no-op rather than an error – this lets backend-neutral code call with_precision uniformly across engines.

Parameters:

precision (Any)

Return type:

SymbolicEngine

constant(value)[source]

Return value as a symbolic constant node.

Parameters:

value (Any)

Return type:

SymbolicExpression

symbol(name)[source]

Return a named symbolic expression variable.

Parameters:

name (str)

Return type:

SymbolicExpression

asarray(x, dtype=None)[source]

Convert scalars/arrays/strings into symbolic expression objects.

Parameters:
Return type:

Any

zeros(shape, dtype=None)[source]

Return an object array filled with symbolic zero constants.

Parameters:
Return type:

Any

empty(shape, dtype=None)[source]

Return an uninitialized object array for symbolic expressions.

Parameters:
Return type:

Any

arange(*args, **kwargs)[source]

Return symbolic constants corresponding to np.arange values.

Parameters:
Return type:

Any

to_numpy(x)[source]

Return x as a NumPy object array without numeric evaluation.

Parameters:

x (Any)

Return type:

Any

evaluate(x, values)[source]

Evaluate a scalar expression or object-array expression tree.

Parameters:
Return type:

Any

symbols(x)[source]

Return sorted symbolic variable names referenced by x.

Parameters:

x (Any)

Return type:

tuple[str, …]

op_counts(x)[source]

Return aggregate operation counts over a scalar or array expression.

Parameters:

x (Any)

Return type:

dict[str, int]

diagnostics(x)[source]

Return a compact diagnostic summary for generated-kernel inspection.

Parameters:

x (Any)

Return type:

dict[str, Any]

stack(arrays, axis=0)[source]

Stack symbolic arrays with NumPy object-array semantics.

Parameters:
Return type:

Any

concatenate(arrays, axis=0)[source]

Join symbolic arrays with NumPy object-array semantics.

Parameters:
Return type:

Any

static sum(x, axis=None, dtype=None, keepdims=False, initial=None, where=None)[source]

Return a symbolic sum reduction over axis, matching NumPy’s sum contract.

keepdims, initial, and where are implemented exactly as np.sum defines them: reduced axes are kept as size-1 dimensions, initial seeds each reduction group (the additive identity 0.0 when omitted), and where excludes masked-out elements by substituting the additive identity in their place before folding. dtype has no symbolic analogue – an expression node carries no numeric storage to cast – so a non-None value is rejected rather than silently ignored (matching the “reject unsupported arguments” half of the declared reduction contract, since honoring it is not meaningful here).

Parameters:
Return type:

Any

static logsumexp(x, axis=None, keepdims=False)[source]

Return a symbolic log-sum-exp reduction over axis, matching SciPy’s logsumexp.

keepdims is implemented exactly as scipy.special.logsumexp defines it. SciPy’s own logsumexp (which NumpyEngine.logsumexp forwards to directly) has no dtype/initial/where parameter either, so there is no declared behavior to honor for them here; passing them raises the same TypeError SciPy’s own function would, for the same reason (an unexpected keyword argument). An all--inf input (every term has zero probability in log-space) correctly evaluates to -inf, not NaN – see _logsumexp_values(). An empty reduction is likewise defined as -inf (the identity: the log of a sum of zero terms), matching SciPy.

Parameters:
Return type:

Any

static max(x, axis=None, keepdims=False, initial=None, where=None)[source]

Return a symbolic max reduction over axis, matching NumPy’s max contract.

keepdims and initial are implemented exactly as np.max defines them (np.max itself has no dtype parameter, so none is declared here either). Max has no universal identity element, so – exactly like np.max – passing where without also passing initial is rejected rather than silently reducing an under-specified group. An empty reduction with no initial raises (there is nothing to seed it with), matching np.max; with initial it returns initial.

Parameters:
Return type:

Any

static where(cond, x, y)[source]

Return a symbolic elementwise conditional expression.

Parameters:
Return type:

Any

static maximum(x, y)[source]

Return a symbolic elementwise maximum expression.

Parameters:
Return type:

Any

static less(x, y)[source]

Return a symbolic less-than comparison.

Parameters:
Return type:

Any

static less_equal(x, y)[source]

Return a symbolic less-than-or-equal comparison.

Parameters:
Return type:

Any

static greater(x, y)[source]

Return a symbolic greater-than comparison.

Parameters:
Return type:

Any

static greater_equal(x, y)[source]

Return a symbolic greater-than-or-equal comparison.

Parameters:
Return type:

Any

static equal(x, y)[source]

Return a symbolic equality comparison.

Parameters:
Return type:

Any

static not_equal(x, y)[source]

Return a symbolic inequality comparison.

Parameters:
Return type:

Any

static logical_and(x, y)[source]

Return a symbolic elementwise logical-and expression.

Parameters:
Return type:

Any

static logical_or(x, y)[source]

Return a symbolic elementwise logical-or expression.

Parameters:
Return type:

Any

static logical_not(x)[source]

Return a symbolic elementwise logical-not expression.

Parameters:

x (Any)

Return type:

Any

static clip(x, a_min=None, a_max=None)[source]

Return a symbolic clipped expression.

Parameters:
Return type:

Any

static cumsum(x, axis=None, dtype=None, out=None)[source]

Return a symbolic cumulative sum over axis, matching NumPy’s cumsum contract.

axis behaves exactly as np.cumsum defines it (None flattens first). dtype has no symbolic analogue – an expression node carries no numeric storage to cast into, the same reason sum() rejects it – and out cannot be honored either, because the result is a freshly built object array of expression nodes rather than a buffer that can be written into. Both are therefore rejected rather than silently discarded (MXR-080-1567: the previous implementation was a lambda that swallowed dtype, out, and any other argument into unused *args/**kwargs and passed only axis through, so a caller asking for a float32 accumulation or an output buffer got neither and no signal). An unexpected keyword now raises TypeError from ordinary signature binding, as the other engines’ cumsum does.

Parameters:
Return type:

Any

index_add(out, index, values)[source]

Return a symbolic index-add operation node.

Parameters:
Return type:

Any

static to_sympy(x)[source]

Lower a symbolic expression (or object array) to a sympy expression.

Parameters:

x (Any)

Return type:

Any

static to_sage(x)[source]

Lower a symbolic expression (or object array) to a sage expression.

Parameters:

x (Any)

Return type:

Any

static to_latex(x)[source]

Return a LaTeX string for a symbolic expression via sympy.

Parameters:

x (Any)

Return type:

str

class SymbolicExpression(op, args=())[source]

Bases: object

A compact immutable symbolic expression tree.

Parameters:
static symbol(name)[source]

Create a named symbolic variable node.

Parameters:

name (str)

Return type:

SymbolicExpression

static constant(value)[source]

Create a symbolic constant node.

Parameters:

value (Any)

Return type:

SymbolicExpression

static call(op, *args)[source]

Create an operation node with symbolic arguments.

Parameters:
Return type:

SymbolicExpression

evaluate(values)[source]

Evaluate the expression with a mapping from symbol names to values.

A where node with a SCALAR condition evaluates only the selected branch (MXR-080-1565); see _evaluate_where(). Every other node evaluates all of its arguments.

Parameters:

values (dict[str, Any])

Return type:

Any

symbols()[source]

Return sorted symbolic variable names referenced by this expression.

Return type:

tuple[str, …]

op_counts()[source]

Return operation counts for this expression tree.

Return type:

dict[str, int]

depth()[source]

Return expression-tree depth, counting this node.

Return type:

int

node_count()[source]

Return the total number of expression nodes in this tree.

Return type:

int

class TorchEngine(device=None, dtype=None, compile=False, mesh=None, shard=None)[source]

Bases: ComputeEngine

Torch tensor engine for device placement, autograd, and optional DTensor sharding.

Parameters:

  • device (str | None)

  • dtype (Any)

  • compile (bool)

  • mesh (Any)

  • shard (str | None)

property accumulator_dtype: Any

High-precision dtype for sufficient-statistic reductions (float64, or float32 on MPS).

Reductions that aggregate over observations accumulate in float64 even when scoring runs in reduced precision, so a float32 fit does not drift on large N. MPS has no float64, so there the accumulator falls back to float32 (its max precision) — fits on very large N may drift slightly.

with_precision(precision)[source]

Return a Torch engine with the same placement and a new dtype policy.

Parameters:

precision (Any)

Return type:

TorchEngine

asarray(x, dtype=None)[source]

Convert x to a Torch tensor or DTensor on the configured device.

Contract: float inputs are force-cast to the engine’s float dtype (e.g. a float32 numpy array becomes the engine’s float64) unless dtype is given; this differs from numpy’s asarray, which preserves the input dtype. Complex NumPy input is rejected when dtype is not given explicitly: the dtype inference below has no complex-precision policy of its own, and letting it fall through to the generic non-floating/non-Boolean case would silently discard the imaginary component and collapse the real part to an integer. Pass an explicit dtype=torch.complex64/torch.complex128 for a genuine complex tensor.

Parameters:
Return type:

Any

zeros(shape, dtype=None)[source]

Allocate a zero tensor with this engine’s dtype/device/placement.

Parameters:
Return type:

Any

empty(shape, dtype=None)[source]

Allocate an uninitialized tensor with this engine’s dtype/device/placement.

Parameters:
Return type:

Any

arange(*args, **kwargs)[source]

Return torch.arange on the configured device.

Parameters:
Return type:

Any

to_numpy(x)[source]

Gather/detach a Torch tensor or DTensor to a host NumPy array.

Parameters:

x (Any)

Return type:

ndarray

stack(arrays, axis=0)[source]

Stack tensors with torch.stack under the engine’s dtype and device policy.

Elements route through asarray() first. Leaf scorers legitimately hand back float64 NumPy arrays; stacking them raw produced a float64 tensor that later met float32 engine tensors – invisible on CPU, but a hard “mat1 double, mat2 float” error inside the HMM forward pass on the first real CUDA execution (rented RTX 3060, 2026-08-04). The engine’s dtype policy has to apply at every ingestion point, not only the explicit asarray calls.

Parameters:
Return type:

Any

concatenate(arrays, axis=0)[source]

Join tensors with torch.cat under the same ingestion policy as stack().

Parameters:
Return type:

Any

replicate(x)[source]

Return x replicated across the configured DeviceMesh.

Parameters:

x (Any)

Return type:

Any

place_component_axis(x, axis=0)[source]

Place a stacked component parameter tensor on Shard(axis).

TorchEngine(mesh=..., shard='components') is the model-parallel configuration from the compute-engine design. Encoded data remains replicated, while homogeneous stacked-kernel parameter tensors are sharded along their component dimension. With no mesh, or with no component sharding requested, this is just asarray.

Parameters:
Return type:

Any

requires_grad(x)[source]

Return whether x is a Torch tensor/DTensor requiring gradients.

Parameters:

x (Any)

Return type:

bool

compile(fn)[source]

Compile fn with torch.compile when enabled and available.

Parameters:

fn (Callable)

Return type:

Callable

sum(x, *args, **kwargs)[source]

Return torch.sum accepting either axis or dim.

Promotes a floating input to accumulator_dtype when the caller doesn’t pass an explicit dtype= – mirroring NumpyEngine.sum(). Without this, torch.sum accumulates in the input tensor’s own dtype by default, so a float32-precision fit on this engine would silently drift on large N (the exact catastrophic-cancellation risk accumulator_dtype exists to guard against) while the numpy engine, which already promotes, stayed accurate.

static max(x, *args, **kwargs)[source]

Return torch.max accepting either axis or dim (incl. a tuple of axes).

matmul(x, y)[source]

Matrix-multiply under the engine’s dtype policy.

Both operands route through asarray(). Mixed float32/float64 operands reached this point through paths that never crossed an explicit ingestion call, and torch.matmul promotes on CPU but hard-errors on CUDA (“mat1 double, mat2 float”) – found and probed on the first real CUDA execution of the HMM forward pass (rented RTX 3060, 2026-08-04: mat1 float64 [100, 2] against mat2 float32 [2, 2]). The engine’s declared dtype IS the policy for floating operands, so applying it here is a correctness boundary, not a cast of convenience.

Parameters:
Return type:

Any

static cumsum(x, *args, **kwargs)[source]

Return torch.cumsum accepting axis/dim and defaulting to a flattened scan.

static logsumexp(x, *args, **kwargs)[source]

Return torch.logsumexp accepting either axis or dim.

torch < 2.5 registers no DTensor sharding strategy for logsumexp, so reducing over a sharded axis (the mixture E-step’s log-partition over component-sharded scores) raises NotImplementedError. Fall back to redistributing the DTensor to replicated first – the reduction then runs locally and is correct on any torch version. torch >= 2.5 has the strategy and takes the fast native path unchanged; non-DTensor inputs re-raise the original error.

hmm_forward_ll(log_emit, log_w, log_a, mask)[source]

Fused log-space forward: per-sequence emission log-likelihood (freeze-alpha padding).

Parameters:
Return type:

Any

hmm_forward_backward(log_emit, log_w, log_a, mask, weights=None)[source]

Fused log-space Baum-Welch recurrences, exactly mirroring the generic engine-op version: alpha freezes across padded steps, beta carries, gamma is mask-SELECTED (NaN-safe for empty sequences), xi is computed for all transitions at once, and per-sequence weights scale gamma / xi / pi. Returns (ll, gamma, xi_sum, pi).

Parameters:
Return type:

tuple[Any, Any, Any, Any]

index_add(out, index, values)[source]

Add values into out along axis 0 using Torch index_add.

Contract: return-value-only – torch.Tensor.index_add (no trailing underscore) does not mutate out in place, unlike numpy’s add.at; callers must use the returned tensor.

Parameters:
Return type:

Any

class JaxEngine(device=None, dtype=None, compile=False)[source]

Bases: ComputeEngine

JAX array engine: XLA-compiled ops, float64, optional jax.jit compilation, GPU/TPU via JAX.

Parameters:

  • device (str | None)

  • dtype (Any)

  • compile (bool)

property accumulator_dtype: Any

High-precision dtype for sufficient-statistic reductions (float64, or float32 when the ambient jax_enable_x64 is off – see the module docstring; JAX would silently truncate a float64 accumulator request to float32 there anyway, so this avoids claiming a precision the runtime cannot actually give).

with_precision(precision)[source]

Return a JAX engine with the same placement and a new dtype policy.

Parameters:

precision (Any)

Return type:

JaxEngine

asarray(x, dtype=None)[source]

Convert x to a JAX array on this engine’s device. Float inputs are force-cast to the engine dtype (float64 by default) unless dtype is given – matching the Torch engine’s contract, not NumPy’s.

Complex input is rejected when dtype is not given explicitly (MXR-080-1560), exactly as mixle.engines.torch_engine.TorchEngine.asarray() does: this engine’s dtype policy has no complex-precision concept, so the generic non-floating/non-Boolean fallback (jnp.int64) would silently discard the imaginary component and collapse the real part to an integer. Pass dtype=jnp.complex64/jnp.complex128 for a genuine complex array.

Parameters:
Return type:

Any

zeros(shape, dtype=None)[source]

Allocate a zero array with this engine’s dtype, placed on this engine’s device.

Parameters:
Return type:

Any

empty(shape, dtype=None)[source]

Allocate an array on this engine’s device (JAX has no uninitialized empty; zeros is the safe equivalent).

Parameters:
Return type:

Any

arange(*args, **kwargs)[source]

Return jnp.arange placed on this engine’s device; float arguments select the engine float dtype.

Parameters:
Return type:

Any

is_concrete(x)[source]

False while x is a jit tracer: its values do not exist yet, so to_numpy would raise.

Parameters:

x (Any)

Return type:

bool

to_numpy(x)[source]

Move a JAX array back to a host NumPy array.

Parameters:

x (Any)

Return type:

ndarray

stack(arrays, axis=0)[source]

Stack arrays with jnp.stack, placed on this engine’s device.

Parameters:
Return type:

Any

concatenate(arrays, axis=0)[source]

Join arrays with jnp.concatenate on this engine’s device.

Parameters:
Return type:

Any

requires_grad(x)[source]

Always False: JAX autograd is functional (jax.grad), not tensor-tagged.

Parameters:

x (Any)

Return type:

bool

compile(fn)[source]

Compile fn with jax.jit when enabled.

Parameters:

fn (Callable)

Return type:

Callable

index_add(out, index, values)[source]

Add values into out along axis 0 via the functional .at[idx].add update, placed on this engine’s device.

Contract: return-value-only – JAX arrays are immutable, so this returns a new array; callers must use the return value (the same contract the Torch engine documents).

Parameters:
Return type:

Any

auto_precision(data=None, *, engine=None, sample_size=512)[source]

Recommend 'float32' or 'float64' from the data and the target hardware.

float32 only helps on a GPU Torch engine (on CPU/NumPy it is a no-op or slower), and even there only when the data is well conditioned for single precision. Sufficient-statistic accumulation is already float64-safe (see ComputeEngine.accumulator_dtype), so this guards the remaining risk – the ~7 significant digits of float32 scoring – by inspecting the data’s magnitude and dynamic range. Returns 'float64' whenever a numeric sample is unavailable, non-finite, or looks risky.

Parameters:

  • data (Any) – A representative sample of the raw observations (or an iterable of them).

  • engine (Any) – The target compute engine; float32 is only recommended for a GPU Torch engine.

  • sample_size (int) – How many observations to inspect, at most. Must be a positive integer.

Returns:

'float32' or 'float64'.

Return type:

str

engine_of(x, default=NUMPY_ENGINE)[source]

Return the ComputeEngine associated with an array or encoded payload.

Nested encodings are scanned recursively. Mixing arrays whose engines disagree on backend, device, mesh/placement, or an explicitly-chosen dtype policy is an error, because silent host/device/precision mixing is almost always a performance or correctness bug (see _engines_compatible()).

Parameters:

  • x (Any)

  • default (ComputeEngine)

Return type:

ComputeEngine

engine_with_precision(engine, precision)[source]

Return engine adjusted to the requested floating precision.

Parameters:

  • engine (Any)

  • precision (Any)

Return type:

Any

is_symbolic_payload(x)[source]

Return True for a symbolic node or a NumPy object array of ALL-symbolic nodes.

Every element of an object array is checked, not just the first: classifying the whole array from a single element lets a mixed-ownership payload – some elements genuine SymbolicExpression nodes, others unrelated plain objects – be misrouted based on whichever kind happens to occupy index 0 (numeric-first routes the symbolic elements through NumPy, which was never built to handle them; symbolic-first routes the unrelated elements through the symbolic engine instead). Mixed ownership is rejected outright rather than silently guessed at from one element.

Scans left to right and stops at the first element whose kind disagrees with the first element’s, so a genuinely uniform array (the overwhelmingly common case) still only pays for a single early-exit scan, while a mixed one is still caught regardless of which kind happens to sit at index 0.

Parameters:

x (Any)

Return type:

bool

normalize_numpy_dtype(precision)[source]

Normalize a precision specifier to a NumPy floating dtype.

Parameters:

precision (Any)

Return type:

dtype | None

normalize_torch_dtype(precision, torch_module)[source]

Normalize a precision specifier to a Torch floating dtype.

Parameters:

  • precision (Any)

  • torch_module (Any)

Return type:

Any

precision_name(precision)[source]

Return a readable canonical precision name.

Parameters:

precision (Any)

Return type:

str

register_array_type(array_type, engine, *, override=False)[source]

Register an array/tensor type with its owning engine.

Existing ownership cannot be replaced unless override=True is explicit. Registering object is forbidden because it would capture every otherwise-unregistered Python value process-wide.

Parameters:

  • array_type (type[Any])

  • engine (ComputeEngine)

  • override (bool)

Return type:

None

to_latex(expr)[source]

Return a LaTeX string for expr via sympy.

Parameters:

expr (Any)

Return type:

str

to_numpy(x)[source]

Convert an engine array/tensor payload to NumPy at an explicit boundary.

A dict, list, or tuple that contains an engine-owned value is converted leaf by leaf, preserving container structure and resolving each leaf’s own owning engine independently. Resolving a single engine for the whole container up front and handing it the container as-is – the previous behavior – fails to stack a ragged list of Torch tensors, leaves dictionary values unconverted, and can hand a device tensor to np.asarray without transferring it to host memory first. Plain nested data with no engine-owned value is still handed to the resolved engine as one unit, unchanged.

Parameters:

x (Any)

Return type:

Any

to_sage(expr)[source]

Convert a SymbolicExpression (or object array of them) to sage.

Mirrors to_sympy() against sage.all. Raises ImportError if sage is unavailable, and NotImplementedError for non-symbolic kernels.

Parameters:

expr (Any)

Return type:

Any

to_sympy(expr)[source]

Convert a SymbolicExpression (or object array of them) to sympy.

Scalar nodes become a sympy expression; object arrays map elementwise into a NumPy object array of sympy expressions (shape preserved). Genuinely non-symbolic kernels (bincount/unique/searchsorted/ index_add) raise NotImplementedError. Raises ImportError if sympy is unavailable.

Parameters:

expr (Any)

Return type:

Any

class DoubleDouble(hi, lo=0.0)[source]

Bases: object

An (almost) ~106-bit float as two non-overlapping float64 arrays hi + lo.

Scalars or numpy arrays; operations broadcast. The invariant is |lo| <= 0.5 * ulp(hi). A component may be +inf or -inf – exactly representable in float64, carried with lo == 0 – but never NaN, and never an overflow of finite operands.

Parameters:

  • hi (Any)

  • lo (Any)

classmethod from_float(x)[source]

Create a double-double value from float64 data.

Parameters:

x (Any)

Return type:

DoubleDouble

to_float()[source]

Collapse back to the nearest float64 (hi rounded with lo).

Return type:

ndarray

dd_sum(x)[source]

Accurate sum of a float64 array in double-double precision – vectorized, no Python loop over elements.

Pairwise tree reduction with an error-free two_sum() combine at every node: O(n) work in O(log n) vectorized passes, accumulating the rounding error into the lo component. The result is correct to ~106 bits even for catastrophically cancelling inputs that float64 sums get wrong.

An infinite term short-circuits to the IEEE answer (-inf for a zero-probability log-density) instead of raising; NaN, and +inf mixed with -inf, still fail closed – see _infinite_sum().

Parameters:

x (Any)

Return type:

DoubleDouble

dd_dot(a, b)[source]

Accurate dot product sum(a_i * b_i) in double-double precision.

Uses the compiled hardware-FMA kernel when available (one fma per element, ~3x faster than the pure-numpy Veltkamp-split path, and exactly as accurate – each product’s error is exact either way, though the two paths accumulate in different orders, sequential versus pairwise tree, so their results agree to double-double precision rather than bit for bit). Otherwise each product is split error-free by _product_and_residual() and the products + errors are summed by dd_sum(). Defeats the cancellation that wrecks a naive float64 dot.

Every check that decides whether an input is acceptable is made once, before dispatch, so the two implementations cannot drift apart on what they accept: equal length (see _require_equal_length(), ValueError), finite operands (ValueError), and a result that is representable at all (OverflowError). A single product that underflows to zero is not one of them: it contributes less than one smallest-subnormal to the sum, so it is summed as the zero it rounds to rather than failing the whole reduction – dd_dot([1e-200, 1.0], [1e-200, 1.0]) is 1.0, which is the right answer, not an error. (two_prod() on its own still fails closed there, because a lone product has nothing to be negligible against.)

Parameters:
Return type:

DoubleDouble

class FloatFormat(mantissa_bits, exp_bits=11)[source]

Bases: NumericFormat

Round-to-nearest MANTISSA rounding at mantissa_bits, nominally styled after an IEEE-ish fpN split (see fp()).

quantize rounds each value’s significand to mantissa_bits explicit bits (round-to-nearest); the exponent passes through UNCHANGED, with no range limit. exp_bits therefore only feeds name / bits_per_value (the advertised bit budget an allocation pass compares against) – it is NOT enforced: there is no overflow/underflow, no saturation or flush-to-zero, no subnormals, and no NaN/Inf bit-pattern encoding. This is a deliberate scope choice, not an oversight: an unbounded exponent is what makes max_rel_error a genuine bound for any finite magnitude (see that property’s docstring), so a caller can read the round trip as a pure mantissa-rounding measurement independent of range.

A consequence: this codec does NOT faithfully emulate a real bounded fpN storage format. A real 8-bit float cannot represent 1e100 at all, but FloatFormat.fp(8).round_trip([1e100]) rounds 1e100’s mantissa to 3 bits and returns a value that is still ~1e100 in magnitude – by design, per the above, not a bug. Do not use fp(8)/fp(16)/etc. where actual overflow/underflow/saturation behavior matters; they measure mantissa-rounding error only.

Parameters:

  • mantissa_bits (int)

  • exp_bits (int)

classmethod fp(total_bits)[source]

Build a format whose mantissa_bits/exp_bits follow an IEEE-like split of total_bits (n = 1..1024+). Only the resulting mantissa_bits affects quantize’s numeric behavior – exp_bits is bit-budget bookkeeping only (see the class docstring); this does NOT build a range-limited fpN codec.

Parameters:

total_bits (int)

Return type:

FloatFormat

property max_rel_error: float

Return the worst-case relative rounding error – a universal bound, true for every finite input regardless of magnitude (the exponent is never range-limited; see the class docstring).

quantize(x)[source]

Round values to this mantissa precision; the exponent passes through exactly, unclamped.

Parameters:

x (Any)

Return type:

ndarray

dequantize(q)[source]

Decode quantized floating values to float64.

Parameters:

q (Any)

Return type:

ndarray

class FixedPointFormat(frac_bits, int_bits=31)[source]

Bases: NumericFormat

Fixed-point: store round(x * 2**frac_bits) as an integer; real compression + a hard error bound.

int_bits sets the representable magnitude [-2**int_bits, 2**int_bits); out-of-range clamps. in_range_max_abs_error == 2**-(frac_bits+1) bounds round-to-nearest only for values inside that range. The universal max_abs_error is infinite because saturation error is unbounded. Storage is a plain int32/int64 – there is no arbitrary-width packed storage – so the declared total width 1 + int_bits + frac_bits cannot exceed 64 bits; wider requests raise at construction instead of silently overflowing the int64 cast (and falsifying max_abs_error).

Parameters:

  • frac_bits (int)

  • int_bits (int)

property max_abs_error: float

Return the universal error bound; saturation makes it unbounded.

property in_range_max_abs_error: float

Return the half-step error bound for inputs inside the representable range.

quantize(x)[source]

Encode values as clipped scaled integers.

Parameters:

x (Any)

Return type:

ndarray

dequantize(q)[source]

Decode scaled integers back to float64 values.

Parameters:

q (Any)

Return type:

ndarray

class CodebookFormat(codebook)[source]

Bases: NumericFormat

Scalar vector-quantization: store an index into a learned codebook; log2(K) bits per value.

The genuine pure-numpy compression codec – quantize gathers the nearest code (an unsigned index array), dequantize gathers the code values back. Fit the codebook to data with fit().

The codebook must be finite and nonempty (validated at construction, ValueError otherwise). Duplicate entries are allowed: a repeated code costs a wasted, unreachable index but does not break nearest-code assignment or decoding (both entries decode to the same, correct value), and fit()’s Lloyd iteration can legitimately converge two cluster means onto the same value on degenerate/skewed data – rejecting that outcome would be pure loss, not a correctness fix.

Parameters:

codebook (Any)

classmethod fit(data, n_codes, iters=25, seed=0)[source]

Learn n_codes codes by 1-D k-means (Lloyd) on data; codes are the cluster means.

Raises ValueError if n_codes or iters is not a positive integer.

Parameters:
Return type:

CodebookFormat

quantize(x)[source]

Map values to nearest codebook indices.

Parameters:

x (Any)

Return type:

ndarray

dequantize(q)[source]

Map codebook indices back to representative values.

Parameters:

q (Any)

Return type:

ndarray

compress(x)[source]

Quantize x and pack the indices to bytes: returns (packed_uint8, count).

For K <= 256 codes the indices are sub-byte or byte-exact and pack_bits() bit-packs them (e.g. 16 codes -> 4-bit indices -> 2 values/byte -> 16x vs float64). Larger codebooks need 16- or 32-bit indices, which are not sub-byte-packable – those are stored byte-aligned (little-endian), still returned as a flat uint8 array.

Parameters:

x (Any)

Return type:

tuple[ndarray, int]

decompress(packed, count)[source]

Inverse of compress(): unpack indices and gather the codebook back to float64.

Parameters:
Return type:

ndarray

class Interval(lo, hi)[source]

Bases: object

A guaranteed enclosure [lo, hi] of a value (scalar or numpy array), outward-rounded.

Parameters:

  • lo (Any)

  • hi (Any)

classmethod exact(x)[source]

A degenerate interval [x, x] for an exactly-represented value.

Parameters:

x (Any)

Return type:

Interval

classmethod from_quantized(original, fmt)[source]

Enclose original using the quantized format’s error bound.

max_rel_error is a duck-typed, self-reported claim: nothing here enforces that an implementer’s static bound actually holds for every value fmt might quantize, only that it is offered as one. A format that flushes underflowing values to exactly zero, or saturates on overflow, would still report a fixed max_rel_error while its true error at those magnitudes no longer scales with it – the analytic pad below would then be too narrow, silently returning an interval that does not actually enclose original ( FloatFormat avoids this today by staying mantissa-only with an unbounded exponent, so its max_rel_error genuinely is universal – see that class’s docstring – but nothing here depends on every format that will ever be passed in honoring that same discipline). Guard against a format that breaks the assumption by cross-checking the analytic pad against the error actually measured on original – already in hand, it is quantized right here – and padding by whichever is WIDER. This needs nothing new from fmt beyond round_trip, which every caller already requires, and is a no-op whenever max_rel_error genuinely does bound this data (the analytic pad is then always >= the measured one, by construction).

Parameters:
Return type:

Interval

width()[source]

The guaranteed error bound: hi - lo (outward-rounded).

A degenerate point exactly at +/-inf (lo == hi == +/-inf, legal per the constructor) has zero width by definition – but IEEE-754 computes inf - inf == NaN for it, which would otherwise leak a NaN out of this function with no constructor downstream to catch it. Special-case that one pattern directly; every other case (including an ordinary, non- degenerate unbounded interval like [0, inf], whose width is correctly inf) is unaffected.

Return type:

ndarray

max_width()[source]

Return the largest interval width.

Return type:

float

midpoint()[source]

Return interval midpoints.

Same inf - inf gap as width(): a degenerate point at +/-inf is its own midpoint by definition, computed directly rather than via the subtraction that would otherwise produce NaN.

Return type:

ndarray

contains(value)[source]

Return a boolean mask for values inside the interval.

Parameters:

value (Any)

Return type:

ndarray

sum_error_bound(x)[source]

Return a certified bound on the float64 error of sum(x).

The standard a-priori bound |fl(sum) - sum| <= gamma_{n-1} * sum|x_i| with gamma_k = k*u / (1 - k*u) and u = 2**-53. It is sound for any summation order. A bound that is large relative to |sum x| means the sum is ill-conditioned (cancellation) and warrants the double-double mixle.engines.extended.dd_sum(); a tight one means float64 already suffices and no extra compute is justified.

Parameters:

x (Any)

Return type:

float

float64_sum_is_accurate(x, target_rel_error=1e-12)[source]

Return whether float64 summation is accurate to target_rel_error.

Reads the certified bound relative to the magnitude of the result – precision allocation in one call.

Parameters:
Return type:

bool

class AffineForm(center, terms=None)[source]

Bases: object

center + sum_i coeff_i * eps_i over shared noise symbols eps_i in [-1, 1].

Parameters:

  • center (Any)

  • terms (dict[int, np.ndarray] | None)

classmethod constant(x)[source]

Create an affine form with no uncertainty terms.

Parameters:

x (Any)

Return type:

AffineForm

classmethod uncertain(x, radius=0.0)[source]

An input known only to within +/- radius – one fresh noise symbol of that radius.

Parameters:
Return type:

AffineForm

radius()[source]

Half-width = sum_i |coeff_i| (one outward ULP of slop for the f64 summation).

Return type:

ndarray

max_radius()[source]

Return the largest interval half-width across entries.

Return type:

float

to_interval()[source]

Convert the affine form to an interval enclosure.

Return type:

Any

contains(value)[source]

Return a boolean mask for values contained in the affine enclosure.

Parameters:

value (Any)

Return type:

ndarray

inject_roundoff(dtype)[source]

Add the roundoff a dtype-dtype evaluation introduces: a fresh symbol of u*|center|.

Parameters:

dtype (Any)

Return type:

AffineForm

allocate_precision(center_magnitude, op_count, target_abs_error)[source]

Suggest the lowest-cost dtype whose endpoint-magnitude error heuristic is under target.

The estimate op_count*u*center_magnitude does not observe intermediate values or an operation graph and is therefore only a planning heuristic. It must never be used as an accuracy certificate.

Raises ValueError unless center_magnitude and op_count are both finite and nonnegative and target_abs_error is finite and strictly positive. Before this validation existed, a negative op_count flipped the estimated error negative, which spuriously satisfies every dtype’s budget comparison and misselects float16 – the least precise dtype – regardless of how large the magnitude or how tight the target; a negative target_abs_error made every dtype’s comparison fail and fell through to silently returning qd even though no dtype can certify a negative bound (impossible by definition); and a NaN/Inf value anywhere in the formula rode comparison-with-NaN-is-always-False (or Inf’s unpredictable propagation) into whatever the fallthrough dtype happened to be, not a principled decision.

Parameters:
Return type:

PrecisionResult

accurate_sum(x, target_rel_error=1e-12)[source]

Sum x to target_rel_error relative accuracy using the lowest-cost sufficient backend.

Returns a SumResult certifying the accuracy actually achieved. Escalates from float64 to double-double to arbitrary precision only when the certified error bound says the lower-cost backend cannot meet the budget – so the common well-conditioned case never leaves vectorized float64 – and fails closed (status="insufficient") rather than silently returning a value that does not meet target_rel_error when no backend available in this process can certify it.

Raises ValueError if target_rel_error is not a positive, finite number.

Parameters:
Return type:

SumResult

cast(x, precision)[source]

Cast x onto the spectrum: a native dtype name, "dd"/"fp<bits>", or an integer bit width.

"dd", "fp128", and the integer 128 are the same request (DoubleDouble); "fp96" and the integer 96 are likewise the same request (both resolve to 96 total bits via _precision_bits()) and route to the identical representation via _cast_by_bits(), regardless of which spelling the caller used. A nonpositive bit width (e.g. "fp0", -5) raises ValueError rather than reaching format construction. Anything else (e.g. "float64", np.float32) is handed straight to numpy.dtype().

Returns a numpy array (native / simulated low-bit), a DoubleDouble, or an MPFR object array – see _cast_by_bits() for the exact tiering.

Parameters:
Return type:

Any

sum_certificate(x)[source]

Report the certified float64 summation error and the condition number, without choosing a backend.

Parameters:

x (Any)

Return type:

dict[str, float | str | bool]

Submodules