onnxruntime-ep-mlx 0.31.1

//! Engine core: the plan description (`NodeDesc`/`TensorRef`/`OutRef`) and the `TranslationContext`
//! a handler uses to Resolve inputs, Bind outputs, and emit MLX ops. It supports eager translation,
//! compiled decode/prefill/general closures, and recursively inlined control-flow subgraphs.

use std::collections::{HashMap, HashSet};
use std::os::raw::c_void;

use crate::mlx::{self, Array, VectorArray};
use crate::sys::mlx as mlxsys;
use crate::sys::ort;

/// Where a node input resolves from (mirrors the C++ `Src`).
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub enum Src {
    CtxInput,
    Initializer,
    Intermediate,
    Absent,
}

struct InputWriteback {
    destination: *mut c_void,
    array: mlxsys::mlx_array,
    copies: Vec<(std::ops::Range<usize>, usize)>,
}

/// A constant initializer payload surfaced at compile time (session-owned storage). For a
/// control-flow BODY initializer the bytes are COPIED into `owned` (the subgraph handle is released
/// after the compile-time walk), and `data` points into that owned buffer; for an ordinary
/// session-owned initializer `owned` is `None` and `data` points at ORT's storage.
/// String metadata has no raw payload; array resolution must not interpret it as numeric bytes.
#[derive(Clone)]
pub struct InitData {
    pub data: *const c_void,
    pub shape: Vec<i64>,
    pub dtype: ort::ONNXTensorElementDataType,
    /// ORT-reported element count, used to bound typed initializer reads safely.
    pub count: usize,
    /// Owned copy of the bytes (control-flow body initializers), keeping `data` valid.
    #[allow(dead_code)]
    // RAII: backs the `data` pointer above; must outlive it even if not read.
    pub owned: Option<std::sync::Arc<Vec<u8>>>,
}

pub(crate) fn packed4_signed(dtype: ort::ONNXTensorElementDataType) -> Option<bool> {
    match dtype {
        ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_INT4 => Some(true),
        ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT4 => Some(false),
        _ => None,
    }
}

/// Packed storage is never an ordinary MLX tensor of the logical ONNX shape.
pub(crate) struct Packed4 {
    pub data: mlxsys::mlx_array,
    pub shape: Vec<i32>,
    pub signed: bool,
}

/// A single node input reference.
#[derive(Clone)]
pub struct TensorRef {
    pub name: String,
    pub source: Src,
    pub ctx_index: usize,
    /// True when a CtxInput is a hoisted constant initializer (wrapped/cached once).
    pub constant: bool,
    /// True when this tensor's VALUE derives purely from input SHAPES + constants (a `Shape`/`Size`
    /// output, or a deterministic chain over such) — independent of runtime DATA. Such a value is a
    /// genuine constant inside a shape-keyed `mlx_compile` trace (no tracer dependency), so it can be
    /// eval'd mid-trace to feed a static reshape/expand/slice/range target.
    pub shape_const: bool,
    pub init: Option<InitData>,
}

impl TensorRef {
    pub(crate) fn cache_key(&self) -> String {
        if self.source == Src::Initializer
            && let Some(init) = &self.init
            && init.owned.is_some()
        {
            format!("{}\0{:p}", self.name, init.data)
        } else {
            self.name.clone()
        }
    }

    pub fn absent() -> Self {
        TensorRef {
            name: String::new(),
            source: Src::Absent,
            ctx_index: 0,
            constant: false,
            shape_const: false,
            init: None,
        }
    }
}

/// A single node output reference.
#[derive(Clone)]
pub struct OutRef {
    pub name: String,
    /// A subgraph boundary output routed to `KernelContext_GetOutput(ctx_index)`.
    pub external: bool,
    pub ctx_index: usize,
    pub otype: ort::ONNXTensorElementDataType,
}

/// A partial copy-out descriptor for a shared-buffer KV `present` output. In shared-buffer mode
/// `present` is bound to the SAME ORT buffer as `past`, so the rows outside `[offset, offset+count)`
/// on `axis` are already correct (they are last step's `past`); only the `count` newly written rows
/// need to be copied back. This turns the per-token copy-out from O(capacity) into O(new-tokens).
///
/// `alias_ptr` is the ORT `past` input buffer address this `present` is expected to alias. The
/// partial write is applied ONLY when the actual `present` output buffer equals `alias_ptr` — i.e.
/// the runtime really bound `present` onto `past` (share-buffer IoBinding). When they differ (e.g.
/// the op-test harness, or any non-aliasing runtime) the full copy is used, so correctness never
/// depends on an assumption about ORT's binding.
#[derive(Clone, Copy)]
pub struct DeltaWrite {
    pub axis: usize,
    pub offset: i64,
    pub count: i64,
    pub alias_ptr: usize,
}

/// A TENSOR-valued attribute payload (Constant `value`, ConstantOfShape `value`) surfaced at compile
/// time. Raw host bytes are copied into `data` (owned, kept alive for the plan's lifetime) so the
/// handler can materialize an MLX array from it. Mirrors the C++ `CopyScalarAttrs` TENSOR path.
#[derive(Clone)]
#[allow(dead_code)] // Populated for TENSOR-valued attributes; consumed by upcoming handler waves.
pub struct ConstTensor {
    pub data: Vec<u8>,
    pub shape: Vec<i64>,
    pub dtype: ort::ONNXTensorElementDataType,
    pub count: usize,
}

impl ConstTensor {
    /// Raw bytes of a single-element tensor, for one-element attributes such as
    /// `ConstantOfShape`'s `value`. Returns `None` unless the tensor holds exactly one element.
    pub fn single_element_bytes(&self) -> Option<&[u8]> {
        if self.count != 1 || self.data.is_empty() {
            return None;
        }
        Some(self.data.as_slice())
    }
}

/// A control-flow node's body subgraph (If then/else branch, Scan/Loop body), captured recursively so
/// the translator can realize the control flow by translating the body inline. `input_names` /
/// `output_names` are the body graph's FORMAL inputs/outputs (positional). `nodes` are the body's
/// nodes in topological order, whose input TensorRefs already resolve against the body scope with a
/// fall-through to the enclosing scope (implicit inputs). Faithful port of the C++ `SubgraphDesc`.
#[derive(Clone)]
pub struct SubgraphDesc {
    pub attr_name: String,
    pub input_names: Vec<String>,
    pub output_names: Vec<String>,
    pub nodes: Vec<NodeDesc>,
}

/// One ONNX node with just the metadata the MLX translator needs.
#[derive(Clone)]
pub struct NodeDesc {
    pub node_id: usize,
    pub name: String,
    pub op_type: String,
    pub domain: String,
    pub since_version: i32,
    pub compile_shape_safety: crate::registry::CompileShapeSafety,
    pub ints: HashMap<String, i64>,
    pub floats: HashMap<String, f32>,
    pub int_arrays: HashMap<String, Vec<i64>>,
    pub float_arrays: HashMap<String, Vec<f32>>,
    pub strings: HashMap<String, String>,
    pub string_arrays: HashMap<String, Vec<String>>,
    /// TENSOR-valued attributes (Constant/ConstantOfShape `value`), keyed by attribute name.
    #[allow(dead_code)] // Consumed by upcoming Constant/ConstantOfShape handler waves.
    pub tensors: HashMap<String, ConstTensor>,
    pub inputs: Vec<TensorRef>,
    pub outputs: Vec<OutRef>,
    /// Body subgraphs for a control-flow node (If/Scan/Loop). Empty for every ordinary op.
    pub subgraphs: Vec<SubgraphDesc>,
    /// Recognized DQ operands for a private quantized kernel, never a public ONNX op.
    pub(crate) qdq: Option<Box<crate::ops::qdq::QuantizedMatMul>>,
    pub(crate) lut4: Option<Box<crate::ops::lut4::LookupMatMul>>,
}

impl NodeDesc {
    pub fn new(
        op_type: String,
        domain: String,
        since_version: i32,
        compile_shape_safety: crate::registry::CompileShapeSafety,
    ) -> Self {
        NodeDesc {
            node_id: 0,
            name: String::new(),
            op_type,
            domain,
            since_version,
            compile_shape_safety,
            ints: HashMap::new(),
            floats: HashMap::new(),
            int_arrays: HashMap::new(),
            float_arrays: HashMap::new(),
            strings: HashMap::new(),
            string_arrays: HashMap::new(),
            tensors: HashMap::new(),
            inputs: Vec::new(),
            outputs: Vec::new(),
            subgraphs: Vec::new(),
            qdq: None,
            lut4: None,
        }
    }
}

/// A per-token dynamic subgraph input (non-constant `CtxInput`) fed as a compiled-closure argument.
#[derive(Clone)]
pub struct DynInput {
    pub name: String,
    pub ctx_index: usize,
}

/// A distinct RoPE cos/sin cache whose per-step row is pre-sliced OUTSIDE the compiled graph and fed
/// as a synthetic closure input (so the compiled graph never slices a cache at a runtime position —
/// which shapeless `mlx_compile` cannot shape-infer). `key` is the placeholder env key.
#[derive(Clone)]
pub struct SynthRope {
    pub key: String,
    pub cache_name: String,
}

/// Placeholder env key for the pre-sliced cos/sin row of the cache named `cache_name`.
pub fn rope_row_key(cache_name: &str) -> String {
    format!("__rope_row__{cache_name}")
}

/// Placeholder env key for the shared-buffer `valid_past` scalar. In the compiled shared-buffer
/// contract the number of valid past keys (`total_sequence_length - S`) is DATA that advances every
/// step, so it is fed as a live `[1]` int32 closure input (appended after the synth RoPE rows)
/// instead of being resolved from the in-graph `total_sequence_length` — which shapeless compile
/// would freeze at trace time. The GQA op reads it to drive the in-place write offset + causal mask.
pub const GQA_VALID_PAST_KEY: &str = "__gqa_valid_past";

/// Opaque payload handed to the `mlx_closure` trace thunk. Holds the raw pointers the thunk needs to
/// (re)build a `TranslationContext` over the plan for the ONE-TIME shapeless trace: the plan itself
/// (nodes + constant cache), the live ORT api/kernel-context (for constant host reads during the
/// trace), and the mlx stream. `kctx` is refreshed before every apply (only read during the trace).
///
/// It lives in a stable `Box` on the plan (`CompiledSubgraph::payload`) so its address — captured by
/// the base closure at build time — stays valid for the plan's lifetime.
pub struct TracePayload {
    pub plan: *mut Plan,
    pub ort_api: *const ort::OrtApi,
    pub kctx: *mut ort::OrtKernelContext,
    pub stream: mlxsys::mlx_stream,
    /// Which `CompiledSubgraph` slot on the plan this closure is tracing (so the unified core writes
    /// its discovered schema / transient handles back to the right slot).
    pub slot: Slot,
}

/// Selects one of the plan's `CompiledSubgraph` slots. Each slot carries a distinct [`CompiledConfig`]
/// and its own compiled closure + schema, but shares the ONE core in [`crate::compiled`].
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub enum Slot {
    /// Decode (`plan.compiled`): shapeless, all decode features.
    Decode,
    /// Prefill (`plan.prefill`): shape-keyed, all decode features (Phase 2).
    Prefill,
    /// General static-shape subgraph (`plan.general`): shape-keyed, no attention.
    General,
}

impl Slot {
    #[inline]
    pub fn get(self, plan: &Plan) -> &CompiledSubgraph {
        match self {
            Slot::Decode => &plan.compiled,
            Slot::Prefill => &plan.prefill,
            Slot::General => &plan.general,
        }
    }
    #[inline]
    pub fn get_mut(self, plan: &mut Plan) -> &mut CompiledSubgraph {
        match self {
            Slot::Decode => &mut plan.compiled,
            Slot::Prefill => &mut plan.prefill,
            Slot::General => &mut plan.general,
        }
    }
}

/// How `mlx_compile` keys the fused closure — the single knob that used to distinguish the two
/// historical fast-path modules.
#[derive(Clone, Copy, PartialEq, Eq)]
pub enum ShapeMode {
    /// Compile SHAPELESS: a growing/changing input shape never retraces (decode's growing KV length,
    /// which would otherwise recompile every token). One closure serves every step.
    Shapeless,
    /// Compile SHAPE-KEYED: `mlx_compile` keys on the input shapes/dtypes and transparently retraces
    /// (re-invokes the thunk) when they change, so a different shape recompiles rather than
    /// miscomputes (general static-shape subgraphs; prefill's varying query length).
    ShapeKeyed,
}

/// Static configuration selecting how a [`CompiledSubgraph`] traces, compiles, and applies. The two
/// historical fast paths — and the new prefill path — are now expressed as configurations of this
/// ONE core (the point of the unification: one place to trace → compile → cache → apply, with the
/// decode specialisations as composable opt-in features rather than a separate module):
///   * decode  = `{ Shapeless,  kv_alias, rope_as_data, delta_copyout }`
///   * general = `{ ShapeKeyed, contiguous_outputs }`
///   * prefill = `{ ShapeKeyed, kv_alias, rope_as_data, delta_copyout }` (Phase 2)
#[derive(Clone, Copy)]
pub struct CompiledConfig {
    /// Shapeless vs shape-keyed `mlx_compile` (see [`ShapeMode`]).
    pub shape_mode: ShapeMode,
    /// Support a fixed-capacity SHARED KV buffer (`present` aliased onto `past` at a runtime-owned max
    /// length): GQA writes the new K/V in place at a data-driven valid-past offset and masks the
    /// buffer tail, instead of the growing past/present concat. Detected per session at build time.
    pub kv_alias: bool,
    /// Feed the RoPE cos/sin rows (pre-sliced OUTSIDE the graph) and the `valid_past` offset as DATA
    /// closure inputs, and trace RoPE in its slice-free `[hd,hd]` matmul form (`rope_dynamic`), so the
    /// per-step position never bakes into the graph. Requires the decoder GQA shape (RoPE == head_dim).
    pub rope_as_data: bool,
    /// Copy back only the `S` new KV rows of a shared-buffer `present` (O(new-tokens) delta copy-out)
    /// rather than the whole `[B,kv,cap,hd]` buffer.
    pub delta_copyout: bool,
    /// Materialise each boundary output to row-major contiguous in the trace (so a general
    /// static-shape copy-out is a straight typed memcpy). The RoPE/KV paths leave it unset (they mirror
    /// the decode trace exactly).
    pub contiguous_outputs: bool,
}

impl CompiledConfig {
    /// Decode: shapeless (growing KV never retraces) with all decode specialisations on.
    pub const fn decode() -> Self {
        CompiledConfig {
            shape_mode: ShapeMode::Shapeless,
            kv_alias: true,
            rope_as_data: true,
            delta_copyout: true,
            contiguous_outputs: false,
        }
    }
    /// General static-shape subgraph: shape-keyed, no attention/KV specialisation, contiguous outputs.
    pub const fn general() -> Self {
        CompiledConfig {
            shape_mode: ShapeMode::ShapeKeyed,
            kv_alias: false,
            rope_as_data: false,
            delta_copyout: false,
            contiguous_outputs: true,
        }
    }
    /// Prefill (Phase 2): the SAME decoder subgraph as decode with EVERY decode specialisation
    /// (`kv_alias` / `rope_as_data` / `delta_copyout`), but a variable query length S>1. The causal
    /// mask's query-position `arange(0, S)` extent and the KV write width are read as Rust ints during
    /// the trace, so `S` bakes into the graph as a constant — a single shapeless closure traced at one
    /// S would miscompute at another. Shape-keyed compilation retraces per distinct S (re-baking it
    /// correctly) while replaying the fused closure for repeated prompt lengths.
    pub const fn prefill() -> Self {
        CompiledConfig {
            shape_mode: ShapeMode::ShapeKeyed,
            kv_alias: true,
            rope_as_data: true,
            delta_copyout: true,
            contiguous_outputs: false,
        }
    }
}

