Function gated_delta 

pub fn gated_delta(
    q: &Array,
    k: &Array,
    v: &Array,
    g: &Array,
    beta: &Array,
    state: &Array,
) -> (Array, Array)

The gated delta rule, Gated DeltaNet’s recurrence (as mlx-lm’s gated_delta_update once its gates are computed). For each step t and value head, with S its state ([Dv, Dk]):

S ← g_t · S, then S ← S + β_t (v_t − S k_t) k_tᵀ, and y_t = S q_t.

Shapes: q and k are [B, T, Hk, Dk], v is [B, T, Hv, Dv] (value head h reads key head h / (Hv / Hk)), g and beta are [B, T, Hv] and state is [B, Hv, Dv, Dk]. Returns y ([B, T, Hv, Dv], in q’s dtype) and the final state (in state’s dtype); the recurrence runs in float32. On the GPU it is one kernel (mlx-lm’s, for Dk a multiple of 32), else the steps as ordinary operations, which also define the derivatives.

use nuvec_core::{DType, fast, ops::{ones, zeros}};

let q = ones([1, 2, 1, 32], DType::F32);
let v = ones([1, 2, 1, 4], DType::F32);
let (g, beta) = (ones([1, 2, 1], DType::F32), ones([1, 2, 1], DType::F32));
let state = zeros([1, 1, 4, 32], DType::F32);
let (y, s) = fast::gated_delta(&(&q * 0.25), &(&q * 0.125), &v, &g, &(&beta * 0.5), &state);
// k · k = 0.5 and k · q = 1. Step 1: S = 0.5 · v kᵀ, so y = 0.5. Step 2: S k = 0.25,
// so S gains (1 − 0.25) · 0.5 · v kᵀ, and y = 0.875.
assert_eq!(y.to_vec::<f32>().unwrap(), [0.5, 0.5, 0.5, 0.5, 0.875, 0.875, 0.875, 0.875]);
assert_eq!(s.shape(), [1, 1, 4, 32]);