Abstract compute backend interface.
One concrete IBackend exists per Device slot (CPU,
GPU); Dispatcher holds them and routes every op based on the
tensor's device tag. Implementations dispatch the actual math to
their device-native library — Apple Accelerate (vDSP / vForce / BLAS
/ LAPACK) for CPU, MLX for GPU.
See Also
CpuBackend — Accelerate-backed concrete CPU.
GpuBackend — MLX-backed concrete GPU.
Dispatcher — selects which to use.
Notes
Ownership. Each concrete backend is created once and stored
inside the Dispatcher singleton via std::unique_ptr; the
lifetime of every IBackend therefore equals the lifetime of
the process.
Thread safety. Methods are thread-compatible but not internally synchronised: the engine never calls into the same backend from multiple threads simultaneously on the same op, though distinct ops on distinct tensors may overlap. Callers that share a backend across threads must provide their own serialisation, or rely on the underlying framework (e.g. MLX's own thread-safety guarantees).
Design rationale. All methods take Storage by
const-reference and return a new Storage; no in-place
mutation is exposed at this layer. Concrete implementations choose
the most efficient internal representation — CpuStorage (Apple
aligned host buffer) for CPU, GpuStorage (mlx::core::array)
for GPU.
Destructor
1Virtual methods
13cross_entropy_loss
ClassLossForwardResult cross_entropy_loss(const int & input, const int & target, const int * weight, const int & input_shape, const int & target_shape, Dtype dt, double eps, int ignore_index, int reduction)Cross-entropy loss with optional per-class weights and ignore_index. saved_aux inside the result holds the log-softmax activations needed by the backward pass. weight is optional (may be nullptr for unweighted CE).
Returns the Device tag associated with this backend.
Returns
DeviceDevice::CPU for CpuBackend, Device::GPU for GpuBackend.
linalg_eig
StoragePair linalg_eig(const int & a, const int & shape, const int & values_shape, const int & vectors_shape, Dtype dt)General eigendecomposition for non-symmetric matrices; returns (eigenvalues, eigenvectors). Eigenvalues may be complex; CPU backend returns packed real+imag components.
linalg_eigh
StoragePair linalg_eigh(const int & a, const int & shape, const int & values_shape, const int & vectors_shape, Dtype dt)Symmetric/Hermitian eigendecomposition (real eigenvalues guaranteed). Returns (eigenvalues, eigenvectors); eigenvectors are column-ordered.
linalg_ldl_factor
StoragePair linalg_ldl_factor(const int & a, const int & shape, Dtype dt)LDL^T factorization of symmetric matrix. Returns {LD_packed, pivots_i32}.
linalg_lu_factor
StoragePair linalg_lu_factor(const int & a, const int & shape, Dtype dt)LU factorisation (packed format). Returns {LU, pivots} where LU is the packed n×n matrix (LAPACK dgetrf_ format) and pivots is an n-element I32 tensor of 1-based pivot indices.
linalg_qr
StoragePair linalg_qr(const int & a, const int & shape, const int & q_shape, const int & r_shape, Dtype dt)Reduced QR decomposition; returns (Q, R).
nll_loss
ClassLossForwardResult nll_loss(const int & input, const int & target, const int * weight, const int & input_shape, const int & target_shape, Dtype dt, int ignore_index, int reduction)Negative log-likelihood loss; operates on log-probabilities directly (unlike cross_entropy which applies log-softmax internally).
nonzero_forward
CpuStorage nonzero_forward(const int & input, const int & input_shape, Dtype input_dtype, int & numel_out)Returns a CpuStorage containing the flat indices of all nonzero elements. numel_out is set to the number of nonzero elements found.
permute_cpu
CpuStorage permute_cpu(const CpuStorage & src, const int & src_shape, const int & perm, Dtype dt)Permutes a CpuStorage buffer in-place and returns the result as a new CpuStorage. Used for data layout conversions before GPU upload.
rms_norm_backward
StoragePair rms_norm_backward(const int & x, const int & gamma, const int & saved_rstd, const int & grad, int outer, int normalized_size, const int & x_shape, const int & gamma_shape, Dtype dt)RMSNorm backward; returns (grad_x, grad_gamma).
rms_norm_forward
StoragePair rms_norm_forward(const int & x, const int & gamma, int outer, int normalized_size, double eps, const int & x_shape, Dtype dt)RMSNorm forward pass. outer = product of batch/sequence dims; normalized_size is the last (feature) dimension. Returns (output, saved_rstd); saved_rstd holds per-row reciprocal standard deviations needed by the backward pass.
CpuStorage to_cpu(const int & a, const int & shape)Materialises the tensor's data in CPU-accessible memory.
Parameters
aconst Storage&shapeconst Shape&a.Returns
CpuStorageHost-side aligned buffer containing the evaluated data.
Notes
For CpuBackend this is effectively a copy of the existing
CpuStorage; for GpuBackend it calls
mlx::core::array::eval() first (triggering execution of the
pending lazy graph) and then copies the result into a fresh
CpuStorage on the host.
Methods
227Elementwise arithmetic on two broadcast-prepared operands.
Each of the five primitives below computes out[i] = a[i] op b[i]
for every element, with both a and b already broadcast (by
the caller) to the common output shape.
Math
a[i] + b[i] & \text{add} \\ a[i] - b[i] & \text{sub} \\ a[i] \cdot b[i] & \text{mul} \\ a[i] / b[i] & \text{div} \\ a[i]^{b[i]} & \text{pow} \end{cases}$$Parameters
a, bconst Storage&shape.shapeconst Shape&numel computation).dtDtypea, b, and the result).Returns
StorageNewly-allocated output buffer of the same shape and dtype.
Broadcasts a scalar to the array and adds/multiplies element-wise.
affine_grid_backward
int affine_grid_backward(const int & grad_grid, int N, int H, int W, bool align_corners, Dtype dt)Affine grid backward: returns grad_theta of shape (N, 2, 3).
affine_grid_forward
int affine_grid_forward(const int & theta, int N, int H, int W, bool align_corners, Dtype dt)Affine grid generation for spatial transformer networks. theta is an (N, 2, 3) affine matrix; output has shape (N, H, W, 2).
Boolean reductions — reduce the full tensor to a scalar bool.
arg_reduce_index
int arg_reduce_index(const int & a, const int & shape, int axis, bool keepdims, Dtype dt, bool is_min)Returns the index of the max (is_min=false) or min (is_min=true) along axis. Output dtype is I64.
Returns the permutation indices that would sort a along axis.
Cast elements to a different dtype; shape is unchanged. CPU: element-wise static_cast loop. GPU: mlx::core::astype.
avg_pool_nd_backward
int avg_pool_nd_backward(const int & grad_out, const int & x_shape, const int & out_shape, const PoolOpts & opts, Dtype dt)N-D average pooling backward; distributes grad uniformly across each window.
avg_pool_nd_forward
int avg_pool_nd_forward(const int & x, const int & x_shape, const int & out_shape, const PoolOpts & opts, Dtype dt)N-D average pooling forward (count-include-pad semantics by default).
batch_norm_backward
int batch_norm_backward(const int & x, const int & gamma, const int & saved_mean, const int & saved_rstd, const int & saved_xnorm, const int & grad, int batch, int channels, int spatial, int ndim, const int & x_shape, Dtype dt, double eps)BatchNorm backward; returns [grad_x, grad_gamma, grad_beta]. eps matches the forward's epsilon — GPU backend needs it to reconstruct variance from saved_rstd (var = 1/rstd^2 - eps) for the MPSGraph normalizationGradient* path. CPU backend may ignore.
3.4+ Phase A.4: saved_xnorm optionally carries the forward's
normalised input (x - mean) * rstd as a full-size tensor (same
shape as x). When present (i.e. non-empty Storage) the MLX-path
backward consumes it directly and skips two element-wise ops per
BN backward (centered = x - mean; xnorm = centered * rstd). Passing
an empty Storage triggers the prior recomputation path — required
for the MPSGraph dispatch, which has its own xnorm management.
batch_norm_eval_backward
int batch_norm_eval_backward(const int & x, const int & mean, const int & gamma, const int & rstd, const int & grad_out, const int & x_shape, int C, int spatial, Dtype dt)BatchNorm inference-mode backward; returns [grad_x, grad_gamma, grad_beta].
batch_norm_eval_forward
int batch_norm_eval_forward(const int & x, const int & mean, const int & var, const int & gamma, const int & beta, const int & x_shape, int C, int spatial, double eps, Dtype dt)BatchNorm inference-mode forward; uses running statistics (mean, var) rather than computing batch statistics. Returns [output, rstd].
batch_norm_forward
int batch_norm_forward(const int & x, const int & gamma, const int & beta, int batch, int channels, int spatial, int ndim, double eps, const int & x_shape, Dtype dt)BatchNorm training forward. ndim distinguishes 1-D, 2-D, and 3-D spatial tensors. Returns [output, saved_mean, saved_rstd, saved_xnorm]. The 4th element (xnorm = (x - mean) * rstd) is the normalised input that backward needs; the MLX forward path already materialises it as a lazy intermediate, so exposing it costs nothing at forward time and saves a recomputation at backward time. Backends without a meaningful xnorm intermediate (e.g. the MPSGraph dispatch) may return an empty Storage for slot 3 — the backward will then fall back to recomputing.
bitwise_binary
→Storageint bitwise_binary(const int & a, const int & b, const int & shape, Dtype dt, int op)Elementwise bitwise operation selected by integer op code.
Parameters
a, bconst Storage&shape.shapeconst Shape&dtDtypeopint0=AND, 1=OR, 2=XOR (GPU also supports 3=left_shift, 4=right_shift).Returns
StorageNewly-allocated result of the same shape and dtype.
broadcast
int broadcast(const int & a, const int & src_shape, const int & dst_shape, Dtype dt)Broadcasts a from src_shape to dst_shape following NumPy rules.
broadcast_back_for_reduce
int broadcast_back_for_reduce(const int & grad, const int & grad_shape, const int & input_shape, const int & axes, bool keepdims, Dtype dt)Broadcasts grad back to input_shape after a reduction. The inverse of reduce_sum when keepdims=false.
Type-casts the buffer element-by-element from src_dt to dst_dt.
Clamps each element to [min_v, max_v].
Returns a deep copy of src reinterpreted with the given shape/dtype.
Parameters
srcconst Storage&shape and dt.shapeShapedtDtypeReturns
StorageA fresh, independent allocation containing the same bytes.
compare_binary
→Storageint compare_binary(const int & a, const int & b, const int & shape, Dtype dt, int op)Elementwise comparison returning a Bool Storage of the same shape.
Parameters
a, bconst Storage&shape.shapeconst Shape&dtDtypeopint0=EQ, 1=NE, 2=GT, 3=GE, 4=LT, 5=LE.Returns
StorageBool buffer of shape shape.
Build a C64 array from two F32 arrays of the same shape; the result satisfies complex_real(out) == re and complex_imag(out) == im.
Element-wise complex conjugate. No-op for real dtypes (returned unchanged); negates the imaginary part for C64.
Imaginary part of a complex (C64) input — output dtype is F32.
Real part of a complex (C64) input — output dtype is F32.
Concatenates the tensors in xs along axis; each shape in shapes corresponds element-wise to xs.
contiguous
→Storageint contiguous(const int & src, const int & shape, const int & stride, int storage_offset, bool already_contiguous, Dtype dt)Returns a contiguous copy of a (possibly strided or offset) view.
Walks the stride / offset layout to produce a densely-packed output buffer suitable for kernels that assume row-major contiguous data.
Parameters
srcconst Storage&shapeShapestrideconst Stride&storage_offsetstd::size_tsrc where the view begins.already_contiguousboolmemcpy path when combined with storage_offset == 0.dtDtypeReturns
StorageDensely-packed copy of the view.
cross_entropy_backward
int cross_entropy_backward(const int & saved_softmax, const int & target, const int * weight, const int & valid_count, const int & grad, const int & input_shape, Dtype dt, int ignore_index, int reduction)Cross-entropy backward; saved_softmax is the softmax output stored during the forward. valid_count accounts for ignored-index samples in mean reduction.
ctc_loss_forward
int ctc_loss_forward(const int & log_probs, const int & targets, const int & input_lengths, const int & target_lengths, const int & lp_shape, int blank, bool zero_infinity, Dtype dt)CTC (Connectionist Temporal Classification) loss. log_probs: (T, N, C) log-probabilities; targets: (N, S) or flat (sum_S,) input_lengths: (N,); target_lengths: (N,). Returns per-sample losses of shape (N,) before reduction.
Scan operations along a single axis.
Extracts a 1-D diagonal at offset k from a 2-D matrix, or places a 1-D vector on a 2-D diagonal. out_shape is written with the result shape.
diagonal
int diagonal(const int & a, const int & input_shape, int offset, int axis1, int axis2, Dtype dt)Extracts a generalized diagonal between axis1 and axis2 with an integer offset from the main diagonal (positive = above, negative = below).
diagonal_backward
int diagonal_backward(const int & grad, const int & input_shape, const int & output_shape, int offset, int axis1, int axis2, Dtype dt)Backward pass for diagonal: scatters the diagonal gradient back into a zero-filled tensor of input_shape.
embedding_backward
int embedding_backward(const int & grad_out, const int & indices, const int & weight_shape, const int & indices_shape, int padding_idx, Dtype dt)Backward pass for embedding: scatter-adds grad_out rows into a zero-filled weight gradient of weight_shape.
embedding_bag_forward
int embedding_bag_forward(const int & weight, const int & indices, const int & offsets, const int & weight_shape, const int & indices_shape, int mode, int padding_idx, bool include_last_offset, Dtype dt)EmbeddingBag: gather rows from weight at indices, then reduce per bag. mode: 0=sum, 1=mean, 2=max. For 1-D indices, offsets marks bag starts.
embedding_forward
int embedding_forward(const int & weight, const int & indices, const int & weight_shape, const int & indices_shape, const int & out_shape, int padding_idx, Dtype dt)Embedding table lookup: out[i] = weight[indices[i]]. padding_idx entries are zeroed in the output.
Elementwise unary math ops.
The group below — exp, log, sqrt,
rsqrt, abs, neg, sign,
floor, ceil, round, sin,
cos, tanh, sigmoid, relu — all
share the same call shape. Each applies the named scalar
function element-by-element to the flattened buffer; shape is
used only for numel computation.
Parameters
aconst Storage&shapeconst Shape&dtDtypeReturns
StorageNewly-allocated output of the same shape and dtype.
Notes
Dispatch: CPU implementations call vForce (transcendentals) or
vDSP (algebraic ops); GPU implementations call the matching
mlx::core::* primitive.
Constructs an M×N identity-like matrix with the diagonal at offset k.
Flip (reverse) along the given axes.
Integer floor division: out[i] = floor(a[i] / b[i]).
Transfers a CPU buffer into the backend's native storage format.
Parameters
cpuCpuStorageshapeShapemlx::core::array shape; CPU is shape-agnostic at this layer.Returns
StorageBackend-native storage variant.
Notes
For CpuBackend this is a no-op move into the
CpuStorage slot of the variant; for GpuBackend it
copies the data to GPU-private memory via mlx::core::copy.
Allocates a tensor filled with fill_value.
fused_linear_relu_forward
int fused_linear_relu_forward(const int & x, const int & w, const int & b, const int & out_shape, Dtype dt)Optional fused kernels; default implementations call the base backend's linear() followed by relu/gelu and throw if the dtype is unsupported. Backends that can fuse these two ops (e.g. ANE) override them.
gather
int gather(const int & a, const int & indices, const int & input_shape, const int & output_shape, int axis, Dtype index_dtype, Dtype dt)Gathers slices from a along axis using integer indices. index_dtype is the dtype of the indices tensor (I32 or I64).
gather_backward
int gather_backward(const int & grad, const int & indices, const int & input_shape, const int & output_shape, int axis, Dtype index_dtype, Dtype dt)Backward pass for gather: scatters grad values back to the source positions identified by indices. Result has input_shape.
Element-wise masks that return 0.0/1.0 float results. ge_mask: out[i] = (a[i] >= b[i]) ? 1 : 0 lt_mask: out[i] = (a[i] < b[i]) ? 1 : 0
gelu_exact — Gaussian-CDF formulation: 0.5x(1 + erf(x/sqrt(2))). Distinct from gelu (tanh approximation) so callers can pick either numerical recipe; Python F.gelu(approximate="none") routes here.
global_response_norm_backward
int global_response_norm_backward(const int & x, const int & gamma, const int & beta, const int & saved_Nx, const int & grad_out, const int & x_shape, double eps, Dtype dt)GRN backward; returns [grad_x, grad_gamma, grad_beta].
global_response_norm_forward
int global_response_norm_forward(const int & x, const int & gamma, const int & beta, const int & x_shape, double eps, Dtype dt)Global Response Normalization (GRN) from ConvNeXt-v2. Returns [output, saved_Nx] where Nx is the per-channel L2 norm.
grid_sample_backward
int grid_sample_backward(const int & grad_out, const int & input, const int & grid, const int & in_shape, const int & grid_shape, int mode, int padding_mode, bool align_corners, Dtype dt)Grid sample backward; returns [grad_input, grad_grid].
grid_sample_forward
int grid_sample_forward(const int & input, const int & grid, const int & in_shape, const int & grid_shape, int mode, int padding_mode, bool align_corners, Dtype dt)Samples input at grid locations using bilinear (mode=0) or nearest (mode=1) interpolation. padding_mode: 0=zeros, 1=border, 2=reflection.
group_norm_backward
int group_norm_backward(const int & x, const int & gamma, const int & saved_mean, const int & saved_rstd, const int & grad, int batch, int channels, int spatial, int groups, const int & spatial_dims, const int & x_shape, Dtype dt)GroupNorm backward; returns [grad_x, grad_gamma, grad_beta].
group_norm_forward
int group_norm_forward(const int & x, const int & gamma, const int & beta, int batch, int channels, int spatial, int groups, const int & spatial_dims, double eps, const int & x_shape, Dtype dt)GroupNorm forward; groups partitions channels into equal groups before normalizing each group independently. Returns [output, saved_mean, saved_rstd].
histogram_forward
int histogram_forward(const int & input, const int & input_shape, Dtype input_dtype, double lo, double hi, int bins, bool density)Computes a histogram over input with bins equal-width buckets in [lo, hi]. density=true normalises the result to a probability density.
huber_loss
int huber_loss(const int & input, const int & target, const int & shape, Dtype dt, double delta, int reduction)Huber loss (smooth L1): quadratic for |r|<=delta, linear otherwise. delta controls the transition point between the two regimes. Backward returns (grad_input, grad_target).
in_range_mask
int in_range_mask(const int & a, const int & shape, Dtype dt, double lo, double hi)Returns a 0/1 mask where a[i] is in the closed interval [lo, hi].
inner
int inner(const int & a, const int & b, const int & a_shape, const int & b_shape, const int & out_shape, Dtype dt)Inner product generalised to batches: contracts the last axis of a against the last axis of b.
insert_axis_slice
int insert_axis_slice(const int & a, const int & src_shape, const int & dst_shape, int axis, int offset, Dtype dt)Inverse of slice_axis: scatters a into a zero-filled tensor of dst_shape at the given offset along axis. Used in slice backward passes.
interpolate_bilinear_forward
int interpolate_bilinear_forward(const int & input, const int & in_shape, int H_out, int W_out, bool align_corners, Dtype dt)Bilinear interpolation forward/backward. align_corners=true maps corner pixels to corner coordinates; false scales by factor (in_size / out_size).
interpolate_nearest_2d_forward
int interpolate_nearest_2d_forward(const int & input, const int & in_shape, int H_out, int W_out, Dtype dt)Nearest-neighbor upsampling for 2-D and 3-D feature maps.
interpolate_trilinear_forward
int interpolate_trilinear_forward(const int & input, const int & in_shape, int D_out, int H_out, int W_out, bool align_corners, Dtype dt)Trilinear interpolation for volumetric feature maps.
Bitwise NOT.
Parameters
aconst Storage&shapeconst Shape&dtDtypeReturns
StoragePer-element bitwise complement.
Floating-point predicate ops (output dtype is always Bool).
layer_norm_backward
int layer_norm_backward(const int & x, const int & gamma, const int & saved_mean, const int & saved_rstd, const int & grad, int outer, int normalized_size, const int & x_shape, const int & gamma_shape, const int & beta_shape, Dtype dt)LayerNorm backward; returns [grad_x, grad_gamma, grad_beta].
layer_norm_forward
int layer_norm_forward(const int & x, const int & gamma, const int & beta, int outer, int normalized_size, double eps, const int & x_shape, Dtype dt)LayerNorm forward. Returns [output, saved_mean, saved_rstd].
Returns slope*a[i] when a[i]<0, else a[i]. Used inside leaky-relu backward.
linalg_cholesky
int linalg_cholesky(const int & a, const int & shape, bool upper, Dtype dt)Cholesky decomposition. upper=true returns the upper triangular factor.
Determinant (scalar per batch element).
linalg_householder_product
int linalg_householder_product(const int & H, const int & tau, const int & h_shape, Dtype dt)Reconstruct Q (m×k) from Householder reflectors H (m×n) and tau (k,).
Matrix inverse using LU factorisation.
linalg_lstsq
int linalg_lstsq(const int & a, const int & b, const int & a_shape, const int & b_shape, Dtype dt)Least-squares: solve min||AX-B||_2. Returns [solution, residuals, rank, svd]. a_shape=(m,n), b_shape=(m,nrhs). Solution shape=(n,nrhs).
linalg_lu_solve
int linalg_lu_solve(const int & LU, const int & pivots, const int & b, const int & lu_shape, const int & b_shape, Dtype dt)Solve AX=B given LU+pivots from linalg_lu_factor. Returns X (same shape as B).
linalg_matrix_power
int linalg_matrix_power(const int & a, const int & shape, int power, Dtype dt)A^power computed via repeated squaring or eigendecomposition.
linalg_norm
int linalg_norm(const int & a, const int & shape, double ord, const int & axes, bool keepdims, Dtype dt)Vector/matrix/Frobenius norm; ord controls the norm type.
Moore-Penrose pseudoinverse via SVD.
linalg_solve
int linalg_solve(const int & a, const int & b, const int & a_shape, const int & b_shape, Dtype dt)Solves a linear system Ax = b; b may have multiple right-hand sides.
linalg_solve_triangular
int linalg_solve_triangular(const int & a, const int & b, const int & a_shape, const int & b_shape, bool upper, bool unitriangular, Dtype dt)Triangular solve: solve A X = B where A is triangular. upper=true → upper triangular; unit=true → unit diagonal.
linalg_svd
int linalg_svd(const int & a, const int & shape, bool compute_uv, const int & u_shape, const int & s_shape, const int & vt_shape, Dtype dt)Singular value decomposition. compute_uv=false returns only singular values S. Returns [U, S, Vt] or [S] depending on compute_uv.
linear
→Storageint linear(const int & x, const int & weight, const int & bias, const int & x_shape, const int & weight_shape, const int & out_shape, Dtype dt)Fused linear projection .
Parameters
xconst Storage&(*, K).weightconst Storage&(N, K) (note: pre-transposed).biasconst Storage&(N,).x_shape, weight_shape, out_shapeconst Shape&dtDtypeReturns
StorageOutput activations of shape (*, N).
linear_backward
int linear_backward(const int & grad, const int & x, const int & weight, const int & x_shape, const int & weight_shape, const int & bias_shape, Dtype dt)Backward pass for linear; returns [grad_x, grad_weight, grad_bias].
log_softmax
→Storageint log_softmax(const int & a, const int & shape, int axis, Dtype dt)Numerically stable log-softmax along a single axis.
Math
computed via the max-subtraction trick to avoid overflow.
Parameters
aconst Storage&shapeconst Shape&a.axisintdtDtypeReturns
StorageLog-probabilities of the same shape and dtype as a.
log_softmax_backward
→Storageint log_softmax_backward(const int & y, const int & grad_out, const int & shape, int axis, Dtype dt)Backward pass for log_softmax.
Math
\frac{\partial \mathcal{L}}{\partial y} - \exp(y) \sum_{\text{axis}}\frac{\partial \mathcal{L}}{\partial y},$$ where `y = log_softmax(x)` is the saved forward output.Parameters
yconst Storage&grad_outconst Storage&shapeconst Shape&y, grad_out, and the returned input gradient.axisintdtDtypeReturns
StorageInput gradient .
Additional elementwise unary ops.
Logarithmic (log2), error function (erf,
erfinv), reciprocal / power (reciprocal,
square, cube, cube_root), trigonometric
(tan, asin, acos, atan,
sinh, cosh) that are not part of the core unary
set above. Same call shape and dispatch contract as that group.
lp_normalize_backward
int lp_normalize_backward(const int & x, const int & saved_norm, const int & grad_out, const int & x_shape, double ord, int axis, Dtype dt)Lp-normalize backward; saved_norm is the per-row norm from the forward.
lp_normalize_forward
int lp_normalize_forward(const int & x, const int & x_shape, double ord, int axis, double eps, Dtype dt)Lp-normalizes x along axis using the given norm order ord. Returns [output, saved_norm] so the backward can reuse the computed norm.
lstm_backward
int lstm_backward(const int & grad_output, const int & grad_hn, const int & grad_cn, const int & input, const int & h0, const int & weights, const int & gates_all, const int & cells_all, const LstmOpts & opts, Dtype dt)LSTM BPTT backward. gates_all and cells_all are the saved activations from lstm_forward_train. Returns [dX, dh0, dc0, dW_ih, dW_hh, db_ih, db_hh].
lstm_forward
int lstm_forward(const int & input, const int & h0, const int & c0, const int & weights, const LstmOpts & opts, const int & out_shape, Dtype dt)LSTM inference forward. Returns [output, hn, cn]. Default implementation throws not_implemented; concrete backends override.
lstm_forward_train
int lstm_forward_train(const int & input, const int & h0, const int & c0, const int & weights, const LstmOpts & opts, Dtype dt)LSTM training forward; saves gate and cell activations for BPTT. Returns [output, hn, cn, gates_all, cells_all].
masked_fill
int masked_fill(const int & a, const int & mask, const int & shape, Dtype dt, double value)Replaces elements of a where mask is nonzero with value.
masked_select_count
int masked_select_count(const int & mask, const int & shape, Dtype dt)Masked select: extract elements where bool mask == true. Returns a flat 1-D Storage of n_true elements.
matmul
→Storageint matmul(const int & a, const int & b, const MatmulOpts & opts, Dtype dt)General batched matrix multiplication.
Math
where a and b may have a leading batch dimension and may
be transposed before contraction.
Parameters
a, bconst Storage&optsconst MatmulOpts&M, K, N, batch) and transpose flags.dtDtypeReturns
StorageResult of shape (batch?, M, N).
Notes
CPU dispatches to cblas_sgemm / cblas_dgemm via
Accelerate; GPU dispatches to mlx::core::matmul (which selects
an MPS or Metal-shader kernel internally).
max_pool_nd_backward
int max_pool_nd_backward(const int & grad_out, const int & saved_argmax, const int & x_shape, const int & out_shape, const PoolOpts & opts, Dtype dt)N-D max pooling backward; scatters grad_out to the argmax positions.
max_pool_nd_forward
int max_pool_nd_forward(const int & x, const int & x_shape, const int & out_shape, const PoolOpts & opts, Dtype dt)N-D max pooling forward. Returns [output, argmax]; argmax stores the flat index within the spatial window that produced each maximum, required by max_pool_nd_backward to route gradients.
Elementwise max(a, b) and min(a, b).
Parameters
a, bconst Storage&shape.shapeconst Shape&dtDtypeReturns
StorageElement-wise max or min, same shape and dtype.
Constructs N-D coordinate grids from 1-D input vectors. indexing_xy=true uses x-column / y-row convention (MATLAB-style).
mse_loss
int mse_loss(const int & input, const int & target, const int & shape, Dtype dt, int reduction)Mean-squared error: 0.5*(input - target)^2 summed/averaged. Backward returns (grad_input, grad_target).
nan_to_num
int nan_to_num(const int & a, const int & shape, Dtype dt, double nan_val, double posinf_val, double neginf_val)Replace NaN/Inf values with finite substitutes. Output dtype matches input.
nll_loss_backward
int nll_loss_backward(const int & target, const int * weight, const int & valid_count, const int & grad, const int & input_shape, Dtype dt, int ignore_index, int reduction)NLL backward; returns the input gradient.
nn_fold
int nn_fold(const int & x, const int & x_shape, const int & out_shape, const int & kernel_size, const int & stride, const int & padding, const int & dilation, Dtype dt)Fold (col2im): inverse of unfold/im2col. Accumulates (N, CkHkW, L) patches back into (N, C, outH, outW) using scatter-add.
one_hot_forward
int one_hot_forward(const int & indices, const int & indices_shape, int num_classes, Dtype out_dtype)One-hot encoding: produces a (num_classes,)-wide indicator tensor for each index in indices. Output dtype is typically F32 or I32.
Allocates a one-filled buffer of shape shape with element type dt.
Parameters
shapeShapedtDtypeReturns
StorageNewly-allocated buffer where every element equals 1.
Zero-pads (or reflects/wraps) the tensor with pad_width[d] = (before, after) extra elements for each dimension d. constant specifies the fill value.
Permutes dimensions according to perm (equivalent to np.transpose).
Returns 1.0 where a[i] > 0, else 0.0. Used in relu backward.
Raises each element to a scalar exponent: a**exp.
reduce_broadcast
int reduce_broadcast(const int & grad, const int & input_shape, const int & output_shape, Dtype dt)Sums grad over the broadcast dimensions so that the result has output_shape (the pre-broadcast shape).
reduce_grad_to_shape
int reduce_grad_to_shape(const int & grad, const int & grad_shape, const int & target_shape, Dtype dt)Reduces grad (of shape grad_shape) to target_shape by summing over the broadcast dimensions. Used in broadcast backward passes.
reduce_sum
int reduce_sum(const int & a, const int & in_shape, const ReduceOpts & opts, Dtype dt)Axis reductions. opts.axes specifies which dimensions to collapse; the output shape follows keepdims semantics. variance() uses the biased (population) variance formula to match NumPy/reference defaults.
relu_backward
int relu_backward(const int & g, const int & x, const int & shape, Dtype dt)Fused ReLU backward: returns g where x > 0, else 0. Single- kernel alternative to multiply(g, positive_mask(x)) — collapses the 3-op chain (greater + astype + multiply) into one MLX expression that MLX can fuse natively (or that GPU dispatches as where). Falls back to the 2-step composition on backends without a fused path.
ReLU clamped to [0, 6], used in MobileNet-style networks.
Repeats elements repeats times along axis (equivalent to np.repeat).
repeat_backward
int repeat_backward(const int & grad_out, const int & input_shape, const int & output_shape, int axis, int repeats, Dtype dt)Backward pass for repeat(): reduces repeated slices by summing them back to the original input_shape.
reverse_along_axis
int reverse_along_axis(const int & a, const int & shape, int axis, Dtype dt)Reverses elements along axis (equivalent to a[::-1] on that dimension).
rope_backward
int rope_backward(const int & grad_out, const int & saved_cos, const int & saved_sin, const int & x_shape, bool interleaved, Dtype dt)RoPE backward; applies the inverse rotation using cached cos/sin values.
rope_forward
int rope_forward(const int & x, const int * position_ids, const int & x_shape, bool interleaved, Dtype pos_dtype, Dtype dt)Rotary Position Embedding (RoPE) forward. Returns [output, cos_cache, sin_cache]. interleaved=true uses GPT-NeoX layout; false uses Llama layout. position_ids is optional; if null, positions 0..seq_len-1 are assumed.
rotate_forward
int rotate_forward(const int & input, const int & shape, double angle_rad_neg, double cx, double cy, Dtype dt)Rotates a 2-D image by angle_rad_neg (stored negated for efficiency) around center (cx, cy) using bilinear interpolation.
Raises a scalar base to the power of each element: base**a.
run_custom_metal_kernel
int run_custom_metal_kernel(const int & kernel_source, const int & function_name, const int & inputs, const int & output_shape, Dtype output_dtype, const int & grid, const int & threads)Compiles and launches a user-provided Metal Shading Language (MSL) kernel. Default implementation throws not_implemented; only GpuBackend overrides this. grid and threads follow Metal's threadgroups / threadsPerThreadgroup convention.
scatter_add
int scatter_add(const int & base, const int & indices, const int & src, const int & base_shape, const int & idx_shape, int dim, Dtype dt)User-facing scatter-add: out[..., index[i], ...] += src[..., i, ...] along dim. base is unchanged except at index positions where src is added.
scatter_add_axis
int scatter_add_axis(const int & grad, const int & indices, const int & output_shape, const int & grad_shape, int axis, Dtype dt)Accumulates grad values into output_shape positions specified by indices along axis (sparse scatter-add / gather backward).
scatter_amax
int scatter_amax(const int & base, const int & indices, const int & src, const int & base_shape, const int & idx_shape, int dim, Dtype dt)Scatter-reduce variants: base is pre-initialised with the appropriate neutral element (−∞ for amax, +∞ for amin, 1 for prod) when include_self is false; the callers set this up before dispatching.
sdpa_backward
int sdpa_backward(const int & grad_out, const int & q, const int & k, const int & v, const int & saved_weights, const int & q_shape, const int & k_shape, const int & v_shape, double scale, Dtype dt)SDPA backward; returns [grad_q, grad_k, grad_v].
sdpa_forward
int sdpa_forward(const int & q, const int & k, const int & v, const int * attn_mask, const int & q_shape, const int & k_shape, const int & v_shape, Dtype mask_dtype, int mask_numel, double scale, bool is_causal, Dtype dt)Scaled dot-product attention forward. Returns [output, saved_weights] where saved_weights holds the post-softmax attention scores for backward. is_causal=true applies a lower-triangular mask to prevent attending to future tokens without materialising the full mask tensor.
Activation functions with paired forward / backward kernels.
Each forward op below has a matching *_backward that receives
the original pre-activation input a together with the upstream
gradient grad and returns the input gradient .
Parameters
aconst Storage&gradconst Storage&*_backward variants).shapeconst Shape&dtDtypeReturns
StorageForward output, or input gradient for the *_backward form.
Notes
Backward kernels exist as fused single-op delegates so the GPU backend can route to MPSGraph and the CPU backend to hand-rolled scalar loops; previously the autograd composed these from primitive storage ops.
silu_backward
int silu_backward(const int & a, const int & grad, const int & shape, Dtype dt)silu_backward — dL/dx = σ(x) * (1 + x*(1 - σ(x))) * dL/dy. Lucid's autograd node previously composed this from storage primitives; the single-op delegate lets the GPU backend dispatch to MPSGraph and CPU to a hand-rolled scalar loop.
sinusoidal_pos_embedding
int sinusoidal_pos_embedding(int seq_len, int embed_dim, Dtype dt)Generates a fixed sinusoidal positional encoding of shape (seq_len, embed_dim) following the original Transformer paper formula.
Stable numerically-safe softmax forward and its vector-Jacobian product. softmax_backward receives the softmax output z (not the pre-activation input) because the gradient formula only needs the output: dx = z*(dz - (z·dz)).
sort_select
int sort_select(const int & a, const int & input_shape, const int & output_shape, int axis, Dtype dt, bool descending)Returns (sorted_values, sorted_indices) along axis. descending=true sorts from largest to smallest.
split_at
int split_at(const int & a, const int & shape, int axis, const int & indices, Dtype dt)Splits a at the given index boundaries along axis (like np.split).
split_equal
int split_equal(const int & a, const int & shape, int axis, int num_splits, Dtype dt)Splits a into num_splits equal-sized pieces along axis.
Stacks tensors along a new axis inserted at position axis. All inputs must have the same input_shape.
tensordot
int tensordot(const int & a, const int & b, const int & a_shape, const int & b_shape, const int & out_shape, const int & axes_a, const int & axes_b, Dtype dt)General tensor contraction (generalised matmul) along axes_a and axes_b.
Tiles the tensor by replicating it reps[i] times along each dimension.
tile_backward
int tile_backward(const int & grad_out, const int & input_shape, const int & padded_shape, const int & output_shape, const int & reps, Dtype dt)Backward pass for tile(): accumulates gradients from all tiled copies back into the original input_shape.
Moves src to MTLResourceStorageModeShared memory so that both CPU and GPU can access it without a copy. Default is a no-op; GpuBackend overrides.
Sum of the main diagonal elements; trace_backward scatters the upstream gradient back to the diagonal positions of the input shape.
Applies a triangular mask: upper=true zeroes the lower triangle, upper=false zeroes the upper triangle, relative to diagonal offset k.
unfold_dim
int unfold_dim(const int & a, const int & in_shape, int dim, int size, int step, Dtype dt)Sliding-window view along a single dimension. Returns shape (*base.shape[:dim], L, *base.shape[dim+1:], size) where L = (dim_size - size) / step + 1.
where_branch
int where_branch(const int & grad, const int & cond, const int & shape, Dtype dt, bool true_branch)Selects grad where cond is true (true_branch=true) or false, zeroing the other branch. Used by the Where autograd backward.
where_op
int where_op(const int & cond, const int & x, const int & y, const int & shape, Dtype dt)Element-wise selection: out[i] = cond[i] ? x[i] : y[i].
Allocates a zero-filled buffer of shape shape with element type dt.
Parameters
shapeShapedtDtypeReturns
StorageNewly-allocated zero buffer.