/// Unified compiled fast-path state: ONE parameterised core (`config`) plus the trace/compile/cache
/// machinery that every mode shares. What used to be the decode-only `CompiledDecode` and the
/// static-shape `GeneralCompiled` structs are now just two [`CompiledConfig`]s of this type (plus the
/// Phase 2 prefill config). The RoPE/KV fields below are only populated when the matching feature
/// flag is set; they stay empty/default for a plain general subgraph. Any doubt at build or apply
/// time flips `valid` off and the caller falls back to the always-correct eager translator.
pub struct CompiledSubgraph {
    /// Static mode + opt-in features (see [`CompiledConfig`]).
    pub config: CompiledConfig,
    /// Kill-switch off and no route-specific compile barrier.
    pub enabled: bool,
    /// Have we tried to build the compiled closure yet? (one-shot; failure => eager forever).
    pub attempted: bool,
    /// Is `closure` usable?
    pub valid: bool,
    /// Host decisions baked into the current general trace: control flow and boundary shape values.
    /// MLX's native shape-keyed cache does not observe changes to their data.
    pub host_specialization_key: Option<String>,
    /// Whether this general plan contains static control flow and therefore needs a host-decision
    /// specialization key. False for ordinary graphs, avoiding a full node walk on every call.
    pub control_flow_specialization_enabled: bool,
    /// Boundary parameters read by Reshape/Expand/Slice, checked on EVERY general apply.
    pub boundary_shape_inputs: Vec<usize>,
    /// Stable synthetic shape token per host-value specialization. The token is appended as an
    /// unused closure input so MLX's native shape-keyed cache retains every visited branch/trip
    /// combination and can replay it when execution returns to that decision.
    pub host_specializations: HashMap<String, i32>,
    /// The compiled closure.
    pub closure: Option<crate::mlx::Closure>,
    /// Ordered, de-duplicated dynamic ctx inputs = closure inputs `[0..n)`.
    pub dyn_inputs: Vec<DynInput>,
    /// External boundary outputs, in closure append order.
    pub ext_outputs: Vec<OutRef>,
    /// Transient MLX handles created during the trace; the compiled graph references them after the
    /// thunk returns, so they must outlive it (freed once with the plan).
    pub trace_transient: Vec<Array>,
    /// Stable payload for the trace thunk (self-referential to the enclosing plan).
    pub payload: Option<Box<TracePayload>>,
    // ---- rope_as_data / kv_alias state (empty/default when those features are off) ----
    /// Pre-sliced RoPE cos/sin row inputs, appended AFTER `dyn_inputs` (`rope_as_data`).
    pub synth_ropes: Vec<SynthRope>,
    /// Ctx input to read the RoPE start (past KV length) from, and its sequence axis (`rope_as_data`).
    pub rope_past_ctx_index: i32,
    pub rope_past_axis: i32,
    /// This session drives a fixed-capacity SHARED KV buffer (detected once at build time; `kv_alias`).
    pub shared_kv: bool,
    /// Ctx index of `attention_mask` — its width gives the valid-keys count used to derive the RoPE
    /// start (valid_past) at apply time in shared-buffer mode. `-1` when not shared / not found.
    pub mask_ctx_index: i32,
    /// Ctx index of GQA `total_sequence_length`; preferred over the mask width when available.
    pub total_seq_ctx_index: i32,
    /// GQA `present` KV outputs (shared-buffer mode) as `(present_output_name, past_input_ctx_index)`;
    /// their copy-out is a delta write of only the `S` new rows (`delta_copyout`).
    pub kv_present_names: Vec<(String, usize)>,
    /// Cross-attention K/V boundary inputs copied once per generation and reused for every token.
    /// Keyed by ORT KernelContext input index; cleared when a new self-attention sequence starts.
    pub stable_cross_inputs: HashMap<usize, Array>,
    /// Address of the first cross-attention K tensor for the active generation.
    pub stable_generation_key: Option<usize>,
}

impl CompiledSubgraph {
    pub fn new(config: CompiledConfig) -> Self {
        CompiledSubgraph {
            config,
            enabled: false,
            attempted: false,
            valid: false,
            host_specialization_key: None,
            control_flow_specialization_enabled: false,
            boundary_shape_inputs: Vec::new(),
            host_specializations: HashMap::new(),
            closure: None,
            dyn_inputs: Vec::new(),
            ext_outputs: Vec::new(),
            trace_transient: Vec::new(),
            payload: None,
            synth_ropes: Vec::new(),
            rope_past_ctx_index: -1,
            rope_past_axis: 2,
            shared_kv: false,
            mask_ctx_index: -1,
            total_seq_ctx_index: -1,
            kv_present_names: Vec::new(),
            stable_cross_inputs: HashMap::new(),
            stable_generation_key: None,
        }
    }
}

/// Persistent per-subgraph MLX state: the topo-ordered nodes plus the persistent cache of
/// wrapped/repacked constant arrays (keyed by name, reused across runs — freed with the plan).
pub struct Plan {
    pub nodes: Vec<NodeDesc>,
    pub cache: HashMap<String, Array>,
    /// Keeps shared immutable storage alive while per-plan handles may be used or traced.
    pub(crate) shared_initializers:
        Vec<std::sync::Arc<crate::shared_initializers::SharedInitializer>>,
    pub native_attention_decode: bool,
    pub dedicated_decode_stream: bool,
    /// This subgraph carries float64 tensors or computes float64 statistics and must run on an MLX
    /// CPU stream. GetCapability isolates CPU-required nodes from ordinary GPU work.
    pub requires_cpu_stream: bool,
    pub stable_cross_cache_enabled: bool,
    /// Boundary K/V inputs used directly by cross-attention MHA nodes.
    pub stable_cross_ctx_indices: Vec<usize>,
    /// Compiled-decode fast-path state (shapeless + decode features; see [`CompiledConfig::decode`]).
    pub compiled: CompiledSubgraph,
    /// Compiled-prefill fast-path state (shape-keyed + decode features; see [`CompiledConfig::prefill`]).
    pub prefill: CompiledSubgraph,
    /// General static-shape compiled fast-path state (see [`CompiledConfig::general`]).
    pub general: CompiledSubgraph,
}

impl Drop for Plan {
    fn drop(&mut self) {
        crate::trace::forget_partition(&self.nodes);
    }
}

pub(crate) fn is_separate_qkv_attention_op(domain: &str, op_type: &str) -> bool {
    op_type == "MultiHeadAttention" || (domain.is_empty() && op_type == "Attention")
}

pub(crate) fn is_separate_qkv_attention(node: &NodeDesc) -> bool {
    is_separate_qkv_attention_op(&node.domain, &node.op_type)
}

pub(crate) fn attention_past_key(node: &NodeDesc) -> Option<&TensorRef> {
    let index = if node.op_type == "MultiHeadAttention" {
        6
    } else if node.domain.is_empty() && node.op_type == "Attention" {
        4
    } else {
        return None;
    };
    node.inputs.get(index)
}

pub(crate) fn attention_cross_kv(node: &NodeDesc) -> Option<(&TensorRef, &TensorRef)> {
    if !is_separate_qkv_attention(node) {
        return None;
    }
    let [key, value] = node.inputs.get(1..3)? else {
        return None;
    };
    (key.source == Src::CtxInput && value.source == Src::CtxInput).then_some((key, value))
}

impl Plan {
    pub fn new(nodes: Vec<NodeDesc>) -> Self {
        let dynamic_ranges: Vec<&NodeDesc> = nodes
            .iter()
            .filter(|node| {
                node.op_type == "Range" && node.inputs.iter().any(|input| !input.constant)
            })
            .collect();
        let dynamic_tiles: Vec<&NodeDesc> = nodes
            .iter()
            .filter(|node| {
                node.op_type == "Tile" && node.inputs.get(1).is_some_and(|input| !input.constant)
            })
            .collect();
        let position_graph = match (dynamic_ranges.as_slice(), dynamic_tiles.as_slice()) {
            ([range], [tile]) => {
                let producer = |name: &str| {
                    nodes
                        .iter()
                        .find(|node| node.outputs.iter().any(|output| output.name == name))
                };
                let limit_is_start_plus_len = producer(&range.inputs[1].name).is_some_and(|add| {
                    add.op_type == "Add"
                        && add
                            .inputs
                            .iter()
                            .any(|input| input.name == range.inputs[0].name)
                });
                let tile_uses_range = producer(&tile.inputs[0].name).is_some_and(|expand| {
                    expand.op_type == "Expand"
                        && expand.inputs.first().is_some_and(|input| {
                            producer(&input.name).is_some_and(|unsqueeze| {
                                unsqueeze.op_type == "Unsqueeze"
                                    && unsqueeze
                                        .inputs
                                        .first()
                                        .is_some_and(|input| input.name == range.outputs[0].name)
                            })
                        })
                });
                let repeats_are_built =
                    producer(&tile.inputs[1].name).is_some_and(|concat| concat.op_type == "Concat");
                range.inputs[2].constant
                    && limit_is_start_plus_len
                    && tile_uses_range
                    && repeats_are_built
            }
            _ => false,
        };
        let optimized_position_graph = match (dynamic_ranges.as_slice(), dynamic_tiles.as_slice()) {
            ([range], [tile]) => {
                let producer = |name: &str| {
                    nodes
                        .iter()
                        .find(|node| node.outputs.iter().any(|output| output.name == name))
                };
                let limit_is_start_plus_len = producer(&range.inputs[1].name).is_some_and(|add| {
                    add.op_type == "Add"
                        && add
                            .inputs
                            .iter()
                            .any(|input| input.name == range.inputs[0].name)
                });
                let tile_uses_range = producer(&tile.inputs[0].name).is_some_and(|unsqueeze| {
                    unsqueeze.op_type == "Unsqueeze"
                        && unsqueeze
                            .inputs
                            .first()
                            .is_some_and(|input| input.name == range.outputs[0].name)
                });
                let repeats_are_built =
                    producer(&tile.inputs[1].name).is_some_and(|concat| concat.op_type == "Concat");
                limit_is_start_plus_len && tile_uses_range && repeats_are_built
            }
            _ => false,
        };
        let has_separate_qkv_attention = nodes.iter().any(is_separate_qkv_attention);
        let native_attention_decode = position_graph
            && nodes.len() > 32
            && has_separate_qkv_attention
            && nodes.iter().any(|node| {
                node.op_type == "MatMulNBits" && node.ints.get("bits").copied().unwrap_or(4) == 8
            });
        let dedicated_decode_stream = (native_attention_decode || optimized_position_graph)
            && nodes.len() > 32
            && has_separate_qkv_attention
            && nodes.iter().any(|node| {
                node.op_type == "MatMulNBits" && node.ints.get("bits").copied().unwrap_or(4) == 8
            });
        let mut stable_cross_ctx_indices = Vec::new();
        for (key, value) in nodes.iter().filter_map(attention_cross_kv) {
            for input in [key, value] {
                if !stable_cross_ctx_indices.contains(&input.ctx_index) {
                    stable_cross_ctx_indices.push(input.ctx_index);
                }
            }
        }
        let has_self_past_reset = nodes.iter().any(|node| {
            attention_past_key(node).is_some_and(|input| input.source == Src::CtxInput)
        });
        if !has_self_past_reset {
            stable_cross_ctx_indices.clear();
        }
        let stable_cross_cache_enabled =
            !std::env::var_os("ONNXRUNTIME_EP_MLX_NO_STABLE_CROSS_CACHE")
                .map(|value| value != "0" && !value.is_empty())
                .unwrap_or(false);
        // Best-effort fp64 detection from descriptors (outputs, initializers, and stash precision).
        // `build_plan` ORs in the authoritative claim-time answer, which also sees
        // input dtypes; this covers the plans built directly in tests and control-flow bodies.
        let requires_cpu_stream = nodes.iter().any(|node| {
            crate::registry::uses_float64_statistics(
                &node.domain,
                &node.op_type,
                node.ints.get("stash_type").copied().unwrap_or(1),
            ) || node
                .outputs
                .iter()
                .any(|o| crate::registry::is_float64(o.otype))
                || node.inputs.iter().any(|i| {
                    i.init
                        .as_ref()
                        .is_some_and(|init| crate::registry::is_float64(init.dtype))
                })
        });
        Plan {
            nodes,
            cache: HashMap::new(),
            shared_initializers: Vec::new(),
            native_attention_decode,
            dedicated_decode_stream: dedicated_decode_stream && !requires_cpu_stream,
            requires_cpu_stream,
            stable_cross_cache_enabled,
            stable_cross_ctx_indices,
            compiled: CompiledSubgraph::new(CompiledConfig::decode()),
            prefill: CompiledSubgraph::new(CompiledConfig::prefill()),
            general: CompiledSubgraph::new(CompiledConfig::general()),
        }
    }
}

pub(crate) fn initializer_array(
    input: &TensorRef,
    data: *const c_void,
    shape: &[i32],
    dtype: mlxsys::mlx_dtype,
) -> Array {
    if !crate::device::enabled() {
        return Array::from_data_managed(data, shape, dtype);
    }
    // Lazy device outputs can outlive the plan that originally owned the initializer.
    if let Some(storage) = input.init.as_ref().and_then(|init| init.owned.as_ref()) {
        Array::from_owned_data(storage.clone(), shape, dtype)
    } else {
        Array::from_data(data, shape, dtype)
    }
}

/// ONNX tensor element type -> MLX dtype (faithful port of `MlxDtypeFromOnnx`). Unknown types fall
/// back to fp32 so a stray dtype never crashes the wrap.
pub fn mlx_dtype_from_onnx(t: ort::ONNXTensorElementDataType) -> mlxsys::mlx_dtype {
    use ort::*;
    #[allow(non_upper_case_globals)]
    match t {
        ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT => {
            mlxsys::mlx_dtype__MLX_FLOAT32
        }
        ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16 => {
            mlxsys::mlx_dtype__MLX_FLOAT16
        }
        ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_BFLOAT16 => {
            mlxsys::mlx_dtype__MLX_BFLOAT16
        }
        ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE => {
            mlxsys::mlx_dtype__MLX_FLOAT64
        }
        ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_INT8 => mlxsys::mlx_dtype__MLX_INT8,
        ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_INT16 => {
            mlxsys::mlx_dtype__MLX_INT16
        }
        ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32 => {
            mlxsys::mlx_dtype__MLX_INT32
        }
        ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64 => {
            mlxsys::mlx_dtype__MLX_INT64
        }
        ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8 => {
            mlxsys::mlx_dtype__MLX_UINT8
        }
        ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT16 => {
            mlxsys::mlx_dtype__MLX_UINT16
        }
        ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT32 => {
            mlxsys::mlx_dtype__MLX_UINT32
        }
        ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT64 => {
            mlxsys::mlx_dtype__MLX_UINT64
        }
        ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_BOOL => mlxsys::mlx_dtype__MLX_BOOL,
        _ => mlxsys::mlx_dtype__MLX_FLOAT32,
    }
}

/// A translation/runtime error (mirrors the C++ `MlxError`, caught in `RunPlan`).
pub type MlxError = String;

/// Narrow an ONNX i64 dimension to MLX's 32-bit `int` shape type, erroring on
/// overflow instead of silently wrapping (`as i32`). ONNX shape/dim values are
/// i64; MLX's C ABI uses `const int*` for shapes, so a dimension that does not
/// fit i32 cannot be represented — real tensors never approach this, so this is
/// defense-in-depth against corrupt/absurd dims. Sentinels `-1` (reshape infer)
/// and `0` (copy/allowzero) fit i32 and pass through unchanged.
pub(crate) fn dim_i32(d: i64) -> Result<i32, MlxError> {
    i32::try_from(d).map_err(|_| format!("MLX: dimension {d} does not fit MLX's i32 shape range"))
}

/// Human-readable MLX dtype name for trace Args (e.g. `"float32"`, `"bfloat16"`).
pub fn dtype_name(dt: mlxsys::mlx_dtype) -> &'static str {
    #[allow(non_upper_case_globals)]
    match dt {
        mlxsys::mlx_dtype__MLX_BOOL => "bool",
        mlxsys::mlx_dtype__MLX_UINT8 => "uint8",
        mlxsys::mlx_dtype__MLX_UINT16 => "uint16",
        mlxsys::mlx_dtype__MLX_UINT32 => "uint32",
        mlxsys::mlx_dtype__MLX_UINT64 => "uint64",
        mlxsys::mlx_dtype__MLX_INT8 => "int8",
        mlxsys::mlx_dtype__MLX_INT16 => "int16",
        mlxsys::mlx_dtype__MLX_INT32 => "int32",
        mlxsys::mlx_dtype__MLX_INT64 => "int64",
        mlxsys::mlx_dtype__MLX_FLOAT16 => "float16",
        mlxsys::mlx_dtype__MLX_FLOAT32 => "float32",
        mlxsys::mlx_dtype__MLX_FLOAT64 => "float64",
        mlxsys::mlx_dtype__MLX_BFLOAT16 => "bfloat16",
        mlxsys::mlx_dtype__MLX_COMPLEX64 => "complex64",
        _ => "unknown",
    }
}

/// Raw host bytes for a constant (initializer / constant-ctx-input) tensor, surfaced at translate
/// time so shape/axes/indices operands (Reshape shape, Slice starts/ends/axes/steps, Pad pads, …)
/// can be read as plain host integers. Faithful port of the C++ `HostBytes` / `RawHost`.
pub struct HostBytes {
    pub data: *const c_void,
    #[allow(dead_code)] // Retained for symmetry/debugging; reads use `count`+`dtype`.
    pub shape: Vec<i64>,
    pub count: usize,
    pub dtype: ort::ONNXTensorElementDataType,
}

