solprop.backends.mlx_gpu.arithmetic#

The protected operations of the mlx backend, and the probes that establish them.

SEP 33 section 6.3’s arithmetic policy. Each protected operation has a mechanism and its evidence:

Matrix products are float32 only with MLX_ENABLE_TF32=0, which the package sets when it is imported and MLX has not been initialised; gemm_probe() measures every dispatch family a run uses against a discriminating bound (2**-17) and refuses a device that computes them in TF32.
The stage propagator and every combination of stage terms are formed in double-float arithmetic from binary64 terms held as pairs of binary32 numbers (DoubleFloatTerms), then rounded once: on the Metal device by one custom kernel compiled in "safe" math mode, whose two-sum and two-product cannot be reassociated; on MLX’s CPU device, available for local tests, in binary64 on the host with NumPy (the memory is shared). The phase is range-reduced in double-float before its binary32 cosine and sine, and the exponential of the real part is exp(hi) (1 + lo).

Nothing outside solprop.backends.mlx_gpu imports MLX.

Module Attributes

GEMM_PROBE_BOUND

The normwise relative error a float32 product must stay below (SEP 33 7.1).

Functions

`fft_bound_32`(samples)	Return the established transform bound at binary32, `512 eps32 ceil(log2 n)`.
`fft_probe`(lengths)	Measure the device's transforms at a run's own kinds and lengths (SEP 33 6.3).
`gemm_probe`(shapes)	Measure float32 matrix products of each dispatch family a run uses.
`on_gpu`()	Return whether MLX's default device is the Metal GPU.
`population`(exposure, preionisation)	Return the ionised fraction and the neutral survival in their stable forms.
`population_parameters`(preionisation)	Return the population kernel's parameters, `(1 - p, p)` rounded to binary32.
`prefix_scan`(values)	Return the oracle's doubling scan along the last axis, in one launch.
`rate_kernel`(magnitude, coefficients, ...)	Evaluate a rate table's spline at field magnitudes in one launch.

Classes

DoubleFloatTerms(shape, n_terms, parts[, ...])

Complex binary64 terms held on the device for double-float combination.

class solprop.backends.mlx_gpu.arithmetic.DoubleFloatTerms(shape, n_terms, parts, host=None, real_bound=())[source]#

Bases: object

Complex binary64 terms held on the device for double-float combination.

On the Metal device each term is four float32 arrays (the real and imaginary heads and tails); on MLX’s CPU device it is one complex128-equivalent pair of float64 arrays. combine() returns sum_t w_t T_t rounded once to complex64, and exponential() exp(sign sum_t w_t T_t) with the phase range-reduced in double-float before the binary32 exponential, cosine and sine (SEP 33 4.2).

Attributes:

shapetuple of int: The shape of one term.
n_termsint: The number of terms.
partstuple of mlx.core.array: The device arrays, each of shape (n_terms * prod(shape),); empty on MLX’s CPU device.
hostnumpy.ndarray or None: The binary64 terms, kept on MLX’s CPU device only.
real_boundtuple of float: Each term’s largest |Re T_t|, so that sum_t |w_t| real_bound_t bounds a combination’s real parts without forming it.

combination(w_hi, w_lo)[source]#

Return sum_t w_t T_t rounded once to complex64, from device weights.

combine() with weights already on the device (device_weights()).

Parameters:

w_hi, w_lomlx.core.array: The weights’ heads and tails.

Returns:

mlx.core.array: The combination.

combine(weights)[source]#

Return sum_t w_t T_t rounded once to complex64.

Parameters:

weightsnumpy.ndarray: One binary64 weight per term.

Returns:

mlx.core.array: The combination.

device_weights(weights)[source]#

Split binary64 weights into the kernel’s float32 heads and tails.

Parameters:

weightsnumpy.ndarray: One binary64 weight per term.

Returns:

tuple of mlx.core.array: The heads and the tails, (n_terms,) each.

Raises:

BackendError: If the weights do not match the terms.

exponential(weights, sign=1)[source]#

Return exp(sign sum_t w_t T_t) to binary64 accuracy, rounded once.

Parameters:

weightsnumpy.ndarray: One binary64 weight per term.
signint, optional: +1 for the propagator, -1 for its inverse.

Returns:

mlx.core.array: The complex64 propagator.

host: ndarray | None = None#

moments(w_hi, w_lo, count)[source]#

Return exp(x) and k! phi_{k+1}(x) for k < count, x = sum w T.

An exponential Adams trial’s flow and moments (SEP 16) in double-float, each rounded once to complex64 (trial() without a compensation). On the Metal device only.

Parameters:

w_hi, w_lomlx.core.array: The weights’ heads and tails (device_weights()).
countint: The moments, at most MAX_PHI_ORDER.

Returns:

tuple of mlx.core.array: The flow, of shape, and the moments, (count, *shape).

n_terms: int#

property nbytes: int[source]#

Return the device bytes the terms hold.

Returns:

int: Their size.

pair(w_hi, w_lo)[source]#

Return exp(sum_t w_t T_t) and its inverse from one combination.

On the Metal device only; each equals exponential() at its sign. Its inputs are device arrays, so a compiled graph may call it.

Parameters:

w_hi, w_lomlx.core.array: The weights’ heads and tails (device_weights()).

Returns:

tuple of mlx.core.array: The propagator and its inverse, complex64.

parts: tuple[mlx.core.array, ...]#

real_bound: tuple[float, ...] = ()#

real_reach(weights)[source]#

Return a bound on max |Re sum_t w_t T_t| from the terms’ own bounds.

Parameters:

weightsnumpy.ndarray: One binary64 weight per term.

Returns:

float: sum_t |w_t| max |Re T_t|.

