Variadic-operand einsum computed by MLX's path-optimised mlx::core::einsum kernel.

subscripts is a standard Einstein-summation equation (e.g. "ij,jk->ik", "bij,bjk->bik", "bhid,bhjd->bhij", "ij,jk,kl->il"). operands is the corresponding list of 2+ tensors.

Eager-only, not defn-callable

Unlike the other helpers in this module, einsum/2 does not emit an Nx.Defn.Expr node — it takes refs directly off EMLX.Backend-backed tensors and calls the NIF eagerly, in the same "direct-call helper" style as EMLX.quantized_matmul/2. Every operand must live on EMLX.Backend; anything else raises ArgumentError.

Examples

iex> a = Nx.iota({2, 3}, backend: EMLX.Backend, type: :f32)
iex> b = Nx.iota({3, 4}, backend: EMLX.Backend, type: :f32)
iex> y = EMLX.Fast.einsum("ij,jk->ik", [a, b])
iex> Nx.shape(y)
{2, 4}

Fused decode-time KV-cache write (2× put_slice) + causal-masked SDPA read, as one native op — the traced-graph analog of the hand-written EMLX.kv_cache_sdpa_update/7 eager NIF, but operating on the full fixed-size preallocated cache with a dynamic (array-valued) offset instead of a host int, so it can sit inside a compiled :while decode loop without forcing a host readback per iteration.

All tensors use Bumblebee-native, heads-NOT-transposed layout (matching Bumblebee.Layers.Decoder.attention_cache/4's {batch, seq, heads, head_dim} cache shape and EMLXAxon's build_sdpa_layer/5 pre-transpose inputs) — the head-transpose SDPA itself needs is done internally, once, on q and on the already-written k_upd/v_upd; the returned caches stay untransposed so they can feed a :while carry directly with no extra transpose.

• q — {B, T_q, N_q, D} (T_q = 1 for decode; not enforced)
• new_k — {B, T_q, N_kv, D}
• new_v — {B, T_q, N_kv, D}
• k_cache — {B, T_max, N_kv, D}
• v_cache — {B, T_max, N_kv, D}
• offset — scalar/{1} int tensor — write position along the T_max axis
• key_mask — {B, T_max} — 1 = valid (readable), 0 = masked
• scale — pre-computed scalar

Returns {attn_out, k_cache_updated, v_cache_updated} — attn_out is head-transposed, {B, N_q, T_q, D} (matching what the unfused scaled_dot_product_attention_causal_key_masked/6 metadata node itself already returns; callers transpose it back same as today), while the updated caches keep k_cache/v_cache's own (untransposed) shape.

Eager: composes plain Nx.put_slice ×2 (clamped host offset) + scaled_dot_product_attention_causal_key_masked/6 — deliberately independent of the traced path's C++ formula below, so it doubles as an correctness oracle for the native multi_op_registry["kv_cache_sdpa_update"] kernel (same combined causal+key_mask formula, reimplemented once in Elixir/Nx and once in C++).

Traced: lowers to one multi_op_registry["kv_cache_sdpa_update"] (emlx_compiler.cpp) native IR instruction.

Fused layer normalisation without bias (mlx::fast::layer_norm, weight-only variant).

• x — input tensor; normalised over the last axis.
• weight — {hidden} scale vector (gamma).
• eps — numerical stability constant (e.g. 1.0e-5).

Output shape and type match x.

Fused layer normalisation (mlx::fast::layer_norm).

• x — input tensor; normalised over the last axis.
• weight — {hidden} scale vector (gamma).
• bias — {hidden} bias vector (beta).
• eps — numerical stability constant (e.g. 1.0e-5).

Output shape and type match x.

Fused RMS normalisation (mlx::fast::rms_norm).

• x — input tensor; normalised over the last axis.
• weight — {hidden} scale vector (same size as last axis of x).
• eps — numerical stability constant (e.g. 1.0e-6).

Output shape and type match x.

Fused rotary position embedding (mlx::fast::rope).