/// Per-Compute execution context: builds the MLX graph for one forward pass, evals once, copies out.
/// Handlers receive `&mut TranslationContext` and use Resolve/Bind + the MLX op helpers.
pub struct TranslationContext<'a> {
    plan: &'a mut Plan,
    ort_api: *const ort::OrtApi,
    kctx: *mut ort::OrtKernelContext,
    stream: mlxsys::mlx_stream,
    /// name -> raw handle (borrowed; owned by `arena` or the plan cache).
    env: HashMap<String, mlxsys::mlx_array>,
    /// Arrays produced during translation. Boundary roots retain their own references before
    /// these handles are released, allowing MLX to recycle intermediates during eager eval.
    arena: Vec<Array>,
    /// Runtime input wrappers stay referenced through eval, preventing donation of caller-owned
    /// zero-copy buffers when intermediate handles are released.
    input_arena: Vec<Array>,
    /// Per-translation cast CSE. This is intentionally scoped to one graph build: runtime inputs
    /// must never be retained across Compute calls, while repeated Q/K/V and gate/up projections
    /// should share one explicit activation cast inside the compiled closure.
    astype_cache: HashMap<(usize, u32), mlxsys::mlx_array>,
    /// Fused-handler intermediates shared only within this graph translation.
    temporary_arrays: HashMap<String, mlxsys::mlx_array>,
    shape_stream: Option<mlx::Stream>,
    cpu_shape_ops: bool,
    /// Immutable Cast/Transpose/Identity provenance; runtime data never enters this map.
    frozen_views: HashMap<String, String>,
    /// Path a handler declared for the CURRENT node (fast vs composed); the dispatcher
    /// resets this before each handler and reads it after (see `registry::translate`).
    path_mark: Option<crate::trace::PathMark>,
    /// Cached tracing-enabled gate so `mark_fast`/`mark_composed` are a cheap no-op (and
    /// never allocate a reason string) when tracing is off.
    trace_enabled: bool,
    /// True while translating INSIDE the compiled-decode closure trace. In this mode RoPE uses the
    /// pre-sliced cos/sin ROW placeholders (fed as extra closure inputs) + a matmul rotate-half, so
    /// the graph carries no dynamic Slice (which shapeless `mlx_compile` cannot shape-infer).
    rope_dynamic: bool,
    /// True when the compiled-decode trace is for a fixed-capacity SHARED KV buffer session (GQA
    /// writes the new K/V in place at the valid-past offset and masks the buffer tail). Mirrors
    /// `plan.compiled.shared_kv`; only meaningful while `rope_dynamic` is set.
    shared_kv: bool,
    /// True when the compiled trace is SHAPE-KEYED (prefill) rather than shapeless (decode). In this
    /// mode the query length `S` and `valid_past` are fixed for the shape key, so the shared-buffer
    /// attention can be STATICALLY narrowed to the valid prefix `[0, valid_past+S)` (a static slice +
    /// causal SDPA) instead of masking over the full KV capacity. Only meaningful with `rope_dynamic`.
    compiled_shape_keyed: bool,
    /// The static `valid_past` (number of already-valid KV rows) for a shape-keyed compiled trace —
    /// known per shape key because the `attention_mask` width (= `valid_past + S`) is part of the key.
    /// Only meaningful when `compiled_shape_keyed` is set.
    compiled_valid_past: i32,
    /// Shared-buffer KV `present` outputs (name -> delta) for the EAGER path: only the `count` new
    /// rows at `offset` need copying back (the rest already alias correct `past` rows). Empty on the
    /// growing path so its copy-out stays a full, bit-for-bit-unchanged memcpy.
    kv_deltas: HashMap<String, DeltaWrite>,
    /// Evaluated MLX arrays that must be copied back into mutable ORT input buffers for
    /// side-effecting ops such as PagedAttention's in-place KV-cache update.
    input_writebacks: Vec<InputWriteback>,
    patched_alias_outputs: HashSet<String>,
    borrowed_runtime_inputs: bool,
    /// True while translating inside the GENERAL compiled-subgraph closure trace. In this mode the
    /// dynamic inputs are shapeless tracer placeholders with no host data, so any mid-graph host eval
    /// (`contiguous_eval` — the host-computed Det/NonZero/Unique ops) is illegal and must fail the
    /// trace cleanly so the caller falls back to the eager translator instead of eval-ing a tracer.
    in_general_trace: bool,
    /// Shared-buffer KV `present` outputs discovered during a COMPILED trace (`rope_dynamic`), as
    /// `(present_output_name, past_input_ctx_index)`. Collected here (not written straight to a plan
    /// slot) so the unified core can hand them to whichever `CompiledSubgraph` slot it is tracing.
    compiled_kv_present: Vec<(String, usize)>,
    /// True while translating inside ANY SHAPE-KEYED compiled subgraph trace (`CompiledConfig::general`
    /// or `CompiledConfig::prefill` — both `ShapeMode::ShapeKeyed`), as opposed to a SHAPELESS decode
    /// trace or plain eager (non-compiled) translation. Set once, unconditionally, from `cfg.shape_mode`
    /// right after `TranslationContext::new` — independent of `compiled_shape_keyed` (which additionally
    /// requires the RoPE-as-data shared-buffer "narrow" detection to have succeeded). Multi-token
    /// matrix-matrix compiled prefill work (M>1) always shows up here; decode (M==1, Shapeless) never
    /// does. Used to gate opt-in fast kernels that must never run during decode.
    shape_keyed_compile: bool,
    /// True only while translating any compiled closure; eager translation leaves this false.
    in_compiled_trace: bool,
    /// Shape templates for a zero-trip Scan are not real body iterations.
    empty_scan_template: bool,
}

impl<'a> TranslationContext<'a> {
    pub fn new(
        plan: &'a mut Plan,
        ort_api: *const ort::OrtApi,
        kctx: *mut ort::OrtKernelContext,
        stream: mlxsys::mlx_stream,
    ) -> Self {
        TranslationContext {
            plan,
            ort_api,
            kctx,
            stream,
            env: HashMap::new(),
            arena: Vec::new(),
            input_arena: Vec::new(),
            astype_cache: HashMap::new(),
            temporary_arrays: HashMap::new(),
            shape_stream: None,
            cpu_shape_ops: std::env::var_os("ONNXRUNTIME_EP_MLX_GPU_SHAPE_OPS")
                .is_none_or(|v| v != "1"),
            frozen_views: HashMap::new(),
            path_mark: None,
            trace_enabled: crate::trace::tracer().is_enabled(),
            rope_dynamic: false,
            shared_kv: false,
            compiled_shape_keyed: false,
            compiled_valid_past: 0,
            kv_deltas: HashMap::new(),
            input_writebacks: Vec::new(),
            patched_alias_outputs: HashSet::new(),
            borrowed_runtime_inputs: false,
            in_general_trace: false,
            compiled_kv_present: Vec::new(),
            shape_keyed_compile: false,
            in_compiled_trace: false,
            empty_scan_template: false,
        }
    }

    fn with_shape_stream<T>(
        &mut self,
        operation: impl FnOnce(&mut Self) -> Result<T, MlxError>,
    ) -> Result<T, MlxError> {
        let saved = self.stream;
        self.stream = self
            .shape_stream
            .get_or_insert_with(mlx::Stream::new_default_cpu)
            .as_raw();
        let result = operation(self);
        self.stream = saved;
        result
    }

    fn host_shape_ops_enabled(&self) -> bool {
        !self.in_compiled_trace
            && !self.in_general_trace
            && !self.rope_dynamic
            && self.cpu_shape_ops
    }

    pub(crate) fn translate_shape_aware(
        &mut self,
        node: &NodeDesc,
        handler: crate::registry::OpHandler,
    ) -> Result<(), MlxError> {
        // Shape-derived integers stay on the host until a GPU consumer needs them.
        // Live payloads and compiled traces retain their existing execution streams.
        let candidate = self.host_shape_ops_enabled()
            && matches!(node.domain.as_str(), "" | "ai.onnx")
            && matches!(
                node.op_type.as_str(),
                "Equal"
                    | "Greater"
                    | "Less"
                    | "GreaterOrEqual"
                    | "LessOrEqual"
                    | "Not"
                    | "And"
                    | "Or"
                    | "Squeeze"
                    | "Unsqueeze"
                    | "Reshape"
                    | "Concat"
                    | "Gather"
                    | "Slice"
                    | "Add"
                    | "Sub"
                    | "Mul"
                    | "Div"
                    | "Cast"
                    | "Identity"
                    | "NonZero"
            )
            && node.inputs.iter().any(|i| i.shape_const)
            && node.inputs.iter().all(|i| {
                i.shape_const || i.constant || matches!(i.source, Src::Absent | Src::Initializer)
            })
            && node.outputs.iter().all(|o| {
                matches!(
                    o.otype,
                    ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64
                        | ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32
                        | ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_BOOL
                )
            });
        if candidate {
            let mut small = true;
            for input in &node.inputs {
                if input.source == Src::Absent {
                    continue;
                }
                let a = self.resolve(input)?;
                small &= self.size_of(a) <= 64
                    && matches!(
                        self.dtype_of(a),
                        mlxsys::mlx_dtype__MLX_INT64
                            | mlxsys::mlx_dtype__MLX_INT32
                            | mlxsys::mlx_dtype__MLX_BOOL
                    );
            }
            if small {
                let result = self.with_shape_stream(|ctx| handler(ctx, node));
                if result.is_ok() {
                    self.mark_fast("mlx_cpu_shape");
                }
                return result;
            }
        }
        handler(self, node)
    }

    /// Declare that the current node took its **fused MLX kernel** fast path (green/normal).
    /// `kernel` is the MLX kernel name (e.g. `"mlx_fast_sdpa"`). Cheap no-op when tracing is off.
    #[inline]
    pub fn mark_fast(&mut self, kernel: &'static str) {
        if self.trace_enabled {
            self.path_mark = Some(crate::trace::PathMark::Fast(kernel));
        }
    }

    /// Declare that the current node fell back to a slower **composed/generic** path even though
    /// a fused MLX kernel exists — this is flagged prominently in the trace. `reason` explains why
    /// (e.g. `"block_size 16 → dequant+dense matmul"`). Cheap no-op (no string alloc) when tracing
    /// is off.
    #[inline]
    pub fn mark_composed(&mut self, reason: impl Into<String>) {
        if self.trace_enabled {
            self.path_mark = Some(crate::trace::PathMark::Composed(reason.into()));
        }
    }

    /// Clear the per-node path slot before dispatching a handler.
    #[inline]
    pub fn reset_path_mark(&mut self) {
        self.path_mark = None;
    }

    /// Take the path a handler declared for the node just dispatched (see `mark_fast`/`mark_composed`).
    #[inline]
    pub fn take_path_mark(&mut self) -> Option<crate::trace::PathMark> {
        self.path_mark.take()
    }

    #[inline]
    #[allow(dead_code)]
    pub fn stream(&self) -> mlxsys::mlx_stream {
        self.stream
    }

    pub(crate) fn empty_scan_template(&self) -> bool {
        self.empty_scan_template
    }

    /// Register a freshly produced array for teardown at end of run; returns its raw handle.
    pub fn keep(&mut self, a: Array) -> mlxsys::mlx_array {
        let raw = a.as_raw();
        self.arena.push(a);
        raw
    }

    fn keep_input(&mut self, a: Array) -> mlxsys::mlx_array {
        let raw = a.as_raw();
        self.input_arena.push(a);
        raw
    }

    /// Look up a persistent (plan-cached) array by key — the borrowed raw handle if present. Used by
    /// weight-repack handlers (MatMulNBits) to reuse a once-built constant across runs.
    pub fn cache_get(&self, key: &str) -> Option<mlxsys::mlx_array> {
        self.plan.cache.get(key).map(|a| a.as_raw())
    }

    /// Insert an owning array into the persistent plan cache under `key` (freed with the plan) and
    /// return its borrowed raw handle. Use only for genuinely constant (initializer) data.
    pub fn cache_put(&mut self, key: String, a: Array) -> mlxsys::mlx_array {
        let raw = a.as_raw();
        self.plan.cache.insert(key, a);
        raw
    }

    fn runtime_input_array(&mut self, data: *const c_void, shape: &[i32], dtype: u32) -> Array {
        let count = shape.iter().try_fold(1usize, |count, &dim| {
            count.checked_mul(usize::try_from(dim).ok()?)
        });
        let small = count
            .and_then(|count| {
                onnx_runtime_ir::DataType::from_onnx(dtype as i32)
                    .and_then(|dtype| dtype.checked_storage_bytes(count))
            })
            .is_some_and(|bytes| bytes <= 4096);
        if small
            && std::env::var_os("ONNXRUNTIME_EP_MLX_OWN_SMALL_INPUTS")
                .is_none_or(|value| value != "0" && !value.is_empty())
        {
            Array::from_data(data, shape, mlx_dtype_from_onnx(dtype))
        } else {
            self.borrowed_runtime_inputs = true;
            Array::from_data_managed(data, shape, mlx_dtype_from_onnx(dtype))
        }
    }

    pub(crate) fn initializer_view(
        &mut self,
        input: &TensorRef,
        data: *const c_void,
        shape: &[i32],
        dtype: mlxsys::mlx_dtype,
    ) -> Result<Array, MlxError> {
        if input.source == Src::Initializer
            && let Some(init) = &input.init
            && let Some(shared) = crate::shared_initializers::get_view(init, shape, dtype)?
        {
            let handle = shared.array.retained()?;
            self.plan.shared_initializers.push(shared);
            return Ok(handle);
        }
        Ok(initializer_array(input, data, shape, dtype))
    }

    pub(crate) fn derived_initializer_array(
        &mut self,
        input: &TensorRef,
        data: &[u8],
        shape: &[i32],
        dtype: u32,
    ) -> Result<Array, MlxError> {
        if input.source == Src::Initializer
            && let Some(shared) = crate::shared_initializers::get_derived(data, shape, dtype)?
        {
            let handle = shared.array.retained()?;
            self.plan.shared_initializers.push(shared);
            return Ok(handle);
        }
        Ok(Array::from_data(
            data.as_ptr().cast(),
            shape,
            mlx_dtype_from_onnx(dtype),
        ))
    }

    /// Borrow an arena-owned intermediate from this graph translation.
    pub fn temporary_get(&self, key: &str) -> Option<mlxsys::mlx_array> {
        self.temporary_arrays.get(key).copied()
    }

    /// Share an already arena-owned intermediate until the evaluation boundary.
    pub fn temporary_put(&mut self, key: String, value: mlxsys::mlx_array) {
        self.temporary_arrays.insert(key, value);
    }

    /// Bind a node output name to a produced MLX array (visible to downstream nodes and CopyOut).
    pub fn bind(&mut self, o: &OutRef, a: mlxsys::mlx_array) {
        self.env.insert(o.name.clone(), a);
        self.frozen_views.remove(&o.name);
    }

    /// Prove that a value comes from immutable initializer data, excluding shape-derived data.
    pub fn frozen_key(&self, input: &TensorRef) -> Option<String> {
        if input.source == Src::Initializer || (input.source == Src::CtxInput && input.constant) {
            Some(format!("leaf:{:?}", input.cache_key()))
        } else if input.source == Src::Intermediate {
            self.frozen_views.get(&input.name).cloned()
        } else {
            None
        }
    }

    /// Propagate only unary views with immutable data and fixed operator attributes.
    pub fn track_frozen_view(&mut self, node: &NodeDesc) {
        if !matches!(node.domain.as_str(), "" | "ai.onnx")
            || node.inputs.len() != 1
            || node.outputs.len() != 1
        {
            return;
        }
        if !matches!(node.op_type.as_str(), "Identity" | "Cast" | "Transpose") {
            return;
        }
        let Some(origin) = self.frozen_key(&node.inputs[0]) else {
            return;
        };
        let key = match node.op_type.as_str() {
            "Identity" => origin,
            "Cast" => format!("cast({origin},{})", node.outputs[0].otype),
            "Transpose" => format!("transpose({origin},{:?})", node.int_arrays.get("perm")),
            _ => unreachable!(),
        };
        self.frozen_views.insert(node.outputs[0].name.clone(), key);
    }

    /// Retain one contiguous layout of proven immutable data for repeated executions.
    pub fn cache_contiguous(
        &mut self,
        key: String,
        a: mlxsys::mlx_array,
    ) -> Result<mlxsys::mlx_array, MlxError> {
        if let Some(cached) = self.cache_get(&key) {
            return Ok(cached);
        }
        let mut raw = unsafe { mlxsys::mlx_array_new() };
        let rc = unsafe { mlxsys::mlx_contiguous(&mut raw, a, false, self.stream) };
        let packed = Array::from_raw(raw);
        if rc != 0 {
            return Err("MLX frozen layout packing failed".into());
        }
        packed.eval();
        Ok(self.cache_put(key, packed))
    }

