NLSE.backends.mlx_backend¶

MLX backend implementation for Apple Silicon GPU acceleration.

Classes¶

MLXBackend

MLX backend for Apple Silicon GPU acceleration.

Module Contents¶

class NLSE.backends.mlx_backend.MLXBackend¶

Bases: NLSE.backends.backend.Backend

MLX backend for Apple Silicon GPU acceleration.

Fuses via mx.compile, which traces a whole step into one graph. MLX always normalizes its inverse FFT, so the 1/N cannot be folded into the propagator.

has_linear_step = True¶

normalizes_on_host = True¶

has_fused_split_step = True¶

broadcasts_parameters_natively = True¶

has_fused_rk4_step = True¶

has_fused_coupled_split_step = True¶

has_fused_coupled_rk4_rhs = True¶

property name: str¶: Backend name identifier.

allocate_field(shape: tuple, dtype: numpy.typing.DTypeLike) → Any¶: Allocate complex array on MLX device.

allocate_real_field(shape: tuple, dtype: numpy.typing.DTypeLike) → Any¶: Allocate real array on MLX device.

property convolution¶

Return an FFT convolution with scipy’s oaconvolve signature.

MLX has no convolution of its own, which is what previously kept the non-local interaction off this backend: nl_length > 0 is gated on this property being non-None. The kernel is a Bessel profile of a few hundred cells at most, so an overlap-add scheme buys nothing over a single padded transform.

synchronize(array=None) → None¶

Force the lazy graph for array.

Parameters:: array (Any, optional) – The array whose value is needed. Nothing to do without one.

to_numpy(array: Any) → numpy.ndarray¶: Transfer from MLX device to CPU.

from_numpy(array: numpy.ndarray) → Any¶

Transfer from CPU to MLX device.

Downcasts float64/complex128 to float32/complex64 since MLX does not support double precision.

fft(array: Any, plan: list) → Any¶: Perform forward FFT.

ifft(array: Any, plan: list, normalize: bool = True) → Any¶

Perform inverse FFT.

MLX’s inverse always normalizes, so this backend does not declare supports_unnormalized_ifft and is never asked to skip it.

exp(array: Any) → Any¶: Exponentiate without leaving this backend.

sum(array: Any) → float¶: Reduce without leaving this backend.

property kernels: Any¶: Return MLX kernels module.

supports_double_precision() → bool¶: MLX does not support double precision.