• a — input {B, ..., T, D}; ... dims are passed through, but T must be the second-to-last axis (e.g. {B, H, T, D}, heads-transposed) — see EMLX.fast_rope/6's @doc for why a heads-axis placed before T (Bumblebee's {B, T, H, D}) silently miscomputes for more than one head (elixir-nx/emlx#121). Use rope_with_positions/6 or rope_with_freqs/6 for that layout instead.
• dims — number of feature dims to rotate (≤ last-axis size, must be even).
• traditional — false for split-half (Qwen3); true for interleaved.
• base — angular frequency base (e.g. 10_000 or 1_000_000).
• scale — position scale (1.0 unless using NTK-aware scaling).
• offset — integer position offset (tokens already in the KV cache).

traditional must match the model checkpoint's convention. For Qwen3 (split-half): traditional: false.

Output shape and type match a.

Fused RoPE with precomputed inverse-frequency vector (mlx::fast::rope, freqs overload).

Use this variant when the model's RoPE scaling strategy produces a fixed {dims/2} inv-frequency tensor that can be baked at graph-rewrite time (e.g. :llama3 smooth-interpolation). Strategies that are seq-len conditional or require cos/sin post-multiply (:linear, :dynamic, :longrope) should use rope_with_positions/6 instead.

• a — input {B, T, ..., D} (Bumblebee convention: heads NOT yet transposed)
• position_ids — {B, T} integer tensor. For decode (T = 1) the fast path uses position_ids[b,0] as the per-batch offset into freqs (same contract as mlx::fast::rope with a scalar offset per batch). For prefill (T > 1) a per-token path runs so arbitrary positions (e.g. left-padded [0,…,0,1,2,…]) are correct; the offset-only entry point cannot represent that.
• dims — number of feature dims to rotate.
• traditional — false for split-half (Bumblebee / Qwen3); true for interleaved.
• scale — position scale (1.0 for most strategies with precomputed freqs).
• freqs — {dims/2} tensor of precomputed inverse frequencies.

Output shape and type match a.

Fused RoPE accepting a position_ids tensor (mlx::fast::rope, array-offset overload).

Use this variant when the calling convention provides position_ids as a tensor (e.g. from Bumblebee's rotary embedding layer) rather than a scalar integer offset.

• a — input {B, T, ..., D} (Bumblebee convention: heads NOT yet transposed)
• position_ids — {B, T} integer tensor; each row holds the token positions for

               one batch example. **Positions must be sequential within each row**
               (standard causal LM). The starting offset for batch item `b` is
               taken as `position_ids[b, 0]`; subsequent positions are inferred
               by MLX as `offset + 0, offset + 1, ...`.

• dims — number of feature dims to rotate.
• traditional — false for split-half (Bumblebee / Qwen3); true for interleaved.
• base — angular frequency base (e.g. 10_000).
• scale — position scale (1.0 unless using NTK-aware scaling).

Output shape and type match a.

Flash-attention SDPA, no mask (mlx::fast::scaled_dot_product_attention).

GQA-native: k/v may have fewer heads than q — no pre-tiling required.

• q — {B, N_q, T_q, D}
• k — {B, N_kv, T_kv, D}
• v — {B, N_kv, T_kv, D}
• scale — scalar (typically 1 / sqrt(D))

Output: {B, N_q, T_q, D} — same dtype as q. Softmax accumulates in float32 internally regardless of input dtype.

Flash-attention SDPA — either an additive/boolean mask tensor (5th arg), or, with no mask, an opts keyword list (disambiguated by is_list/1):

• :sinks — optional learned per-head attention-sink logits tensor, appended as an extra key/value pair the softmax normalises against (see mlx::fast::scaled_dot_product_attention's sinks parameter). Shape must broadcast against {B, N_q} (typically {N_q}).

mask must be broadcast-compatible with {B, N_q, T_q, T_kv}. Boolean false entries are masked out (-∞); float entries are added to the pre-softmax scores.

For causal masking in decode (single query token), prefer the no-mask arity since T_q=1 is always trivially causal.

Flash-attention SDPA with an additive/boolean mask and opts (:sinks — see scaled_dot_product_attention/5).

Flash-attention SDPA with a built-in causal mask (mlx::fast::scaled_dot_product_attention, mask_mode="causal").

MLX constructs the upper-triangular causal mask internally without materialising it, making this equivalent to scaled_dot_product_attention/5 with a causal boolean mask but cheaper: no mask tensor allocation, and the mask is fused into the Metal kernel.

GQA-native: k/v may have fewer heads than q — no pre-tiling required.

Input/output layout matches scaled_dot_product_attention/4:

• q — {B, N_q, T_q, D}
• k — {B, N_kv, T_kv, D}
• v — {B, N_kv, T_kv, D}
• scale — pre-computed scalar (typically 1 / sqrt(D))
• Output — {B, N_q, T_q, D}, same dtype as q

Causal flash-attention SDPA with opts (:sinks — see scaled_dot_product_attention/5).

Causal SDPA with the key_mask check delegated to the C++ NIF (eager) or folded directly into the compiled graph (traced).

When eager, the NIF evaluates all(key_mask == 1):

• true (no padding, e.g. single-sequence decode) → pure causal SDPA, no mask tensor allocated.
• false (padded batch or multi-sequence) → builds a combined causal + key_mask additive mask and calls the masked SDPA kernel.

This avoids the Nx.cond double-evaluation problem: the NIF forces eval of only the small {B, T_kv} key_mask subgraph, then branches in C++.

When traced, the compiled :fast_sdpa_causal_key_masked* opcode always builds the combined causal+key_mask additive mask in-graph (a compiled program can't branch on a runtime all(key_mask) check).

Input/output layout matches scaled_dot_product_attention_causal/4:

• q — {B, N_q, T_q, D}
• k — {B, N_kv, T_kv, D}
• v — {B, N_kv, T_kv, D}
• scale — pre-computed scalar
• key_mask — {B, T_kv} boolean/int tensor (1 = attend, 0 = masked)
• Output — {B, N_q, T_q, D}, same dtype as q

An optional 6th opts keyword list accepts :sinks (see scaled_dot_product_attention/5).

Fused SwiGLU activation: silu(gate) * up where silu(x) = x * sigmoid(x).

Eliminates the two-op silu(gate_proj) * up_proj pattern that appears in Qwen3's FFN layers (28× per decode step).

• gate — gate-projection output; silu is applied element-wise.
• up — up-projection output; same shape as gate.

Output has the same shape and dtype as gate.