    pub fn write_back_input_rows(
        &mut self,
        input: &TensorRef,
        a: mlxsys::mlx_array,
        axis: usize,
        rows: &[usize],
        output: Option<&OutRef>,
    ) -> Result<mlxsys::mlx_array, MlxError> {
        let a = self.contiguous(a)?;
        self.env.insert(input.name.clone(), a);
        self.frozen_views.remove(&input.name);
        if input.source == Src::CtxInput {
            let value = read_ctx_input_value(self.ort_api, self.kctx, input.ctx_index)?;
            let array = std::mem::ManuallyDrop::new(Array::from_raw(a));
            if crate::device::assign(unsafe { &*self.ort_api }, value, &array)? {
                return Ok(a);
            }
        }
        match input.source {
            Src::CtxInput => {
                let (data, shape, _) = self.read_ctx_input(input.ctx_index)?;
                let arr = std::mem::ManuallyDrop::new(Array::from_raw(a));
                let ranges = input_writeback_ranges(&shape, arr.itemsize(), axis, rows)?;
                if arr.shape() != shape {
                    return Err(format!(
                        "MLX: input writeback {} shape {:?} differs from input {shape:?}; preserve input geometry for in-place updates",
                        input.name,
                        arr.shape()
                    ));
                }
                if let Some(output) = output.filter(|out| !out.name.is_empty())
                    && !rows.is_empty()
                {
                    let axis_len = shape[axis] as usize;
                    let first = rows.iter().map(|row| row % axis_len).min().unwrap();
                    let last = rows.iter().map(|row| row % axis_len).max().unwrap();
                    // Copy the enclosing row interval in every slab only when ORT actually
                    // aliases this output to the input. All other rows are unchanged.
                    self.kv_deltas.insert(
                        output.name.clone(),
                        DeltaWrite {
                            axis,
                            offset: first as i64,
                            count: (last - first + 1) as i64,
                            alias_ptr: data as usize,
                        },
                    );
                }
                self.input_writebacks.push(InputWriteback {
                    destination: data as *mut c_void,
                    array: a,
                    copies: ranges.into_iter().map(|r| (r.clone(), r.start)).collect(),
                });
            }
            Src::Intermediate => {}
            _ => return Err("MLX: cannot write back an immutable input".to_string()),
        }
        Ok(a)
    }

    /// Publish small packed cache updates without rooting a full functional scatter.
    pub fn write_back_input_patches(
        &mut self,
        input: &TensorRef,
        updated: mlxsys::mlx_array,
        patches: mlxsys::mlx_array,
        rows: &[usize],
        output: Option<&OutRef>,
    ) -> Result<mlxsys::mlx_array, MlxError> {
        self.env.insert(input.name.clone(), updated);
        self.frozen_views.remove(&input.name);
        if input.source == Src::Intermediate {
            return Ok(updated);
        }
        if input.source != Src::CtxInput {
            return Err("MLX: cache patches require a mutable runtime input".into());
        }
        let value = read_ctx_input_value(self.ort_api, self.kctx, input.ctx_index)?;
        let array = std::mem::ManuallyDrop::new(Array::from_raw(updated));
        if crate::device::assign(unsafe { &*self.ort_api }, value, &array)? {
            return Ok(updated);
        }
        let (data, shape, dtype) = self.read_ctx_input(input.ctx_index)?;
        let patches = self.contiguous(patches)?;
        let array = std::mem::ManuallyDrop::new(Array::from_raw(patches));
        if shape.len() != 4 || array.dtype() != mlx_dtype_from_onnx(dtype) {
            return Err(
                "MLX: cache patches must preserve the rank-4 input dtype and geometry".into(),
            );
        }
        let row_bytes = (shape[2] as usize)
            .checked_mul(shape[3] as usize)
            .and_then(|v| v.checked_mul(array.itemsize()))
            .ok_or("MLX: cache patch row size exceeds address space")?;
        if rows.len().checked_mul(row_bytes) != Some(array.size() * array.itemsize()) {
            return Err("MLX: cache patch storage must contain one complete row per slot".into());
        }
        let slots = (shape[0] as usize)
            .checked_mul(shape[1] as usize)
            .ok_or("MLX: cache slot count exceeds address space")?;
        if rows.iter().any(|&r| r >= slots) {
            return Err("MLX: cache patch slot exceeds the input cache capacity".into());
        }
        if let Some(output) = output.filter(|o| !o.name.is_empty()) {
            let api = unsafe { &*self.ort_api };
            let mut value = std::ptr::null_mut();
            ort_status(
                api,
                unsafe {
                    (api.KernelContext_GetOutput.unwrap())(
                        self.kctx,
                        output.ctx_index,
                        shape.as_ptr(),
                        shape.len(),
                        &mut value,
                    )
                },
                "KernelContext_GetOutput for cache patches",
            )?;
            if value.is_null() {
                return Err(
                    "MLX: ORT returned no requested cache output; check output bindings".into(),
                );
            }
            let (destination, output_shape, output_dtype) = read_ort_tensor_raw(api, value)?;
            if output_shape != shape || output_dtype != dtype {
                return Err(
                    "MLX: cache output binding must preserve the input shape and dtype".into(),
                );
            }
            if destination == data {
                self.patched_alias_outputs.insert(output.name.clone());
            }
        }
        let copies = rows
            .iter()
            .enumerate()
            .map(|(i, &row)| (i * row_bytes..(i + 1) * row_bytes, row * row_bytes))
            .collect();
        self.input_writebacks.push(InputWriteback {
            destination: data as *mut c_void,
            array: patches,
            copies,
        });
        Ok(updated)
    }

    /// Emit the per-op trace detail for the node just translated (only when tracing is on).
    ///
    /// Reads each output's shape/dtype/size from the (lazy) bound arrays — shape and dtype
    /// are graph metadata available WITHOUT eval, so the normal (fine-off) path stays fully
    /// fused. When [`fine_enabled`](crate::trace::MlxTracer::fine_enabled) is set, additionally
    /// forces an `mlx_array_eval` of this node's outputs to time it individually — a
    /// GPU-inclusive per-op bar that BREAKS fusion (debug-only, slower).
    pub fn trace_node(
        &mut self,
        _op_type: &str,
        node: &NodeDesc,
        start: Option<std::time::Instant>,
    ) {
        if !self.trace_enabled {
            return;
        }
        let Some(start) = start else {
            return;
        };
        let tr = crate::trace::tracer();

        // Output metadata (shapes / dtype / elements / bytes) — from whatever is materialized.
        let mut out_shapes: Vec<String> = Vec::new();
        let mut dtype = "";
        let mut elements: u64 = 0;
        let mut bytes: u64 = 0;
        for o in &node.outputs {
            if o.name.is_empty() {
                continue;
            }
            if let Some(&raw) = self.env.get(&o.name) {
                // Borrow the handle for metadata; do NOT free it (owned by arena/cache).
                let arr = std::mem::ManuallyDrop::new(Array::from_raw(raw));
                let sh = arr.shape();
                let cnt: u64 = sh.iter().map(|&d| d.max(0) as u64).product();
                out_shapes.push(format!("{sh:?}"));
                dtype = dtype_name(arr.dtype());
                elements += cnt;
                bytes += cnt * arr.itemsize() as u64;
            }
        }
        // Input shapes, best-effort from whatever is already materialized in the env.
        let mut in_shapes: Vec<String> = Vec::new();
        for inp in &node.inputs {
            if inp.name.is_empty() {
                continue;
            }
            if let Some(&raw) = self.env.get(&inp.name) {
                let arr = std::mem::ManuallyDrop::new(Array::from_raw(raw));
                in_shapes.push(format!("{:?}", arr.shape()));
            }
        }
        let out_s = out_shapes.join(";");
        let in_s = in_shapes.join(";");

        // A build-time span with resource metadata. The fused subgraph runs as a single
        // `mlx.eval` (see finish_boundary); per-kernel GPU detail is the Xcode GPU capture.
        tr.record_op_meta(
            node,
            start,
            start.elapsed(),
            &out_s,
            &in_s,
            dtype,
            elements,
            bytes,
        );
    }

    /// Resolve a node input to a raw MLX array handle (intermediate env / wrapped ctx input /
    /// cached-or-wrapped initializer). Faithful port of `TranslationContext::Resolve`.
    pub fn resolve(&mut self, r: &TensorRef) -> Result<mlxsys::mlx_array, MlxError> {
        match r.source {
            Src::Intermediate => self
                .env
                .get(&r.name)
                .copied()
                .ok_or_else(|| format!("MLX: missing intermediate {}", r.name)),
            Src::CtxInput => {
                if r.constant {
                    if let Some(a) = self.plan.cache.get(&r.name) {
                        return Ok(a.as_raw());
                    }
                } else if let Some(a) = self.env.get(&r.name) {
                    return Ok(*a);
                }
                let value = read_ctx_input_value(self.ort_api, self.kctx, r.ctx_index)?;
                if let Some(array) = crate::device::array(unsafe { &*self.ort_api }, value)? {
                    if r.constant {
                        let raw = array.as_raw();
                        self.plan.cache.insert(r.name.clone(), array);
                        return Ok(raw);
                    }
                    let raw = self.keep_input(array);
                    self.env.insert(r.name.clone(), raw);
                    return Ok(raw);
                }
                let (data, shape, dtype) = self.read_ctx_input(r.ctx_index)?;
                if !crate::registry::is_movable(dtype) {
                    return Err(format!(
                        "MLX: tensor payload dtype {dtype} has no MLX representation"
                    ));
                }
                let ishape: Vec<i32> = shape
                    .iter()
                    .map(|&d| dim_i32(d))
                    .collect::<Result<_, _>>()?;
                if r.constant {
                    // Runtime ctx-input buffers are valid only for this Compute call. Keep the
                    // defensive copy for compatibility with plans built by older ORT behavior.
                    let arr = Array::from_data(data, &ishape, mlx_dtype_from_onnx(dtype));
                    let raw = arr.as_raw();
                    self.plan.cache.insert(r.name.clone(), arr);
                    Ok(raw)
                } else {
                    // Owned small inputs can outlive Compute; borrowed large inputs still require
                    // synchronous evaluation before ORT may reuse their zero-copy buffers.
                    let arr = self.runtime_input_array(data, &ishape, dtype);
                    let raw = self.keep_input(arr);
                    self.env.insert(r.name.clone(), raw);
                    Ok(raw)
                }
            }
            Src::Initializer => {
                let init = r
                    .init
                    .as_ref()
                    .ok_or_else(|| format!("MLX: initializer {} has no data", r.name))?;
                if !crate::registry::is_movable(init.dtype) {
                    return Err(format!(
                        "MLX: initializer payload dtype {} has no MLX representation",
                        init.dtype,
                    ));
                }
                // Body-local initializers may legally shadow an outer/other-branch name. Qualify
                // copied body constants by their stable owned allocation so branch specialization
                // cannot reuse a same-named constant from another scope.
                let cache_key = r.cache_key();
                if let Some(a) = self.plan.cache.get(&cache_key) {
                    return Ok(a.as_raw());
                }
                let ishape: Vec<i32> = init
                    .shape
                    .iter()
                    .map(|&d| dim_i32(d))
                    .collect::<Result<_, _>>()?;
                if let Some(shared) = crate::shared_initializers::get(init, &ishape)? {
                    let handle = shared.array.retained()?;
                    let raw = handle.as_raw();
                    self.plan.shared_initializers.push(shared);
                    self.plan.cache.insert(cache_key, handle);
                    return Ok(raw);
                }
                let arr = initializer_array(r, init.data, &ishape, mlx_dtype_from_onnx(init.dtype));
                let raw = arr.as_raw();
                self.plan.cache.insert(cache_key, arr);
                Ok(raw)
            }
            Src::Absent => Err("MLX: absent input".to_string()),
        }
    }

    pub(crate) fn packed4_host(&self, r: &TensorRef) -> Result<Option<HostBytes>, MlxError> {
        if !matches!(r.source, Src::Initializer | Src::CtxInput) {
            return Ok(None);
        }
        if self.in_general_trace || self.in_compiled_trace {
            if r.init
                .as_ref()
                .is_some_and(|init| packed4_signed(init.dtype).is_some())
            {
                return Err("MLX: packed tensor transport cannot be captured".into());
            }
            return Ok(None);
        }
        let host = self.raw_host(r)?;
        Ok(packed4_signed(host.dtype).map(|_| host))
    }

    pub(crate) fn resolve_packed4(&mut self, r: &TensorRef) -> Result<Option<Packed4>, MlxError> {
        let Some(host) = self.packed4_host(r)? else {
            return Ok(None);
        };
        let Some(signed) = packed4_signed(host.dtype) else {
            return Err("MLX: packed tensor metadata lost its logical type".into());
        };
        let shape = host
            .shape
            .iter()
            .map(|&d| dim_i32(d))
            .collect::<Result<Vec<_>, _>>()?;
        let count = if shape.contains(&0) {
            0
        } else {
            shape.iter().try_fold(1usize, |n, &d| {
                n.checked_mul(d as usize)
                    .ok_or("MLX: packed tensor element count overflows")
            })?
        };
        if count != host.count || count > (i32::MAX - 1) as usize {
            return Err(
                "MLX: packed tensor logical count is inconsistent or exceeds int32 limits".into(),
            );
        }
        let bytes = onnx_runtime_ir::DataType::from_onnx(host.dtype as i32)
            .and_then(|dtype| dtype.checked_storage_bytes(count))
            .ok_or("MLX: packed tensor byte count overflows")?;
        if r.source == Src::CtxInput {
            let value = read_ctx_input_value(self.ort_api, self.kctx, r.ctx_index)?;
            let api = unsafe { &*self.ort_api };
            let mut actual_bytes = 0;
            let size_in_bytes = api
                .GetTensorSizeInBytes
                .ok_or("MLX: packed tensors require ORT GetTensorSizeInBytes")?;
            ort_status(
                api,
                unsafe { size_in_bytes(value, &mut actual_bytes) },
                "GetTensorSizeInBytes",
            )?;
            if actual_bytes != bytes {
                return Err("MLX: packed ORT storage size does not match logical shape".into());
            }
        }
        if let Some(owned) = r.init.as_ref().and_then(|init| init.owned.as_ref())
            && owned.len() != bytes
        {
            return Err("MLX: packed initializer byte count does not match logical shape".into());
        }
        if bytes != 0 && host.data.is_null() {
            return Err("MLX: nonempty packed tensor has no payload".into());
        }
        let storage_shape = [bytes as i32];
        let cache_key = format!("{}\0packed4", r.cache_key());
        let data = if let Some(array) = self.plan.cache.get(&cache_key) {
            array.as_raw()
        } else if bytes == 0 {
            self.zeros(&storage_shape, mlxsys::mlx_dtype__MLX_UINT8)?
        } else if r.source == Src::Initializer || r.constant {
            let array = Array::from_data(host.data, &storage_shape, mlxsys::mlx_dtype__MLX_UINT8);
            let raw = array.as_raw();
            self.plan.cache.insert(cache_key, array);
            raw
        } else {
            let array = self.runtime_input_array(
                host.data,
                &storage_shape,
                ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8,
            );
            self.keep_input(array)
        };
        Ok(Some(Packed4 {
            data,
            shape,
            signed,
        }))
    }

    /// Read a constant/parameter input's HOST bytes at translate time (shape/axes/indices operands).
    /// Handles both a compile-time `Initializer` and a constant/dynamic `CtxInput` (read live from
    /// the kernel context each run). Faithful port of `TranslationContext::RawHost`.
    pub fn raw_host(&self, r: &TensorRef) -> Result<HostBytes, MlxError> {
        match r.source {
            Src::Initializer => {
                let init = r
                    .init
                    .as_ref()
                    .ok_or_else(|| format!("MLX: initializer {} has no data", r.name))?;
                Ok(HostBytes {
                    data: init.data,
                    shape: init.shape.clone(),
                    count: init.count,
                    dtype: init.dtype,
                })
            }

            Src::CtxInput => {
                let (data, shape, dtype) = self.read_ctx_input(r.ctx_index)?;
                let count = if shape.contains(&0) {
                    0
                } else {
                    shape.iter().try_fold(1usize, |count, &dim| {
                        usize::try_from(dim)
                            .ok()
                            .and_then(|dim| count.checked_mul(dim))
                            .ok_or_else(|| "MLX: invalid host tensor element count".to_string())
                    })?
                };
                Ok(HostBytes {
                    data,
                    shape,
                    count,
                    dtype,
                })
            }
            _ => Err(format!("MLX: RawHost on non-constant input {}", r.name)),
        }
    }

    pub fn tensor_shape(&mut self, r: &TensorRef) -> Result<Vec<i64>, MlxError> {
        match r.source {
            Src::CtxInput => {
                if let Some(&array) = self.env.get(&r.name) {
                    return Ok(self.shape_of(array).into_iter().map(i64::from).collect());
                }
                if self.in_general_trace || self.in_compiled_trace {
                    return Err(format!(
                        "MLX: missing traced tensor metadata input {}",
                        r.name
                    ));
                }
                let value = read_ctx_input_value(self.ort_api, self.kctx, r.ctx_index)?;
                let (shape, _) = read_ort_tensor_metadata(unsafe { &*self.ort_api }, value)?;
                Ok(shape)
            }
            Src::Initializer => r
                .init
                .as_ref()
                .map(|init| init.shape.clone())
                .ok_or_else(|| format!("MLX: initializer {} has no metadata", r.name)),
            Src::Intermediate => {
                let array = self.resolve(r)?;
                Ok(self.shape_of(array).into_iter().map(i64::from).collect())
            }
            Src::Absent => Err("MLX: absent tensor metadata input".into()),
        }
    }

