mixle.engines package

Compute engines for backend-neutral mixle kernels.

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.

name = 'base'
supports_autograd = False
dtype = None
device = 'cpu'
REQUIRED_OPS: tuple[str, ...] = ('asarray', 'zeros', 'empty', 'arange', 'to_numpy', 'stack', 'log', 'exp', 'sqrt', 'abs', 'where', 'maximum', 'clip', 'floor', 'isnan', 'isinf', 'sum', 'max', 'dot', 'matmul', 'cumsum', 'logsumexp', 'bincount', 'unique', 'searchsorted', 'index_add', 'gammaln', 'digamma', 'betaln', 'erf')
pi = 3.141592653589793
e = 2.718281828459045
euler_gamma = 0.5772156649015329
inf = inf
zero = 0.0
one = 1.0
two = 2.0
half = 0.5
constant(value)[source]

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

Parameters:

value (Any)

Return type:

Any

supports_numba = False
resident_estep = True
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

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)

name = 'numpy'
supports_autograd = False
supports_numba = True
resident_estep = False
device = 'cpu'
prefer_fused = False
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

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

min

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=np.float64)
>>> 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=np.float64)  # 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=np.float64))
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.

name = 'symbolic'
supports_autograd = False
pi = pi
e = e
euler_gamma = euler_gamma
inf = inf
zero = 0
one = 1
two = 2
half = (1 / 2)
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

static log(x)
static exp(x)
static sqrt(x)
static abs(x)
static floor(x)
static gammaln(x)
static digamma(x)
static erf(x)
static sum(x, axis=None, *args, **kwargs)[source]

Return a symbolic sum reduction over axis.

Parameters:
Return type:

Any

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

Return a symbolic log-sum-exp reduction over axis.

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 isnan(x)
static isinf(x)
static max(x, axis=None, *args, **kwargs)
static dot(x, y)
static matmul(x, y)
static cumsum(x, axis=None, *args, **kwargs)
static bincount(x, *args, **kwargs)
static unique(x, *args, **kwargs)
static searchsorted(x, y, *args, **kwargs)
static betaln(x, y)
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 tiny immutable symbolic expression tree.

Parameters:
op: str
args: tuple[Any, ...] = ()
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.

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)

name = 'torch'
supports_autograd = True
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.

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 and apply default placement.

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

static log(x)
static exp(x)
static sqrt(x)
static abs(x)
static where(*args)
static maximum(x, y)
static clip(x, a_min=None, a_max=None)
static floor(x)
static isnan(x)
static isinf(x)
static sum(x, *args, **kwargs)[source]

Return torch.sum accepting either axis or dim.

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

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

static dot(x, y)
static matmul(x, y)
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.

static bincount(x, *args, **kwargs)
static unique(x, *args, **kwargs)
static searchsorted(x, y, *args, **kwargs)
static gammaln(x)
static digamma(x)
static betaln(x, y)
static erf(x)
static cos(x)
static sin(x)
static arctan2(x, y)
static i0e(x)
static erfcx(x)
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)

name = 'jax'
supports_autograd = True
property accumulator_dtype: Any

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

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. 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.

Parameters:
Return type:

Any

zeros(shape, dtype=None)[source]

Allocate a zero array with this engine’s dtype.

Parameters:
Return type:

Any

empty(shape, dtype=None)[source]

Allocate an array (JAX has no uninitialized empty; zeros is the safe equivalent).

Parameters:
Return type:

Any

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

Return jnp.arange; float arguments select the engine float dtype.

Parameters:
Return type:

Any

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.

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

static log(x)
static exp(x)
static sqrt(x)
static abs(x)
static where(*args)
static maximum(x, y)
static clip(x, a_min=None, a_max=None)
static floor(x)
static isnan(x)
static isinf(x)
static sum(x, *args, **kwargs)
static max(x, *args, **kwargs)
static dot(x, y)
static matmul(x, y)
static cumsum(x, *args, **kwargs)
static logsumexp(x, *args, **kwargs)
static bincount(x, *args, **kwargs)
static unique(x, *args, **kwargs)
static searchsorted(x, y, *args, **kwargs)
static gammaln(x)
static digamma(x)
static betaln(x, y)
static erf(x)
static cos(x)
static sin(x)
static arctan2(x, y)
static i0e(x)
index_add(out, index, values)[source]

