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, which continuous integration uses, 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.
`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[, host])

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

class solprop.backends.mlx_gpu.arithmetic.DoubleFloatTerms(shape, n_terms, parts, host=None)[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.

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.

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#

n_terms: int#

property nbytes: int[source]#

Return the device bytes the terms hold.

Returns:

int: Their size.

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

shape: tuple[int, ...]#

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: p in [0, 1].

Returns:

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

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.