    /// Whether an optional tensor reaches this fused node's input boundary.
    ///
    /// ORT 1.29 may return a non-null placeholder for an empty optional, so pointer-nullness alone
    /// is insufficient. Only a present optional(tensor) can yield tensor shape information.
    pub fn ctx_input_is_present(&self, index: usize) -> Result<bool, MlxError> {
        unsafe {
            let api = &*self.ort_api;
            let mut value: *const ort::OrtValue = std::ptr::null();
            let st = (api.KernelContext_GetInput.unwrap())(self.kctx, index, &mut value);
            if !st.is_null() {
                (api.ReleaseStatus.unwrap())(st);
                return Err(format!("MLX: KernelContext_GetInput({index}) failed"));
            }
            if value.is_null() {
                return Ok(false);
            }
            let mut info: *mut ort::OrtTensorTypeAndShapeInfo = std::ptr::null_mut();
            let st = (api.GetTensorTypeAndShape.unwrap())(value, &mut info);
            if !st.is_null() || info.is_null() {
                if !st.is_null() {
                    (api.ReleaseStatus.unwrap())(st);
                }
                return Ok(false);
            }
            (api.ReleaseTensorTypeAndShapeInfo.unwrap())(info);
            Ok(true)
        }
    }

    fn sequence_input(&self, r: &TensorRef) -> Result<*const ort::OrtValue, MlxError> {
        if r.source != Src::CtxInput {
            return Err("MLX: sequence input must cross the ORT kernel boundary".into());
        }
        unsafe {
            let api = &*self.ort_api;
            let value = read_ctx_input_value(self.ort_api, self.kctx, r.ctx_index)?;
            let mut kind = 0;
            ort_status(
                api,
                (api.GetValueType.unwrap())(value, &mut kind),
                "GetValueType(sequence)",
            )?;
            if kind != ort::ONNXType_ONNX_TYPE_SEQUENCE {
                return Err("MLX: expected an ORT sequence input".into());
            }
            Ok(value)
        }
    }

    pub fn sequence_len(&self, r: &TensorRef) -> Result<usize, MlxError> {
        let value = self.sequence_input(r)?;
        unsafe {
            let api = &*self.ort_api;
            let mut count = 0;
            ort_status(
                api,
                (api.GetValueCount.unwrap())(value, &mut count),
                "GetValueCount(sequence)",
            )?;
            Ok(count)
        }
    }

    pub fn sequence_element(
        &mut self,
        r: &TensorRef,
        index: usize,
        expected_dtype: ort::ONNXTensorElementDataType,
    ) -> Result<mlxsys::mlx_array, MlxError> {
        let value = self.sequence_input(r)?;
        let index = i32::try_from(index)
            .map_err(|_| "MLX: sequence index exceeds the ORT GetValue range".to_string())?;
        unsafe {
            let api = &*self.ort_api;
            let mut allocator = std::ptr::null_mut();
            ort_status(
                api,
                (api.GetAllocatorWithDefaultOptions.unwrap())(&mut allocator),
                "GetAllocatorWithDefaultOptions",
            )?;
            if allocator.is_null() {
                return Err("MLX: ORT returned a null sequence allocator".into());
            }
            let mut element = std::ptr::null_mut();
            let result = ort_status(
                api,
                (api.GetValue.unwrap())(value, index, allocator, &mut element),
                "GetValue(sequence)",
            )
            .and_then(|()| {
                if element.is_null() {
                    return Err("MLX: ORT returned a null sequence element".into());
                }
                let (data, shape, dtype) = read_ort_tensor_raw(api, element)?;
                if dtype != expected_dtype {
                    return Err("MLX: sequence element dtype differs from the output dtype".into());
                }
                let shape = shape
                    .into_iter()
                    .map(dim_i32)
                    .collect::<Result<Vec<_>, _>>()?;
                // GetValue owns a temporary OrtValue: copy before releasing it, never borrow.
                Ok(Array::from_data(data, &shape, mlx_dtype_from_onnx(dtype)))
            });
            if !element.is_null() {
                (api.ReleaseValue.unwrap())(element);
            }
            result.map(|array| self.keep(array))
        }
    }

    /// Read an int32/int64 host parameter exactly, widening int32 values before use.
    pub fn read_ints(&self, r: &TensorRef) -> Result<Vec<i64>, MlxError> {
        let h = self.raw_host(r)?;
        if !matches!(
            h.dtype,
            ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32
                | ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64
        ) {
            return Err(format!(
                "MLX: integer parameter {} must be int32/int64",
                r.name
            ));
        }
        if h.count == 0 {
            return Ok(Vec::new());
        }
        if h.data.is_null() {
            return Err(format!("MLX: integer parameter {} has null data", r.name));
        }
        if h.dtype == ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32 {
            return Ok(
                unsafe { std::slice::from_raw_parts(h.data as *const i32, h.count) }
                    .iter()
                    .map(|&value| i64::from(value))
                    .collect(),
            );
        }
        let p = h.data as *const i64;
        Ok(unsafe { std::slice::from_raw_parts(p, h.count) }.to_vec())
    }

    fn read_ctx_input(
        &self,
        index: usize,
    ) -> Result<(*const c_void, Vec<i64>, ort::ONNXTensorElementDataType), MlxError> {
        read_ctx_input_raw(self.ort_api, self.kctx, index)
    }

    // ---- MLX op helpers (each keeps and returns the raw result) --------------------------------

    /// Apply a unary `mlx_*(res, a, stream)` op.
    pub fn unary(
        &mut self,
        op: unsafe extern "C" fn(
            *mut mlxsys::mlx_array,
            mlxsys::mlx_array,
            mlxsys::mlx_stream,
        ) -> i32,
        a: mlxsys::mlx_array,
    ) -> Result<mlxsys::mlx_array, MlxError> {
        let mut res = Array::new();
        let mut raw = res.as_raw();
        let rc = unsafe { op(&mut raw, a, self.stream) };
        // The op may replace the handle; re-wrap whatever it produced.
        res = Array::from_raw(raw);
        if rc != 0 {
            return Err("mlx unary op failed".to_string());
        }
        Ok(self.keep(res))
    }

    /// Apply a binary `mlx_*(res, a, b, stream)` op.
    pub fn binary(
        &mut self,
        op: unsafe extern "C" fn(
            *mut mlxsys::mlx_array,
            mlxsys::mlx_array,
            mlxsys::mlx_array,
            mlxsys::mlx_stream,
        ) -> i32,
        a: mlxsys::mlx_array,
        b: mlxsys::mlx_array,
    ) -> Result<mlxsys::mlx_array, MlxError> {
        let mut res = Array::new();
        let mut raw = res.as_raw();
        let rc = unsafe { op(&mut raw, a, b, self.stream) };
        res = Array::from_raw(raw);
        if rc != 0 {
            return Err("mlx binary op failed".to_string());
        }
        Ok(self.keep(res))
    }

    /// `astype(a, t)` — cast to another dtype.
    pub fn astype(
        &mut self,
        a: mlxsys::mlx_array,
        t: mlxsys::mlx_dtype,
    ) -> Result<mlxsys::mlx_array, MlxError> {
        let mut res = Array::new();
        let mut raw = res.as_raw();
        let rc = unsafe { mlxsys::mlx_astype(&mut raw, a, t, self.stream) };
        res = Array::from_raw(raw);
        if rc != 0 {
            return Err("mlx_astype failed".to_string());
        }
        Ok(self.keep(res))
    }

    /// Cast once per source handle and target dtype during this graph translation.
    pub fn astype_once(
        &mut self,
        a: mlxsys::mlx_array,
        t: mlxsys::mlx_dtype,
    ) -> Result<mlxsys::mlx_array, MlxError> {
        let key = (a.ctx as usize, t);
        if let Some(&casted) = self.astype_cache.get(&key) {
            return Ok(casted);
        }
        let casted = self.astype(a, t)?;
        self.astype_cache.insert(key, casted);
        Ok(casted)
    }

    /// Materialize and retain a cast of a genuinely constant array in the plan cache.
    pub fn cache_astype(
        &mut self,
        key: String,
        a: mlxsys::mlx_array,
        t: mlxsys::mlx_dtype,
    ) -> Result<mlxsys::mlx_array, MlxError> {
        if let Some(casted) = self.cache_get(&key) {
            return Ok(casted);
        }
        let mut raw = unsafe { mlxsys::mlx_array_new() };
        let rc = unsafe { mlxsys::mlx_astype(&mut raw, a, t, self.stream) };
        if rc != 0 {
            unsafe { mlxsys::mlx_array_free(raw) };
            return Err("mlx_astype failed".to_string());
        }
        let casted = Array::from_raw(raw);
        casted.eval();
        Ok(self.cache_put(key, casted))
    }

    /// `zeros_like(a)`.
    pub fn zeros_like(&mut self, a: mlxsys::mlx_array) -> Result<mlxsys::mlx_array, MlxError> {
        let mut res = Array::new();
        let mut raw = res.as_raw();
        let rc = unsafe { mlxsys::mlx_zeros_like(&mut raw, a, self.stream) };
        res = Array::from_raw(raw);
        if rc != 0 {
            return Err("mlx_zeros_like failed".to_string());
        }
        Ok(self.keep(res))
    }

    /// Generic result-producing MLX op runner: builds a fresh result array, invokes the closure with
    /// `(&mut result, stream)`, re-wraps whatever handle the op produced (RAII), errors on non-zero
    /// return, and keeps + returns the raw result. This replaces the per-signature helper boilerplate
    /// so each op handler is a one-liner regardless of the underlying `mlx_*` arity.
    pub fn emit<F>(&mut self, f: F) -> Result<mlxsys::mlx_array, MlxError>
    where
        F: FnOnce(*mut mlxsys::mlx_array, mlxsys::mlx_stream) -> i32,
    {
        let mut res = Array::new();
        let mut raw = res.as_raw();
        let rc = f(&mut raw, self.stream);
        res = Array::from_raw(raw);
        if rc != 0 {
            return Err("mlx op failed".to_string());
        }
        Ok(self.keep(res))
    }

    /// Shape-preserving cumulative reduction.
    ///
    /// MLX 0.32's backing `Scan` primitive lacks `Primitive::output_shapes`, so every registered
    /// handler that can reach this helper must declare `CompileShapeSafety::ShapeKeyedOnly`.
    pub fn cumsum(
        &mut self,
        a: mlxsys::mlx_array,
        axis: i32,
        reverse: bool,
        inclusive: bool,
    ) -> Result<mlxsys::mlx_array, MlxError> {
        let dtype = mlxsys::mlx_optional_dtype {
            value: mlxsys::mlx_dtype__MLX_BOOL,
            has_value: false,
        };
        self.emit(|res, stream| unsafe {
            mlxsys::mlx_cumsum_axis(res, a, axis, reverse, inclusive, dtype, stream)
        })
    }

    /// Shape-preserving cumulative product with the same MLX Scan shape-inference constraint.
    pub fn cumprod(
        &mut self,
        a: mlxsys::mlx_array,
        axis: i32,
        reverse: bool,
        inclusive: bool,
    ) -> Result<mlxsys::mlx_array, MlxError> {
        let dtype = mlxsys::mlx_optional_dtype {
            value: mlxsys::mlx_dtype__MLX_BOOL,
            has_value: false,
        };
        self.emit(|res, stream| unsafe {
            mlxsys::mlx_cumprod_axis(res, a, axis, reverse, inclusive, dtype, stream)
        })
    }

    // ---- array introspection (borrowed raw handles; ownership stays with the arena/cache) ---------

    pub fn shape_of(&self, a: mlxsys::mlx_array) -> Vec<i32> {
        let nd = unsafe { mlxsys::mlx_array_ndim(a) };
        let sh = unsafe { mlxsys::mlx_array_shape(a) };
        (0..nd).map(|i| unsafe { *sh.add(i) }).collect()
    }

    pub fn ndim(&self, a: mlxsys::mlx_array) -> usize {
        unsafe { mlxsys::mlx_array_ndim(a) }
    }

    pub fn dim(&self, a: mlxsys::mlx_array, i: i32) -> i32 {
        unsafe { mlxsys::mlx_array_dim(a, i) }
    }

    pub fn size_of(&self, a: mlxsys::mlx_array) -> usize {
        unsafe { mlxsys::mlx_array_size(a) }
    }

    pub fn dtype_of(&self, a: mlxsys::mlx_array) -> mlxsys::mlx_dtype {
        unsafe { mlxsys::mlx_array_dtype(a) }
    }

    // ---- constant materialization helpers ---------------------------------------------------------

    /// A kept 0-d float32 scalar array.
    pub fn scalar_f32(&mut self, v: f32) -> mlxsys::mlx_array {
        self.keep(Array::from_raw(unsafe { mlxsys::mlx_array_new_float32(v) }))
    }

    /// A kept 0-d float64 scalar array. Used by the fp64 CPU-stream path so a constant like Selu's
    /// alpha is not first rounded through fp32.
    pub fn scalar_f64(&mut self, v: f64) -> mlxsys::mlx_array {
        let sh: [i32; 0] = [];
        self.keep(Array::from_data(
            &v as *const f64 as *const c_void,
            &sh,
            mlxsys::mlx_dtype__MLX_FLOAT64,
        ))
    }

    /// A kept 0-d int32 scalar array.
    pub fn scalar_i32(&mut self, v: i32) -> mlxsys::mlx_array {
        self.keep(Array::from_raw(unsafe { mlxsys::mlx_array_new_int(v) }))
    }

    /// A kept 0-d int64 scalar array.
    pub fn scalar_i64(&mut self, v: i64) -> mlxsys::mlx_array {
        let sh: [i32; 0] = [];
        self.keep(Array::from_data(
            &v as *const i64 as *const c_void,
            &sh,
            mlxsys::mlx_dtype__MLX_INT64,
        ))
    }

    /// A kept 0-d scalar built byte-exactly from host data of the given MLX dtype.
    ///
    /// Deliberately does not widen through FLOAT64 as a transport: the EP runs on a GPU stream and
    /// MLX rejects float64 there, so a widening round-trip would abort the process.
    pub fn scalar_from_bytes(
        &mut self,
        bytes: &[u8],
        dtype: mlxsys::mlx_dtype,
    ) -> mlxsys::mlx_array {
        let sh: [i32; 0] = [];
        self.keep(Array::from_data(
            bytes.as_ptr() as *const c_void,
            &sh,
            dtype,
        ))
    }

    /// A kept 1-D (or 0-D) int64 array wrapping host values (Shape/Size outputs).
    // builder-style: constructs an mlx_array from host data into the ctx arena.
    #[allow(clippy::wrong_self_convention)]
    pub fn from_host_i64(&mut self, data: &[i64], shape: &[i32]) -> mlxsys::mlx_array {
        self.keep(Array::from_data(
            data.as_ptr() as *const c_void,
            shape,
            mlxsys::mlx_dtype__MLX_INT64,
        ))
    }

    // ---- common data-movement helpers -------------------------------------------------------------

    pub fn reshape(
        &mut self,
        a: mlxsys::mlx_array,
        shape: &[i32],
    ) -> Result<mlxsys::mlx_array, MlxError> {
        self.emit(|res, s| unsafe { mlxsys::mlx_reshape(res, a, shape.as_ptr(), shape.len(), s) })
    }

    pub fn transpose(
        &mut self,
        a: mlxsys::mlx_array,
        axes: &[i32],
    ) -> Result<mlxsys::mlx_array, MlxError> {
        self.emit(|res, s| unsafe {
            mlxsys::mlx_transpose_axes(res, a, axes.as_ptr(), axes.len(), s)
        })
    }

    /// Force a (possibly strided/broadcast) view to row-major contiguous — required before a boundary
    /// output produced by a view op (transpose/slice/expand/split) is memcpy'd across the ORT boundary.
    pub fn contiguous(&mut self, a: mlxsys::mlx_array) -> Result<mlxsys::mlx_array, MlxError> {
        self.emit(|res, s| unsafe { mlxsys::mlx_contiguous(res, a, false, s) })
    }

    pub fn zeros(
        &mut self,
        shape: &[i32],
        dtype: mlxsys::mlx_dtype,
    ) -> Result<mlxsys::mlx_array, MlxError> {
        self.emit(|res, s| unsafe { mlxsys::mlx_zeros(res, shape.as_ptr(), shape.len(), dtype, s) })
    }

    /// Softmax over the last axis (precise), used by the Softmax handler.
    /// Softmax over a single (already-normalized, or negative) axis via `mlx_softmax_axis`.
    ///
    /// The `precise` flag asks MLX to accumulate in float32. That is what we want for fp16/bf16, but
    /// for a float64 input it is a *down*cast: an input near float32's max (3.4e38 is a perfectly
    /// ordinary float64) overflows to inf and the result comes back NaN. Ask for the precise path
    /// only when the input is actually narrower than float32.
    pub fn softmax_axis(
        &mut self,
        a: mlxsys::mlx_array,
        axis: i32,
    ) -> Result<mlxsys::mlx_array, MlxError> {
        let precise = self.dtype_of(a) != mlxsys::mlx_dtype__MLX_FLOAT64;
        let mut res = Array::new();
        let mut raw = res.as_raw();
        let rc = unsafe { mlxsys::mlx_softmax_axis(&mut raw, a, axis, precise, self.stream) };
        res = Array::from_raw(raw);
        if rc != 0 {
            return Err("mlx_softmax_axis failed".to_string());
        }
        Ok(self.keep(res))
    }