Add values into out along axis 0 via the functional .at[idx].add update.

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 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 leading observations to inspect.

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 owned by different engine classes is an error because silent host/device mixing is almost always a performance or correctness bug.

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

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)[source]

Register an array/tensor type with its owning engine.

Parameters:

  • array_type (type[Any])

  • engine (ComputeEngine)

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.

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).

Parameters:

  • hi (Any)

  • lo (Any)

hi
lo
classmethod from_float(x)[source]
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.

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 bit-for-bit identical); otherwise each product is split error-free by two_prod() and the products + errors are summed by dd_sum(). Defeats the cancellation that wrecks a naive float64 dot.

Parameters:
Return type:

DoubleDouble

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

Bases: NumericFormat

A float with mantissa_bits of mantissa – rounds x to that band’s representable set.

Below ~52 mantissa bits this is lossy (fp4/fp8/fp16 storage); the round trip measures the band’s accuracy. max_rel_error == 2**-(mantissa_bits+1) from round-to-nearest.

Parameters:

  • mantissa_bits (int)

  • exp_bits (int)

classmethod fp(total_bits)[source]

Build an n-bit float format with an IEEE-like exponent/mantissa split (n = 1..1024+).

Parameters:

total_bits (int)

Return type:

FloatFormat

property max_rel_error: float
quantize(x)[source]
Parameters:

x (Any)

Return type:

ndarray

dequantize(q)[source]
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. max_abs_error == 2**-(frac_bits+1) (round-to-nearest), independent of x.

Parameters:

  • frac_bits (int)

  • int_bits (int)

property max_abs_error: float
quantize(x)[source]
Parameters:

x (Any)

Return type:

ndarray

dequantize(q)[source]
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().

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.

Parameters:
Return type:

CodebookFormat

quantize(x)[source]
Parameters:

x (Any)

Return type:

ndarray

dequantize(q)[source]
Parameters:

q (Any)

Return type:

ndarray

compress(x)[source]

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

For K <= 16 codes the indices are sub-byte and pack 8//bits per byte, realizing the advertised compression (e.g. 16 codes -> 4-bit indices -> 2 values/byte -> 16x vs float64).

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)

hi
lo
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 given only its round-trip through fmt – i.e. bound the value a consumer of the quantized data is uncertain about, using the format’s relative-error bound.

Parameters:
Return type:

Interval

width()[source]

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

Return type:

ndarray

max_width()[source]
Return type:

float

midpoint()[source]
Return type:

ndarray

contains(value)[source]
Parameters:

value (Any)

Return type:

ndarray

sum_error_bound(x)[source]

Certified bound on the float64 error of sum(x) – vectorized, no per-element loop.

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 – the logic that spends precision only where it is needed.

Parameters:

x (Any)

Return type:

float

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

Whether a float64 sum of x is already accurate to target_rel_error (else use dd_sum).

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)

center
terms
classmethod constant(x)[source]
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 type:

float

to_interval()[source]
Return type:

Any

contains(value)[source]
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]

Cheapest dtype whose accumulated roundoff over op_count ops keeps error under target.

Each op injects ~``u(d) * |magnitude|``; op_count of them accumulate to op_count*u*|mag|. Walk cheapest-first and return the first dtype that fits the budget – spend minimal precision.

Parameters:

  • center_magnitude (float)

  • op_count (int)

  • target_abs_error (float)

Return type:

str

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

Sum x to target_rel_error relative accuracy using the cheapest sufficient backend.

Returns (value, backend) where backend is "float64", "dd" (double-double), or "mpfr<bits>". Escalates only when the certified error bound says the cheaper backend cannot meet the budget – so the common (well-conditioned) case never leaves vectorized float64.

Parameters:
Return type:

tuple[float, str]

cast(x, precision)[source]

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

Returns a numpy array (native), a DoubleDouble (dd/fp128), or an MPFR object array (>= ~fp256 / explicit bit width).

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]

Submodules