shape: tuple[int, ...]#

trial(weights, count, spans, *, endpoint)[source]#

Return an exponential Adams trial’s flow, moments and compensation factors.

One pass over the terms, in double-float, every output rounded once: exp(x) and k! phi_{k+1}(x) – the reference’s far recurrence beyond its join and its scaled series and doubling within it (solprop.backends.exponential.phi_functions()), exp reduced by whole turns of a double-float 2 pi – with x the first row’s combination; then for each span s_j the factor exp(s_j x) R_j, R_j the combination of row 1 + j; then, with endpoint, the last row’s combination alone. On the Metal device only.

Parameters:

weightsnumpy.ndarray: Binary64 weights over the terms, (n_rows, n_terms): x’s, each history sample’s difference of rates, and the endpoint’s.
countint: The moments, at most MAX_PHI_ORDER.
spansnumpy.ndarray: s_j, binary64, one per history sample.
endpointbool: Whether the last row is the endpoint’s.

Returns:

tuple: The flow, of shape; the moments, (count, *shape); and the samples’ factors then the endpoint’s, each of shape.

Raises:

BackendError: If the weights do not match the terms or the rows.

classmethod upload(terms)[source]#

Upload binary64 terms of shape (n_terms, *shape).

Parameters:

termsnumpy.ndarray: Complex binary64 terms, finite.

Returns:

DoubleFloatTerms: The terms on the device.

Raises:

BackendError: If a term is not finite.

solprop.backends.mlx_gpu.arithmetic.GEMM_PROBE_BOUND = 7.62939453125e-06#

The normwise relative error a float32 product must stay below (SEP 33 7.1).

Float32 products of inner dimension up to 10**3 are expected near sqrt(n) u <= 2e-6 and TF32 near its own unit roundoff 2**-11: the bound separates the two by more than a factor of four on each side.

solprop.backends.mlx_gpu.arithmetic.fft_bound_32(samples)[source]#

Return the established transform bound at binary32, 512 eps32 ceil(log2 n).

The ladder’s per-transform envelope (tests/bounds.py, fft_bound) evaluated at binary32’s machine epsilon (SEP 33 7.1).

Parameters:

samplesint: The transform length.

Returns:

float: The bound, relative and normwise.

solprop.backends.mlx_gpu.arithmetic.fft_probe(lengths)[source]#

Measure the device’s transforms at a run’s own kinds and lengths (SEP 33 6.3).

A Gaussian random signal of each length is transformed forward and back on the device and compared with binary64: the forward transform normwise, and the round trip, whose scale bias MLX’s GPU transform carries, against the identity.

Parameters:

lengthslist of tuple: (kind, n), kind "field" (real transforms) or "envelope".

Returns:

float: The largest relative error found.

Raises:

BackendError: If a transform exceeds its bound at binary32.

solprop.backends.mlx_gpu.arithmetic.gemm_probe(shapes)[source]#

Measure float32 matrix products of each dispatch family a run uses.

Gaussian random factors of each (m, k, n) are multiplied on the device and compared with binary64 (SEP 33 6.3, 7.1).

Parameters:

shapeslist of tuple of int: The products’ shapes.

Returns:

float: The largest normwise relative error.

Raises:

BackendError: If any product exceeds GEMM_PROBE_BOUND: the device computes in TF32 or lower, which SEP 25 forbids.

solprop.backends.mlx_gpu.arithmetic.on_gpu()[source]#

Return whether MLX’s default device is the Metal GPU.

Returns:

bool: True on the GPU, False on MLX’s CPU device.

solprop.backends.mlx_gpu.arithmetic.population(exposure, preionisation)[source]#

Return the ionised fraction and the neutral survival in their stable forms.

The oracle’s stable_population (SEP 33 section 4.4): f = p + (1 - p) (-expm1(-X)) and sigma = (1 - p) exp(-X). On the Metal device it is a custom kernel in "safe" math mode with Kahan’s expm1; on MLX’s CPU device, whose exp errs by up to 3e-6 relative at large arguments, NumPy’s binary64 forms rounded once.

Parameters:

exposuremlx.core.array: X, float32, non-negative.
preionisationfloat or mlx.core.array: p in [0, 1], or on the Metal device the kernel’s parameters (population_parameters()) already on it, as a compiled stage passes them.

Returns:

tuple of mlx.core.array: The fraction and the survival, float32.

solprop.backends.mlx_gpu.arithmetic.population_parameters(preionisation)[source]#

Return the population kernel’s parameters, (1 - p, p) rounded to binary32.

Parameters:

preionisationfloat: p in [0, 1].

Returns:

numpy.ndarray: Float32, (2,).

solprop.backends.mlx_gpu.arithmetic.prefix_scan(values)[source]#

Return the oracle’s doubling scan along the last axis, in one launch.

Parameters:

valuesmlx.core.array: Float32 increments.

Returns:

mlx.core.array: The running sums, the oracle’s to the bit.

solprop.backends.mlx_gpu.arithmetic.rate_kernel(magnitude, coefficients, constant_tail, inverse_edges, inverse_tail, constants)[source]#

Evaluate a rate table’s spline at field magnitudes in one launch.

Parameters:

magnitudemlx.core.array: |E|, float32.
coefficientsmlx.core.array: The spline’s (4, n_intervals) coefficients flattened row by row, the constant term’s head last.
constant_tailmlx.core.array: The constant term’s tail.
inverse_edgesmlx.core.array: 1 / E_k’s head.
inverse_tailmlx.core.array: 1 / E_k’s tail.
constantsmlx.core.array: ln E_0, 1 / d ln E, E_min, E_max, W_max, float32.

Returns:

mlx.core.array: The rate, float32, of magnitude’s shape.