    // ---- extra op helpers shared by the signal/random/recurrent/ssm/misc/controlflow ops ---------

    /// `a * b` (elementwise, with MLX broadcast).
    pub fn mul(
        &mut self,
        a: mlxsys::mlx_array,
        b: mlxsys::mlx_array,
    ) -> Result<mlxsys::mlx_array, MlxError> {
        self.binary(mlxsys::mlx_multiply, a, b)
    }

    /// `a + b`.
    pub fn add(
        &mut self,
        a: mlxsys::mlx_array,
        b: mlxsys::mlx_array,
    ) -> Result<mlxsys::mlx_array, MlxError> {
        self.binary(mlxsys::mlx_add, a, b)
    }

    /// `a - b`.
    pub fn sub(
        &mut self,
        a: mlxsys::mlx_array,
        b: mlxsys::mlx_array,
    ) -> Result<mlxsys::mlx_array, MlxError> {
        self.binary(mlxsys::mlx_subtract, a, b)
    }

    /// `a @ b` (matmul).
    pub fn matmul(
        &mut self,
        a: mlxsys::mlx_array,
        b: mlxsys::mlx_array,
    ) -> Result<mlxsys::mlx_array, MlxError> {
        self.emit(|res, s| unsafe { mlxsys::mlx_matmul(res, a, b, s) })
    }

    /// Concatenate two arrays along `axis`.
    pub fn concat2(
        &mut self,
        a: mlxsys::mlx_array,
        b: mlxsys::mlx_array,
        axis: i32,
    ) -> Result<mlxsys::mlx_array, MlxError> {
        let mut vec = VectorArray::new();
        vec.append(a);
        vec.append(b);
        self.emit(|res, s| unsafe { mlxsys::mlx_concatenate_axis(res, vec.as_raw(), axis, s) })
    }

    /// Contiguous strided slice with unit stride (`start`/`stop` per axis).
    pub fn slice(
        &mut self,
        a: mlxsys::mlx_array,
        start: &[i32],
        stop: &[i32],
    ) -> Result<mlxsys::mlx_array, MlxError> {
        let stride = vec![1i32; start.len()];
        self.emit(|res, s| unsafe {
            mlxsys::mlx_slice(
                res,
                a,
                start.as_ptr(),
                start.len(),
                stop.as_ptr(),
                stop.len(),
                stride.as_ptr(),
                stride.len(),
                s,
            )
        })
    }

    /// `expand_dims(a, axis)`.
    pub fn expand_dims(
        &mut self,
        a: mlxsys::mlx_array,
        axis: i32,
    ) -> Result<mlxsys::mlx_array, MlxError> {
        self.emit(|res, s| unsafe { mlxsys::mlx_expand_dims(res, a, axis, s) })
    }

    /// `arange(start, stop, step)` of the given dtype.
    pub fn arange(
        &mut self,
        start: f64,
        stop: f64,
        step: f64,
        dtype: mlxsys::mlx_dtype,
    ) -> Result<mlxsys::mlx_array, MlxError> {
        self.emit(|res, s| unsafe { mlxsys::mlx_arange(res, start, stop, step, dtype, s) })
    }

    /// Elementwise `a <= b` (broadcasting), producing a bool array.
    pub fn less_equal(
        &mut self,
        a: mlxsys::mlx_array,
        b: mlxsys::mlx_array,
    ) -> Result<mlxsys::mlx_array, MlxError> {
        self.binary(mlxsys::mlx_less_equal, a, b)
    }

    /// `where(cond, x, y)` (elementwise select with broadcasting).
    pub fn where_(
        &mut self,
        cond: mlxsys::mlx_array,
        x: mlxsys::mlx_array,
        y: mlxsys::mlx_array,
    ) -> Result<mlxsys::mlx_array, MlxError> {
        self.emit(|res, s| unsafe { mlxsys::mlx_where(res, cond, x, y, s) })
    }

    /// Functional `slice_update` with a DATA start offset (an int array giving the start index for
    /// each of `axes`; unspecified axes start at 0). Returns the full-shape updated `src`.
    pub fn slice_update_dynamic(
        &mut self,
        src: mlxsys::mlx_array,
        update: mlxsys::mlx_array,
        start: mlxsys::mlx_array,
        axes: &[i32],
    ) -> Result<mlxsys::mlx_array, MlxError> {
        self.emit(|res, s| unsafe {
            mlxsys::mlx_slice_update_dynamic(res, src, update, start, axes.as_ptr(), axes.len(), s)
        })
    }

    /// `squeeze(a, axis)`.
    pub fn squeeze(
        &mut self,
        a: mlxsys::mlx_array,
        axis: i32,
    ) -> Result<mlxsys::mlx_array, MlxError> {
        self.emit(|res, s| unsafe { mlxsys::mlx_squeeze_axis(res, a, axis, s) })
    }

    /// Stack a list of same-shaped arrays along a new `axis`.
    pub fn stack(
        &mut self,
        parts: &[mlxsys::mlx_array],
        axis: i32,
    ) -> Result<mlxsys::mlx_array, MlxError> {
        let mut vec = VectorArray::new();
        for &p in parts {
            vec.append(p);
        }
        self.emit(|res, s| unsafe { mlxsys::mlx_stack_axis(res, vec.as_raw(), axis, s) })
    }

    /// A kept 0-d complex64 scalar array (real, imag).
    pub fn scalar_complex(&mut self, re: f32, im: f32) -> mlxsys::mlx_array {
        self.keep(Array::from_raw(unsafe {
            mlxsys::mlx_array_new_complex(re, im)
        }))
    }

    /// A kept 0-d bool scalar array.
    pub fn scalar_bool(&mut self, v: bool) -> mlxsys::mlx_array {
        self.keep(Array::from_raw(unsafe { mlxsys::mlx_array_new_bool(v) }))
    }

    /// A kept array wrapping host bytes of the given `dtype` and `shape`.
    // builder-style: constructs an mlx_array from host data into the ctx arena.
    #[allow(clippy::wrong_self_convention)]
    pub fn from_host(
        &mut self,
        data: *const c_void,
        shape: &[i32],
        dtype: mlxsys::mlx_dtype,
    ) -> mlxsys::mlx_array {
        self.keep(Array::from_data(data, shape, dtype))
    }

    /// Force `a` contiguous, evaluate it, and return `(kept handle, shape, dtype)` so a handler can
    /// read its host bytes mid-graph (the host-computed ops: Det / NonZero / Unique). The kept handle
    /// stays alive for the rest of the run.
    pub fn contiguous_eval(&mut self, a: mlxsys::mlx_array) -> Result<mlxsys::mlx_array, MlxError> {
        if self.in_general_trace {
            // A host-computed op (Det/NonZero/Unique) needs the input's DATA, but in a general
            // compiled trace the inputs are shapeless tracer placeholders with no data. Fail the
            // trace so `try_compiled_general` falls back to the eager translator (which evals real
            // data). Never eval a tracer here — that would abort inside MLX.
            return Err("MLX: host eval not permitted in general compiled trace".to_string());
        }
        let r = self.contiguous(a)?;
        let mut v = VectorArray::new();
        v.append(r);
        mlx::eval(&v)?;
        Ok(r)
    }

    /// Raw host byte pointer of an (evaluated) array — valid until the array is freed.
    pub fn host_ptr(&self, a: mlxsys::mlx_array) -> *const u8 {
        unsafe { mlxsys::mlx_array_data_uint8(a) }
    }

    /// Like [`Self::contiguous_eval`] but permitted inside a general trace. ONLY call when the array
    /// is provably shape-const (a `Shape`/`Size`-derived value or a chain over constants): it has no
    /// tracer dependency, so eval computes a genuine constant rather than aborting on a tracer.
    fn contiguous_eval_const(
        &mut self,
        a: mlxsys::mlx_array,
    ) -> Result<mlxsys::mlx_array, MlxError> {
        let r = self.contiguous(a)?;
        let mut v = VectorArray::new();
        v.append(r);
        mlx::eval(&v)?;
        Ok(r)
    }

    /// Read an int64 vector honoring the ref's SOURCE. Constant initializers / subgraph boundary
    /// inputs are read from host bytes (cheap). A dynamic intermediate is materialized by a mid-graph
    /// eval — permitted inside a general trace ONLY when the ref is `shape_const` (no tracer
    /// dependency); a data-dependent runtime value in a general trace fails so the plan falls back to
    /// eager (where real data exists).
    pub fn read_ints_eval(&mut self, r: &TensorRef) -> Result<Vec<i64>, MlxError> {
        if !matches!(r.source, Src::Intermediate) {
            return self.read_ints(r);
        }
        if self.in_general_trace && !r.shape_const {
            return Err(
                "MLX: data-dependent runtime shape cannot be read in a general trace".to_string(),
            );
        }
        let a = self.resolve(r)?;
        if r.shape_const && self.host_shape_ops_enabled() && self.size_of(a) <= 64 {
            return self.with_shape_stream(|ctx| ctx.read_ints_array(a));
        }
        self.read_ints_array(a)
    }

    fn read_ints_array(&mut self, a: mlxsys::mlx_array) -> Result<Vec<i64>, MlxError> {
        let a = self.astype(a, mlxsys::mlx_dtype__MLX_INT64)?;
        let a = self.contiguous_eval_const(a)?;
        let n = self.size_of(a);
        if n == 0 {
            return Ok(Vec::new());
        }
        let p = unsafe { mlxsys::mlx_array_data_int64(a) };
        if p.is_null() {
            return Ok(Vec::new());
        }
        Ok(unsafe { std::slice::from_raw_parts(p, n) }.to_vec())
    }

    /// Evaluate a 0-d/1-element integer array and read its scalar value host-side (int32 or int64).
    /// Used by the eager GQA path to recover `total_sequence_length` (an in-graph scalar) so the
    /// valid-past length becomes a known integer for static slice bounds. This forces a small
    /// mid-graph eval, mirroring the host-computed ops (Det/NonZero/Unique via `contiguous_eval`).
    pub fn read_scalar_i64(&mut self, a: mlxsys::mlx_array) -> Result<i64, MlxError> {
        let dt = self.dtype_of(a);
        let a = if dt == mlxsys::mlx_dtype__MLX_INT64 {
            a
        } else {
            self.astype(a, mlxsys::mlx_dtype__MLX_INT64)?
        };
        let a = self.contiguous_eval(a)?;
        let mut v: i64 = 0;
        let rc = unsafe { mlxsys::mlx_array_item_int64(&mut v, a) };
        if rc != 0 {
            return Err("mlx_array_item_int64 failed".to_string());
        }
        Ok(v)
    }

    /// In-place KV write: functional `slice_update` of `update` into `src` over `[start, stop)`
    /// (unit stride). Returns the full-shape updated buffer (the untouched region carries `src`'s
    /// data through), matching the runtime's fixed-capacity shared KV buffer contract.
    pub fn slice_update(
        &mut self,
        src: mlxsys::mlx_array,
        update: mlxsys::mlx_array,
        start: &[i32],
        stop: &[i32],
    ) -> Result<mlxsys::mlx_array, MlxError> {
        let strides = vec![1i32; start.len()];
        self.emit(|res, s| unsafe {
            mlxsys::mlx_slice_update(
                res,
                src,
                update,
                start.as_ptr(),
                start.len(),
                stop.as_ptr(),
                stop.len(),
                strides.as_ptr(),
                strides.len(),
                s,
            )
        })
    }

    pub(crate) fn run_empty_scan_template(
        &mut self,
        body: &SubgraphDesc,
        inputs: &[mlxsys::mlx_array],
    ) -> Result<Vec<mlxsys::mlx_array>, MlxError> {
        let previous = std::mem::replace(&mut self.empty_scan_template, true);
        let result = self.run_subgraph(body, inputs);
        self.empty_scan_template = previous;
        result
    }

    /// Translate a control-flow body subgraph inline: bind its formal inputs to `inputs` (positional),
    /// dispatch every body node through the registry (implicit inputs fall through to the enclosing
    /// env), collect + return the arrays bound to the body's formal outputs, then restore the env.
    /// Faithful port of the C++ `TranslationContext::RunSubgraph`.
    pub fn run_subgraph(
        &mut self,
        sg: &SubgraphDesc,
        inputs: &[mlxsys::mlx_array],
    ) -> Result<Vec<mlxsys::mlx_array>, MlxError> {
        if inputs.len() != sg.input_names.len() {
            return Err(format!(
                "MLX RunSubgraph: input arity mismatch for body '{}'",
                sg.attr_name
            ));
        }
        // Names this body binds (formal inputs + every produced output); snapshot shadowed outer ones.
        let mut body_names: std::collections::HashSet<String> = std::collections::HashSet::new();
        for nm in &sg.input_names {
            if !nm.is_empty() {
                body_names.insert(nm.clone());
            }
        }
        for node in &sg.nodes {
            for o in &node.outputs {
                if !o.name.is_empty() {
                    body_names.insert(o.name.clone());
                }
            }
        }
        let mut saved: HashMap<String, mlxsys::mlx_array> = HashMap::new();
        for nm in &body_names {
            if let Some(&a) = self.env.get(nm) {
                saved.insert(nm.clone(), a);
            }
        }
        // Formal inputs shadow outer bindings and must not inherit frozen provenance.
        let mut saved_frozen = HashMap::new();
        for name in &body_names {
            if let Some(key) = self.frozen_views.remove(name) {
                saved_frozen.insert(name.clone(), key);
            }
        }
        // Bind formal inputs, then translate the body in topological order.
        for (i, nm) in sg.input_names.iter().enumerate() {
            if !nm.is_empty() {
                self.env.insert(nm.clone(), inputs[i]);
            }
        }
        for node in &sg.nodes {
            crate::registry::translate(self, node)?;
        }
        // Collect the body's formal outputs (before restoring the env).
        let mut outs = Vec::with_capacity(sg.output_names.len());
        for on in &sg.output_names {
            match self.env.get(on) {
                Some(&a) => outs.push(a),
                None => {
                    return Err(format!(
                        "MLX RunSubgraph: body '{}' did not produce output {on}",
                        sg.attr_name
                    ));
                }
            }
        }
        // Restore the enclosing scope.
        for nm in &body_names {
            self.env.remove(nm);
            self.frozen_views.remove(nm);
        }
        self.frozen_views.extend(saved_frozen);
        for (k, v) in saved {
            self.env.insert(k, v);
        }
        Ok(outs)
    }

    // ---- compiled decode support ---------------------------------------------------------------

    /// True while translating inside the compiled-decode closure trace (see [`Self::rope_dynamic`]).
    #[inline]
    pub fn rope_dynamic(&self) -> bool {
        self.rope_dynamic
    }

    /// True when the compiled-decode trace is for a fixed-capacity shared KV buffer session.
    #[inline]
    pub fn shared_kv(&self) -> bool {
        self.shared_kv
    }

    /// Register a shared-buffer KV `present` output for delta copy-out. `offset`/`count` are the new
    /// K/V rows written this step (axis 2 of `[B,kv,cap,hd]`) and `past_ctx_index` is the ctx index
    /// of the matching `past` input (whose buffer `present` should alias). In the EAGER path this
    /// records the live per-run delta (including the `past` pointer) consumed by `copy_out`. In the
    /// COMPILED trace (`rope_dynamic`) `offset` is still the capacity placeholder, so we only remember
    /// `(name, past_ctx_index)` on the plan; the real per-token offset/count/pointer are supplied at
    /// apply time by `try_compiled_decode`.
    pub fn record_kv_present(
        &mut self,
        out_name: &str,
        offset: i64,
        count: i64,
        past_ctx_index: usize,
    ) {
        if self.rope_dynamic {
            if !self.compiled_kv_present.iter().any(|(n, _)| n == out_name) {
                self.compiled_kv_present
                    .push((out_name.to_string(), past_ctx_index));
            }
        } else {
            let alias_ptr = self
                .read_ctx_input(past_ctx_index)
                .map(|(p, _, _)| p as usize)
                .unwrap_or(0);
            self.kv_deltas.insert(
                out_name.to_string(),
                DeltaWrite {
                    axis: 2,
                    offset,
                    count,
                    alias_ptr,
                },
            );
        }
    }

    /// The live `[1]` int32 `valid_past` array seeded into the compiled shared-buffer trace (see
    /// [`GQA_VALID_PAST_KEY`]). `None` outside the compiled shared-buffer path.
    #[inline]
    pub fn shared_valid_past(&self) -> Option<mlxsys::mlx_array> {
        self.env.get(GQA_VALID_PAST_KEY).copied()
    }

    /// Put this context into the compiled-decode/prefill TRACE mode (`rope_dynamic = true`), recording
    /// whether the session drives a shared KV buffer. Used by the unified core while building the
    /// closure body (replaces the old decode-only `new_trace`). `shape_keyed` marks a prefill trace
    /// (static shapes for the key); `valid_past` is the static valid-past for that key (only used when
    /// `shape_keyed` is set, to statically narrow the shared-buffer attention to the valid prefix).
    pub(crate) fn set_compiled_trace(
        &mut self,
        shared_kv: bool,
        shape_keyed: bool,
        valid_past: i32,
    ) {
        self.rope_dynamic = true;
        self.shared_kv = shared_kv;
        self.compiled_shape_keyed = shape_keyed;
        self.compiled_valid_past = valid_past;
    }

    /// True while translating inside a SHAPE-KEYED compiled (prefill) trace — the shared-buffer
    /// attention can be statically narrowed to the valid prefix. Only meaningful with `rope_dynamic`.
    #[inline]
    pub fn compiled_shape_keyed(&self) -> bool {
        self.compiled_shape_keyed
    }

    /// The static `valid_past` for a shape-keyed compiled (prefill) trace (see `compiled_shape_keyed`).
    #[inline]
    pub fn compiled_valid_past(&self) -> i32 {
        self.compiled_valid_past
    }

    /// Mark this context as a GENERAL compiled-subgraph trace (see [`Self::in_general_trace`]).
    pub(crate) fn set_general_trace(&mut self) {
        self.in_general_trace = true;
    }

    /// Record whether this translation is happening inside a `ShapeMode::ShapeKeyed` compiled
    /// subgraph (`general` or `prefill`) as opposed to a `Shapeless` (decode) one or plain eager
    /// translation. Call once, right after `TranslationContext::new`, from the closure's
    /// `CompiledConfig::shape_mode` (see [`Self::shape_keyed_compile`]).
    pub(crate) fn set_shape_keyed_compile(&mut self, shape_keyed: bool) {
        self.in_compiled_trace = true;
        self.shape_keyed_compile = shape_keyed;
    }

    /// True while translating inside ANY shape-keyed compiled subgraph (general or prefill) — i.e. a
    /// multi-token, matrix-matrix compiled pass, never decode. See the field doc for the distinction
    /// from [`Self::compiled_shape_keyed`] (which is additionally scoped to the RoPE-as-data
    /// shared-buffer narrowing feature).
    #[inline]
    pub fn shape_keyed_compile(&self) -> bool {
        self.shape_keyed_compile
    }

    /// True while tracing the shape-polymorphic token-at-a-time decode closure.
    #[inline]
    pub fn shapeless_compile(&self) -> bool {
        self.in_compiled_trace && !self.shape_keyed_compile
    }

    pub(crate) fn native_attention_decode(&self) -> bool {
        self.plan.native_attention_decode
    }

    pub(crate) fn is_capturing(&self) -> bool {
        self.in_general_trace || self.in_compiled_trace || self.rope_dynamic
    }

    /// Take the shared-buffer KV `present` outputs discovered during the compiled trace, handing them
    /// to the core so it can store them on the traced `CompiledSubgraph` slot. Leaves the list empty.
    pub(crate) fn take_compiled_kv_present(&mut self) -> Vec<(String, usize)> {
        std::mem::take(&mut self.compiled_kv_present)
    }

    /// Seed an env binding (the compiled-closure placeholders: dynamic ctx inputs + pre-sliced RoPE
    /// rows). `raw` is a borrowed handle owned by the trace arena.
    pub(crate) fn seed(&mut self, name: String, raw: mlxsys::mlx_array) {
        self.frozen_views.remove(&name);
        self.env.insert(name, raw);
    }

    /// The raw handle bound to `name` in the env, if any (compiled-closure output collection).
    pub(crate) fn env_get(&self, name: &str) -> Option<mlxsys::mlx_array> {
        self.env.get(name).copied()
    }

    /// Take ownership of this run's transient arena (handing it to the plan so the compiled graph's
    /// handles outlive the trace). Leaves the context arena empty.
    pub(crate) fn take_arena(&mut self) -> Vec<Array> {
        std::mem::take(&mut self.arena)
    }

    /// Translate every node into MLX ops (no eval) and return the cast external boundary outputs as
    /// a fresh vector — the body of a compiled closure trace (decode, prefill, or general). Unlike
    /// [`Self::execute`] this does NOT eval or copy out; the compiled closure defers evaluation to the
    /// single per-step eval. When `contiguous` is set each boundary output is additionally
    /// materialised to row-major contiguous (mirroring [`Self::finish_boundary`]) so a general
    /// static-shape copy-out is a straight typed memcpy; the RoPE/KV paths leave it unset to mirror
    /// the decode trace exactly.
    pub(crate) fn run_trace(
        &mut self,
        ext_outputs: &[OutRef],
        contiguous: bool,
    ) -> Result<VectorArray, MlxError> {
        let nodes = std::mem::take(&mut self.plan.nodes);
        let mut result = Ok(());
        for node in &nodes {
            if let Err(e) = crate::registry::translate(self, node) {
                result = Err(e);
                break;
            }
        }
        self.plan.nodes = nodes;
        result?;
        // Cast each boundary output to its ORT output dtype so the per-step copy-out is a straight
        // typed memcpy (mirrors `finish_boundary`), and append in the fixed closure output order.
        let mut res = VectorArray::new();
        for o in ext_outputs {
            let a = self
                .env_get(&o.name)
                .ok_or_else(|| format!("MLX: compiled trace missing output {}", o.name))?;
            let casted = self.astype(a, mlx_dtype_from_onnx(o.otype))?;
            let casted = if contiguous {
                self.contiguous(casted)?
            } else {
                casted
            };
            res.append(casted);
        }
        Ok(res)
    }

    /// The pre-sliced full-width cos/sin RoPE row placeholder for `cache_name`, reshaped to
    /// `[1,1,S,2*half]` for broadcast over `[B,H,S,2*half]`. Compiled-decode trace only.
    pub fn rope_row_full(
        &mut self,
        cache_name: &str,
        seq: i32,
        half: i32,
    ) -> Result<mlxsys::mlx_array, MlxError> {
        let key = rope_row_key(cache_name);
        let row = self
            .env
            .get(&key)
            .copied()
            .ok_or_else(|| format!("MLX: missing RoPE row placeholder for {cache_name}"))?;
        self.reshape(row, &[1, 1, seq, 2 * half])
    }

    /// Constant `[hd,hd]` matrix `M` such that `x @ M == rotate_half(x)` for non-interleaved RoPE,
    /// i.e. `rotate_half([x1,x2]) = [-x2, x1]` with `half = hd/2`. Built once (fp32) and cached on the
    /// plan. Lets the compiled decode graph do rotate-half with a matmul instead of a Slice (which
    /// shapeless `mlx_compile` cannot shape-infer). The caller casts it to the operand dtype.
    pub fn rotate_half_matrix(&mut self, hd: i32, half: i32) -> mlxsys::mlx_array {
        let key = format!("__rope_rotate_half_{hd}");
        if let Some(a) = self.plan.cache.get(&key) {
            return a.as_raw();
        }
        let n = (hd as usize) * (hd as usize);
        let mut m = vec![0.0f32; n];
        for i in 0..(half as usize) {
            let hd_u = hd as usize;
            let half_u = half as usize;
            m[(i + half_u) * hd_u + i] = -1.0; // col i (<half) picks -x[i+half]
            m[i * hd_u + (i + half_u)] = 1.0; // col i+half picks  x[i]
        }
        let shp = [hd, hd];
        let arr = Array::from_data(
            m.as_ptr() as *const c_void,
            &shp,
            mlxsys::mlx_dtype__MLX_FLOAT32,
        );
        self.cache_put(key, arr)
    }

    /// Translate every node, cast+collect boundary outputs, one `mlx_eval`, copy each output back
    /// across the ORT boundary. Faithful port of `ExecuteEager`.
    pub fn execute(&mut self) -> Result<(), MlxError> {
        let nodes = std::mem::take(&mut self.plan.nodes);
        let mut result = Ok(());
        let tr = crate::trace::tracer();
        {
            // Timing attribution: the eager per-node translation (graph build) phase.
            let _phase = tr.phase("translate");
            for node in &nodes {
                if let Err(e) = crate::registry::translate(self, node) {
                    result = Err(e);
                    break;
                }
            }
        }
        if result.is_ok() {
            result = self.finish_boundary(&nodes);
        }
        // Restore the plan's node list (we only borrowed it).
        self.plan.nodes = nodes;
        result
    }

    fn finish_boundary(&mut self, nodes: &[NodeDesc]) -> Result<(), MlxError> {
        // Collect boundary outputs, cast each to its ORT output dtype BEFORE eval so copy-out is a
        // straight typed memcpy, then eval the whole graph in one shot.
        let mut outs = VectorArray::new();
        let mut ext: Vec<OutRef> = Vec::new();
        for node in nodes {
            for o in &node.outputs {
                if o.external
                    && !self.patched_alias_outputs.contains(&o.name)
                    && let Some(&a) = self.env.get(&o.name)
                {
                    let casted = self.astype(a, mlx_dtype_from_onnx(o.otype))?;
                    // Materialise to row-major contiguous HERE (once, at the boundary) so the
                    // flat copy_out memcpy is valid. Intermediate view ops (transpose/slice/
                    // expand/split) therefore stay zero-copy strided views that MLX folds into
                    // consuming kernels; only actual subgraph outputs pay a contiguous copy.
                    let casted = self.contiguous(casted)?;
                    outs.append(casted);
                    ext.push(o.clone());
                }
            }
        }
        let mut writebacks = Vec::with_capacity(self.input_writebacks.len());
        for writeback in self.input_writebacks.drain(..) {
            let index = outs.size();
            outs.append(writeback.array);
            writebacks.push((writeback.destination, index, writeback.copies));
        }
        // VectorArray owns references to every boundary root, whose lazy graphs keep their
        // dependencies alive. Per-node handles must not pin every intermediate allocation
        // after its final consumer has executed (notably large unpacked lookup matrices).
        self.env.clear();
        self.astype_cache.clear();
        self.temporary_arrays.clear();
        self.frozen_views.clear();
        self.input_writebacks.clear();
        self.arena.clear();
        let mut device_outputs =
            !ext.is_empty() && writebacks.is_empty() && !self.borrowed_runtime_inputs;
        if device_outputs {
            for (index, output) in ext.iter().enumerate() {
                let array = outs.get(index);
                device_outputs &= publish_device_output(self.ort_api, self.kctx, output, &array)?;
            }
        }
        let tr = crate::trace::tracer();
        // A captured boundary must finish before its guard stops Metal capture.
        let capture = tr.begin_gpu_capture();
        if device_outputs {
            if capture.is_some() {
                let _eval = tr.eval_region();
                mlx::eval(&outs)?;
                return Ok(());
            }
            let result = unsafe { mlxsys::mlx_async_eval(outs.as_raw()) };
            if result != 0 {
                return Err(
                    "MLX: asynchronous device output evaluation failed; inspect GPU availability"
                        .into(),
                );
            }
            return Ok(());
        }
        // The single synchronous `mlx_eval` boundary: with tracing on this is wrapped
        // in the `mlx.eval` (cat "gpu") span, whose CPU wall time is the GPU-inclusive
        // time of the whole fused subgraph (MLX blocks here until the GPU work lands).
        // Sample GPU-memory counters just before and after so the curve shows the eval.
        tr.sample_gpu_counters();
        let eval_t0 = if tr.active() {
            Some(std::time::Instant::now())
        } else {
            None
        };
        {
            let _eval = tr.eval_region();
            mlx::eval(&outs)?;
        }
        if let Some(t0) = eval_t0 {
            // Timing attribution: the synchronous GPU-inclusive eval phase (summary).
            tr.record_phase("eval", t0.elapsed());
        }
        tr.sample_gpu_counters();
        for (index, o) in ext.iter().enumerate() {
            let a = outs.get(index);
            self.copy_out(o, a.as_raw())?;
        }
        for (dst, index, copies) in writebacks {
            let arr = outs.get(index);
            let started = tr.active().then(std::time::Instant::now);
            let mut bytes = 0;
            for (range, offset) in copies {
                bytes += range.len();
                unsafe {
                    std::ptr::copy(
                        arr.data_bytes().add(range.start),
                        (dst as *mut u8).add(offset),
                        range.len(),
                    );
                }
            }
            if let Some(started) = started {
                tr.record_copyout(true, bytes as u64, started.elapsed());
            }
        }
        Ok(())
    }

    /// Create the ORT output tensor with the MLX result shape and memcpy on unified memory. For a
    /// shared-buffer KV `present` output only the newly written rows are copied (delta write).
    fn copy_out(&self, o: &OutRef, a: mlxsys::mlx_array) -> Result<(), MlxError> {
        let delta = self.kv_deltas.get(&o.name).copied();
        copy_out_raw_delta(self.ort_api, self.kctx, o, a, delta)
    }
}

/// Read a ctx input's `(data ptr, shape, dtype)` directly from the kernel context (no
/// `TranslationContext` needed — used by the compiled-decode path).
fn ort_status(
    api: &ort::OrtApi,
    status: *mut ort::OrtStatus,
    operation: &str,
) -> Result<(), MlxError> {
    if status.is_null() {
        return Ok(());
    }
    unsafe {
        let detail = std::ffi::CStr::from_ptr((api.GetErrorMessage.unwrap())(status))
            .to_string_lossy()
            .into_owned();
        (api.ReleaseStatus.unwrap())(status);
        Err(format!("MLX: {operation} failed: {detail}"))
    }
}

pub(crate) fn read_ort_tensor_metadata(
    api: &ort::OrtApi,
    value: *const ort::OrtValue,
) -> Result<(Vec<i64>, ort::ONNXTensorElementDataType), MlxError> {
    unsafe {
        let mut info = std::ptr::null_mut();
        let result = ort_status(
            api,
            (api.GetTensorTypeAndShape.unwrap())(value, &mut info),
            "GetTensorTypeAndShape",
        )
        .and_then(|()| {
            if info.is_null() {
                return Err("MLX: ORT returned null tensor metadata".into());
            }
            let mut rank = 0;
            ort_status(
                api,
                (api.GetDimensionsCount.unwrap())(info, &mut rank),
                "GetDimensionsCount",
            )?;
            let mut shape = vec![0; rank];
            ort_status(
                api,
                (api.GetDimensions.unwrap())(info, shape.as_mut_ptr(), rank),
                "GetDimensions",
            )?;
            let mut dtype = 0;
            ort_status(
                api,
                (api.GetTensorElementType.unwrap())(info, &mut dtype),
                "GetTensorElementType",
            )?;
            Ok((shape, dtype))
        });
        if !info.is_null() {
            (api.ReleaseTensorTypeAndShapeInfo.unwrap())(info);
        }
        result
    }
}

fn read_ort_tensor_raw(
    api: &ort::OrtApi,
    value: *const ort::OrtValue,
) -> Result<(*const c_void, Vec<i64>, ort::ONNXTensorElementDataType), MlxError> {
    let (shape, dtype) = read_ort_tensor_metadata(api, value)?;
    if let Some(data) = crate::device::host_data(api, value)? {
        return Ok((data, shape, dtype));
    }
    let mut data = std::ptr::null();
    ort_status(
        api,
        unsafe { (api.GetTensorData.unwrap())(value, &mut data) },
        "GetTensorData",
    )?;
    if !shape.contains(&0) && data.is_null() {
        return Err("MLX: ORT returned null data for a nonempty tensor".into());
    }
    Ok((data, shape, dtype))
}

pub(crate) fn read_ctx_input_value(
    ort_api: *const ort::OrtApi,
    kctx: *mut ort::OrtKernelContext,
    index: usize,
) -> Result<*const ort::OrtValue, MlxError> {
    unsafe {
        let api = &*ort_api;
        let mut val: *const ort::OrtValue = std::ptr::null();
        ort_status(
            api,
            (api.KernelContext_GetInput.unwrap())(kctx, index, &mut val),
            "KernelContext_GetInput",
        )?;
        if val.is_null() {
            return Err(format!(
                "MLX: KernelContext_GetInput({index}) returned no value"
            ));
        }
        Ok(val)
    }
}

pub(crate) fn read_ctx_input_metadata(
    ort_api: *const ort::OrtApi,
    kctx: *mut ort::OrtKernelContext,
    index: usize,
) -> Result<(Vec<i64>, ort::ONNXTensorElementDataType), MlxError> {
    let value = read_ctx_input_value(ort_api, kctx, index)?;
    read_ort_tensor_metadata(unsafe { &*ort_api }, value)
}

pub(crate) fn publish_device_output(
    ort_api: *const ort::OrtApi,
    kctx: *mut ort::OrtKernelContext,
    output: &OutRef,
    array: &Array,
) -> Result<bool, MlxError> {
    if !crate::device::enabled() {
        return Ok(false);
    }
    let api = unsafe { &*ort_api };
    let shape = array.shape();
    let mut value = std::ptr::null_mut();
    ort_status(
        api,
        unsafe {
            (api.KernelContext_GetOutput.unwrap())(
                kctx,
                output.ctx_index,
                shape.as_ptr(),
                shape.len(),
                &mut value,
            )
        },
        "KernelContext_GetOutput for device tensor",
    )?;
    if value.is_null() {
        return Err("MLX: ORT returned no requested device output; check output bindings".into());
    }
    crate::device::assign(api, value, array)
}

pub(crate) fn read_ctx_input_raw(
    ort_api: *const ort::OrtApi,
    kctx: *mut ort::OrtKernelContext,
    index: usize,
) -> Result<(*const c_void, Vec<i64>, ort::ONNXTensorElementDataType), MlxError> {
    let value = read_ctx_input_value(ort_api, kctx, index)?;
    read_ort_tensor_raw(unsafe { &*ort_api }, value)
}

/// Create the ORT output tensor with the MLX result shape and memcpy on unified memory; when
/// `delta` is `Some` only the `count` rows at `offset` along `axis` are memcpy'd back (shared-buffer
/// `axis` are memcpy'd back (shared-buffer KV `present`: the rest of the buffer already holds the
/// correct `past` rows because `present` aliases `past` in ORT memory). The ORT output tensor is
/// still created at the MLX array's FULL shape so ORT sees the whole `[B,kv,cap,hd]` present.
pub(crate) fn copy_out_raw_delta(
    ort_api: *const ort::OrtApi,
    kctx: *mut ort::OrtKernelContext,
    o: &OutRef,
    a: mlxsys::mlx_array,
    delta: Option<DeltaWrite>,
) -> Result<(), MlxError> {
    let arr = std::mem::ManuallyDrop::new(Array::from_raw(a)); // borrow, do not free
    let shape = arr.shape();
    let count: usize = shape.iter().map(|&d| d as usize).product();
    let itemsize = arr.itemsize();
    if publish_device_output(ort_api, kctx, o, &arr)? {
        return Ok(());
    }
    unsafe {
        let api = &*ort_api;
        let mut out: *mut ort::OrtValue = std::ptr::null_mut();
        let st = (api.KernelContext_GetOutput.unwrap())(
            kctx,
            o.ctx_index,
            shape.as_ptr(),
            shape.len(),
            &mut out,
        );
        if !st.is_null() || out.is_null() {
            return Err(format!(
                "MLX: KernelContext_GetOutput({}) failed",
                o.ctx_index
            ));
        }
        let mut dst: *mut c_void = std::ptr::null_mut();
        (api.GetTensorMutableData.unwrap())(out, &mut dst);
        let src = arr.data_bytes();
        if src.is_null() || dst.is_null() {
            return Ok(());
        }
        // Memory + timing view: record whether this copy-out took the delta (new-rows-only) or full
        // path, the bytes moved, and the wall time. Gated so a traced-off run pays one atomic load.
        let tr = crate::trace::tracer();
        let t0 = if tr.active() {
            Some(std::time::Instant::now())
        } else {
            None
        };
        let mut was_delta = false;
        let mut moved_bytes: u64 = (count * itemsize) as u64;
        match delta {
            // A read-only cache can retain the managed input storage through to its aliased output.
            // copy_nonoverlapping forbids even identical pointers.
            _ if src == dst as *const u8 => {
                moved_bytes = 0;
            }
            // Partial copy: only the newly written rows along `axis`, and ONLY when `present` really
            // aliases `past` in ORT memory (dst pointer matches the recorded `past` address). `past`
            // supplies every other row, so they are already correct. For a row-major [outer.., axis,
            // inner..] layout each of the `outer` slabs contributes one contiguous run of
            // `count*inner` elements at `(outer_idx*axis_len + offset)*inner`.
            Some(d)
                if d.count > 0
                    && d.axis < shape.len()
                    && d.alias_ptr != 0
                    && d.alias_ptr == dst as usize =>
            {
                let axis_len = shape[d.axis].max(0) as usize;
                let offset = d.offset.max(0) as usize;
                let n_rows = (d.count as usize).min(axis_len.saturating_sub(offset));
                if n_rows == 0 {
                    return Ok(());
                }
                let outer: usize = shape[..d.axis]
                    .iter()
                    .map(|&x| x as usize)
                    .product::<usize>()
                    .max(1);
                let inner: usize = shape[d.axis + 1..]
                    .iter()
                    .map(|&x| x as usize)
                    .product::<usize>()
                    .max(1);
                let run = n_rows * inner * itemsize; // bytes per outer slab
                let stride = axis_len * inner * itemsize; // bytes between slabs
                let start = offset * inner * itemsize; // byte offset of first row within a slab
                for o_idx in 0..outer {
                    let byte_off = o_idx * stride + start;
                    std::ptr::copy_nonoverlapping(
                        src.add(byte_off),
                        (dst as *mut u8).add(byte_off),
                        run,
                    );
                }
                was_delta = true;
                moved_bytes = (outer * run) as u64;
            }
            // No delta, or `present` did NOT alias `past`: full contiguous copy (growing path is
            // bit-for-bit unchanged). IOBinding can place read-only cache outputs in overlapping
            // input storage, so the full-copy fallback must permit overlap.
            _ => {
                std::ptr::copy(src, dst as *mut u8, count * itemsize);
            }
        }
        if let Some(t0) = t0 {
            tr.record_copyout(was_delta, moved_bytes, t0.elapsed());
        }
    }
    Ok(())
}

fn input_writeback_ranges(
    shape: &[i64],
    itemsize: usize,
    axis: usize,
    rows: &[usize],
) -> Result<Vec<std::ops::Range<usize>>, MlxError> {
    let invalid = || {
        format!(
            "MLX: input writeback rows exceed shape {shape:?} at axis {axis} with {itemsize}-byte elements; check cache geometry and slot mapping"
        )
    };
    if axis >= shape.len() || itemsize == 0 {
        return Err(invalid());
    }
    let dims = shape
        .iter()
        .map(|&dim| usize::try_from(dim).map_err(|_| invalid()))
        .collect::<Result<Vec<_>, _>>()?;
    let row_count = dims[..=axis]
        .iter()
        .try_fold(1usize, |n, &dim| n.checked_mul(dim))
        .ok_or_else(invalid)?;
    let row_bytes = dims[axis + 1..]
        .iter()
        .try_fold(itemsize, |n, &dim| n.checked_mul(dim))
        .ok_or_else(invalid)?;
    row_count.checked_mul(row_bytes).ok_or_else(invalid)?;
    let mut rows = rows.to_vec();
    rows.sort_unstable();
    rows.dedup();
    let mut ranges: Vec<std::ops::Range<usize>> = Vec::new();
    for row in rows {
        if row >= row_count {
            return Err(invalid());
        }
        if row_bytes == 0 {
            continue;
        }
        let start = row * row_bytes;
        let end = start + row_bytes;
        if let Some(last) = ranges.last_mut()
            && last.end == start
        {
            last.end = end;
        } else {
            ranges.push(start..end);
        }
    }
    Ok(ranges)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn input_writeback_rows_are_checked_sorted_and_coalesced() {
        let ranges = input_writeback_ranges(&[3, 4, 2, 8], 2, 1, &[9, 2, 1, 2]).unwrap();
        assert_eq!(ranges, vec![32..96, 288..320]);
        assert_eq!(
            input_writeback_ranges(&[2, 3], 4, 0, &[1]).unwrap(),
            vec![12..24]
        );
        assert!(
            input_writeback_ranges(&[3, 4, 2, 8], 2, 1, &[])
                .unwrap()
                .is_empty()
        );
        assert!(input_writeback_ranges(&[3, 4, 2, 8], 2, 1, &[12]).is_err());
        assert!(input_writeback_ranges(&[3, 4], 2, 2, &[0]).is_err());
        assert!(input_writeback_ranges(&[3, -1], 2, 0, &[0]).is_err());
        assert!(input_writeback_ranges(&[i64::MAX, 4], 8, 0, &[0]).is_err());
    }

    fn empty_context(plan: &mut Plan) -> TranslationContext<'_> {
        TranslationContext::new(
            plan,
            std::ptr::null(),
            std::ptr::null_mut(),
            mlxsys::mlx_stream {
                ctx: std::ptr::null_mut(),
            },
        )
    }

    #[test]
    fn empty_scan_template_restores_scope_on_success_and_error() {
        let mut plan = Plan::new(vec![]);
        let mut ctx = empty_context(&mut plan);
        let mut body = SubgraphDesc {
            attr_name: "body".into(),
            input_names: vec![],
            output_names: vec![],
            nodes: vec![],
        };
        assert!(ctx.run_empty_scan_template(&body, &[]).is_ok());
        assert!(!ctx.empty_scan_template());
        body.input_names.push("missing".into());
        assert!(ctx.run_empty_scan_template(&body, &[]).is_err());
        assert!(!ctx.empty_scan_template());
        ctx.empty_scan_template = true;
        assert!(ctx.run_empty_scan_template(&body, &[]).is_err());
        assert!(ctx.empty_scan_template());
        body.input_names.clear();
        assert!(ctx.run_empty_scan_template(&body, &[]).is_ok());
        assert!(ctx.empty_scan_template());
    }

    #[test]
    fn owned_initializer_cache_keys_preserve_storage_identity() {
        let mut first = ctx_input("cos", 0);
        let data = std::sync::Arc::new(vec![0u8; 4]);
        first.source = Src::Initializer;
        first.init = Some(InitData {
            data: data.as_ptr().cast(),
            shape: vec![1],
            dtype: ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT,
            count: 1,
            owned: Some(data),
        });
        assert_eq!(first.cache_key(), first.clone().cache_key());
        assert_ne!(first.cache_key(), ctx_input("cos", 0).cache_key());
        let mut second = first.clone();
        let data = std::sync::Arc::new(vec![0u8; 4]);
        second.init.as_mut().unwrap().data = data.as_ptr().cast();
        second.init.as_mut().unwrap().owned = Some(data);
        assert_ne!(first.cache_key(), second.cache_key());
    }

    fn input(name: &str) -> TensorRef {
        TensorRef {
            name: name.to_string(),
            source: Src::Intermediate,
            ctx_index: 0,
            constant: false,
            shape_const: false,
            init: None,
        }
    }

    fn ctx_input(name: &str, ctx_index: usize) -> TensorRef {
        TensorRef {
            name: name.to_string(),
            source: Src::CtxInput,
            ctx_index,
            constant: false,
            shape_const: false,
            init: None,
        }
    }

    fn output(name: &str) -> OutRef {
        OutRef {
            name: name.to_string(),
            external: false,
            ctx_index: 0,
            otype: ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT,
        }
    }

    fn node(op_type: &str, inputs: &[&str], output_name: &str) -> NodeDesc {
        let mut node = NodeDesc::new(
            op_type.to_string(),
            String::new(),
            17,
            crate::registry::CompileShapeSafety::Shapeless,
        );
        node.inputs = inputs.iter().map(|name| input(name)).collect();
        node.outputs = vec![output(output_name)];
        node
    }

    fn optimized_whisper_nodes(hidden_size: i64) -> Vec<NodeDesc> {
        let mut nodes = vec![
            node("Add", &["start", "length"], "limit"),
            node("Range", &["start", "limit", "step"], "range"),
            node("Unsqueeze", &["range", "axes"], "positions"),
            node("Concat", &["one", "batch"], "repeats"),
            node("Tile", &["positions", "repeats"], "tiled"),
            node("MultiHeadAttention", &["q", "k", "v"], "attention"),
        ];
        let mut vocab = node("MatMulNBits", &["hidden", "w", "s", "z"], "logits");
        vocab.ints.insert("bits".to_string(), 8);
        vocab.ints.insert("K".to_string(), hidden_size);
        vocab.ints.insert("N".to_string(), 51_865);
        nodes.push(vocab);
        while nodes.len() <= 32 {
            let index = nodes.len();
            nodes.push(node("Add", &["a", "b"], &format!("filler_{index}")));
        }
        nodes
    }

    #[test]
    fn optimized_decoder_position_graph_gets_dedicated_stream() {
        let tiny = Plan::new(optimized_whisper_nodes(384));
        assert!(!tiny.native_attention_decode);
        assert!(tiny.dedicated_decode_stream);

        let base = Plan::new(optimized_whisper_nodes(512));
        assert!(!base.native_attention_decode);
        assert!(base.dedicated_decode_stream);

        let mut ambiguous = optimized_whisper_nodes(384);
        ambiguous.push(node("Range", &["x", "y", "z"], "another_range"));
        assert!(!Plan::new(ambiguous).dedicated_decode_stream);

        let no_attention = optimized_whisper_nodes(384)
            .into_iter()
            .filter(|node| node.op_type != "MultiHeadAttention")
            .collect();
        assert!(!Plan::new(no_attention).dedicated_decode_stream);

        let mut native_attention = optimized_whisper_nodes(384);
        native_attention
            .iter_mut()
            .find(|node| node.op_type == "MultiHeadAttention")
            .unwrap()
            .op_type = "Attention".to_string();
        assert!(Plan::new(native_attention).dedicated_decode_stream);
    }

    #[test]
    fn stable_cross_inputs_require_a_boundary_kv_pair() {
        let mut cross = node("MultiHeadAttention", &["q", "k", "v"], "attention");
        cross.inputs[1] = ctx_input("k", 4);
        cross.inputs[2] = ctx_input("v", 5);
        let mut self_attention = node(
            "MultiHeadAttention",
            &["q", "k", "v", "bias", "mask", "rel", "past_k", "past_v"],
            "self_attention",
        );
        self_attention.inputs[6] = ctx_input("past_k", 6);
        self_attention.inputs[7] = ctx_input("past_v", 7);
        assert_eq!(
            Plan::new(vec![self_attention, cross]).stable_cross_ctx_indices,
            vec![4, 5]
        );

        let mut partial = node("MultiHeadAttention", &["q", "k", "v"], "attention");
        partial.inputs[1] = ctx_input("k", 4);
        assert!(Plan::new(vec![partial]).stable_cross_ctx_indices.is_empty());

        let mut native = node("Attention", &["q", "k", "v"], "attention");
        native.inputs[1] = ctx_input("k", 8);
        native.inputs[2] = ctx_input("v", 9);
        let mut native_self = node(
            "Attention",
            &["q", "k", "v", "mask", "past_k", "past_v"],
            "self_attention",
        );
        native_self.inputs[4] = ctx_input("past_k", 10);
        native_self.inputs[5] = ctx_input("past_v", 11);
        assert_eq!(
            Plan::new(vec![native_self, native]).stable_cross_ctx_indices,
            vec![8, 9]
        );

        let mut contrib = node("Attention", &["q", "w", "bias"], "attention");
        contrib.domain = "com.microsoft".to_string();
        contrib.inputs[1] = ctx_input("w", 8);
        contrib.inputs[2] = ctx_input("bias", 9);
        assert!(Plan::new(vec![contrib]).stable_cross_ctx_indices.is_empty());
    }
}

#[cfg(test)]
mod frozen_view_tests {
    use super::*;
    use crate::mlx::Stream;
    use crate::registry::CompileShapeSafety;

    #[test]
    fn subgraph_formals_shadow_frozen_views_and_restore_outer_provenance() {
        let stream = Stream::new_default_cpu();
        let mut plan = Plan::new(vec![]);
        let mut ctx = TranslationContext::new(
            &mut plan,
            std::ptr::null(),
            std::ptr::null_mut(),
            stream.as_raw(),
        );
        let mut cast = NodeDesc::new("Cast".into(), "".into(), 25, CompileShapeSafety::Shapeless);
        cast.inputs.push(TensorRef {
            name: "weight".into(),
            source: Src::Intermediate,
            ctx_index: 0,
            constant: false,
            shape_const: false,
            init: None,
        });
        cast.outputs.push(OutRef {
            name: "rounded".into(),
            external: false,
            ctx_index: 0,
            otype: ort::ONNXTensorElementDataType_ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT,
        });
        let body = SubgraphDesc {
            attr_name: "body".into(),
            input_names: vec!["weight".into()],
            output_names: vec!["rounded".into()],
            nodes: vec![cast],
        };
        ctx.frozen_views
            .insert("weight".into(), "outer immutable weight".into());
        for value in [1.0f32, 2.0] {
            let data = vec![value; 16384];
            let input = ctx.keep(Array::from_data(
                data.as_ptr().cast(),
                &[16384],
                mlxsys::mlx_dtype__MLX_FLOAT32,
            ));
            let outputs = ctx.run_subgraph(&body, &[input]).unwrap();
            unsafe {
                assert_eq!(mlxsys::mlx_array_eval(outputs[0]), 0);
                assert_eq!(*mlxsys::mlx_array_data_float32(outputs[0]), value);
            }
            assert!(
                ctx.plan.cache.is_empty(),
                "a changing formal must not enter the frozen cache"
            );
            assert_eq!(
                ctx.frozen_views.get("weight").map(String::as_str),
                Some("outer immutable weight")
            );
            assert!(!ctx.frozen_views.contains_key("rounded"));
        }
    }
}