earthsci-ast 0.6.0

//! Step 5 (phase 2): emit the tape IR as an **XLA computation**.
//!
//! This is the Rust half of the compiled-backend plan
//! (`tests/conformance/compiled_rhs/README.md`). It walks a [`TapeProgram`]
//! once and builds, through the `xla` crate's `XlaBuilder`, a single program
//!
//! ```text
//! rhs(u: f64[N], p: f64[M], t: f64[]) -> du: f64[N]
//! ```
//!
//! A program that reads the forcing buffer ([`Instr::LoadForcing`]) takes a
//! fourth parameter, `f: f64[F]`: every forcing it loads, row-major, one after
//! another ([`ForcingFeed`] says where each one sits and packs it).
//!
//! that reproduces what [`super::exec`] / `super::refexec` compute for the
//! same inputs — numerically, within the tier's tolerance classes, never bit
//! for bit (XLA's `exp`/`log`/`pow` are not Rust's libm).
//!
//! ## Hard error, never a fallback
//!
//! A model whose tape declined a rule (there is then no program at all; see
//! [`DeclinedTape`]), or whose program carries ANY instruction this emitter
//! cannot lower, is a **compile error** naming the rule and the reason
//! ([`XlaEmitError`]). There is no mixed execution and no interpreter arm: the
//! adapter turns the error into the tier's `refused` outcome and moves on to
//! the next fixture. This is the ruling of 2026-09-13 made mechanical.
//!
//! ## Layout conventions this file has to honour
//!
//! * **A slot's array value is ROW-MAJOR over `SlotDesc::shape`**, the same
//!   convention `refexec`'s `ArrayD` values use. Scalars are rank-0.
//! * **The flat state vector `u` is COLUMN-MAJOR per variable block**, while
//!   the tape reads state arrays row-major (that is exactly what the fast
//!   executor's `state_rm` mirror is for). So an [`Operand::State`] is
//!   `u[off..off+n]` reshaped to the REVERSED shape and transposed — the
//!   column-major-to-row-major relabelling, which costs nothing in XLA.
//! * **`dy` is column-major per variable block** for the same reason, so the
//!   root does the inverse transpose per variable and concatenates the blocks
//!   in `flat_offset` order.
//!
//! ## Cadence
//!
//! CONST, SEGMENT and CONTINUOUS are all emitted into the ONE program, in
//! program order. That recomputes the parameter-invariant prefix on every
//! call, which the slab executor hoists. Hoisting it here — emitting the
//! CONST/SEGMENT prefix as a separate computation whose outputs are extra
//! parameters of the continuous one, or letting XLA constant-fold it behind a
//! donated buffer — is a later optimization, deliberately not done in this
//! phase: it changes nothing about the numbers and everything about the
//! amount of machinery under test.
//!
//! ## Fusion
//!
//! The emitter is built over the UNFUSED tape: [`emit_rhs`] calls
//! `build_tape_opts(.., None)`, so [`Instr::Fused`] never reaches the match.
//! A fused group is a scalar micro-program over a register file — precisely
//! the shape XLA's own fusion wants to choose for itself.
//!
//! ## Reductions are numerically, not bitwise, equal
//!
//! [`Instr::Reduce`] carries a load-bearing visiting order — row-major over
//! the contraction box, which IS the per-cell oracle's odometer — because
//! floating-point `+` is not associative. XLA's `reduce` does NOT pin the
//! association: the backend is free to split the fold into lanes and combine
//! them in any order. So a compiled reduction is bit-identical to the
//! interpreter only by luck, and this emitter does not claim otherwise. The
//! tier's `reduction` tolerance class (rtol 1e-11, atol scaled by the probe's
//! largest magnitude) is what that lowering is measured against, and it was
//! written for exactly this kind of summation-order difference.
//!
//! ## When the tape grows
//!
//! The instruction match ends in a catch-all arm that REFUSES with a message
//! saying the tape grew and this emitter has to follow. New instruction kinds
//! (array constants, carried observed shapes, rank-0 reductions) therefore
//! surface as named refusals rather than as wrong numbers. That arm is where
//! they plug in.

use super::super::{ArrayCompiled, BinCode, UnCode, VarShape};
use super::ir::*;
use ndarray::ArrayD;
use std::cell::RefCell;
use std::collections::{HashMap, HashSet};
use std::fmt;
use std::rc::Rc;
use xla::{ArrayElement, ElementType, PrimitiveType, XlaBuilder, XlaComputation, XlaOp};

/// Why a recurrence sweep ([`Instr::Sweep`]) is refused: the emitter lowers
/// neither the ordered per-cell loop nor a channel for the causal self-read's
/// fail-closed fault, and CONFORMANCE_SPEC §5.19.3b wants a route that cannot
/// agree to decline loudly rather than return a number.
const RECURRENCE_REFUSAL: &str = "a causal self-reference (recurrence, esm-spec §4.3.1.1) is a \
     sequential sweep whose self-read fails closed at run time; this emitter lowers neither \
     the ordered per-cell loop nor a channel for that fault, so the model is refused \
     (CONFORMANCE_SPEC §5.19.3b)";

/// A model the emitter refuses to lower.
///
/// `rule` names the tape rule (or the pseudo-rule `<program>` for a
/// whole-program property) and `reason` says what could not be lowered. The
/// conformance adapter copies both straight into the tier's `refused` entry,
/// which is why they are separate fields rather than one message.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct XlaEmitError {
    /// The rule whose lowering failed, or `<program>`.
    pub rule: String,
    /// Why it could not be lowered.
    pub reason: String,
}

impl XlaEmitError {
    pub fn new(rule: impl Into<String>, reason: impl Into<String>) -> Self {
        XlaEmitError {
            rule: rule.into(),
            reason: reason.into(),
        }
    }
}

impl fmt::Display for XlaEmitError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "rule {}: {}", self.rule, self.reason)
    }
}

impl std::error::Error for XlaEmitError {}

/// Which way round [`Emitter::interp_count`] compares a knot with the query.
#[derive(Clone, Copy)]
enum CountCmp {
    /// `axis[k] <= q` — the cell search of `interp.linear` / `interp.bilinear`.
    AxisLeQuery,
    /// `xs[k] < q` — the prefix `interp.searchsorted` counts.
    AxisLtQuery,
}

/// Row-major strides (elements) of `shape` — the layout every slot value has.
fn rm_strides(shape: &[usize]) -> Vec<i64> {
    let mut st = vec![0i64; shape.len()];
    let mut acc = 1i64;
    for d in (0..shape.len()).rev() {
        st[d] = acc;
        acc *= shape[d] as i64;
    }
    st
}

/// A lowered right-hand side: the computation plus the shapes its caller has
/// to feed it.
pub struct EmittedRhs {
    computation: XlaComputation,
    n_states: usize,
    params_len: usize,
    n_instrs: usize,
    opcode_counts: Vec<(&'static str, usize)>,
    forcing: Option<ForcingFeed>,
}

/// Where the forcing parameter `f` of a program that reads the forcing buffer
/// comes from, and how it is laid out.
///
/// Each entry is one forcing the program loads, at its offset in `f`, in the
/// order the program first loads them. [`Self::pack`] fills `f` from the
/// model's live buffer with exactly the tape's load semantics, faults
/// included, so a missing entry fails the solve with the interpreter's own
/// error rather than as a device failure.
pub struct ForcingFeed {
    buffer: Rc<RefCell<HashMap<String, ArrayD<f64>>>>,
    declared: HashSet<String>,
    layout: Vec<(ForcingRef, usize)>,
    len: usize,
}

impl ForcingFeed {
    /// Length of `f` (at least one element, like `p`).
    pub fn len(&self) -> usize {
        self.len.max(1)
    }

    /// Always false: `f` is never zero-length (see [`Self::len`]).
    pub fn is_empty(&self) -> bool {
        false
    }

    /// `f` as the buffer stands now.
    pub fn pack(&self) -> Vec<f64> {
        let mut f = vec![0.0f64; self.len()];
        let buffer = self.buffer.borrow();
        for (fr, off) in &self.layout {
            let n = super::exec::forcing_len(fr);
            super::exec::load_forcing(fr, &buffer, &self.declared, false, &mut f[*off..*off + n]);
        }
        f
    }
}

impl EmittedRhs {
    /// The `rhs(u, p, t) -> du` computation.
    pub fn computation(&self) -> &XlaComputation {
        &self.computation
    }
    /// Length of the flat state vector `u` and of the result `du`.
    pub fn n_states(&self) -> usize {
        self.n_states
    }
    /// Length of the `p` buffer the compiled program expects, which is
    /// `max(tape params, 1)`: a model with no parameters still gets a
    /// one-element dummy, because a zero-element PJRT buffer is an awkward
    /// thing to build on every call for no gain. A caller with fewer values
    /// than this zero-fills the rest (`CompiledRhs::eval` does).
    pub fn params_len(&self) -> usize {
        self.params_len
    }
    /// Tape instruction count that was lowered (diagnostics only).
    pub fn n_instrs(&self) -> usize {
        self.n_instrs
    }
    /// How many instructions of each opcode the lowered tape carried
    /// (diagnostics only), in first-seen program order.
    pub fn opcode_counts(&self) -> &[(&'static str, usize)] {
        &self.opcode_counts
    }
    /// The forcing parameter's feed, for a program that reads the forcing
    /// buffer; `None` for the three-parameter form.
    pub fn forcing(&self) -> Option<&ForcingFeed> {
        self.forcing.as_ref()
    }
    /// Take the forcing feed, for the runtime that will own it.
    pub fn take_forcing(&mut self) -> Option<ForcingFeed> {
        self.forcing.take()
    }
    /// HLO text of the computation, for a debug dump. NEVER a gate: the tier
    /// compares numbers, never programs.
    pub fn hlo_text(&self) -> Result<String, XlaEmitError> {
        self.computation
            .proto()
            .to_string()
            .map_err(|e| XlaEmitError::new("<program>", format!("HLO dump failed: {e}")))
    }
}

/// Build this model's tape (fusion OFF) and emit it as an XLA computation.
///
/// The single entry point outside this module. Everything the emitter needs
/// beyond the tape — the flat state layout, the parameter count, the working
/// precision — it reads off `compiled`.
pub fn emit_rhs(compiled: &ArrayCompiled) -> Result<EmittedRhs, XlaEmitError> {
    // Evaluate under the model's own captured precision, exactly as
    // `debug_eval_rhs` does: the tape lowering consults it.
    let _guard = compiled.precision.enter();
    if compiled.precision.is_f32() {
        return Err(XlaEmitError::new(
            "<program>",
            "document working precision is Float32; this emitter lowers the \
             Float64 kernel tables only (the tier's precision-changing \
             fixtures are excluded, CONFORMANCE_SPEC §5.38.4)",
        ));
    }
    if crate::precision::has_variable_overrides() {
        return Err(XlaEmitError::new(
            "<program>",
            "document carries per-variable element types (esm-spec §11.3.1); \
             the tape itself is disabled for such documents",
        ));
    }
    // `None` = the fusion pass off. The emitter is defined over the unfused
    // instruction set; see the module docs.
    // Row-major: the emitter lowers every box in its logical axis order.
    let (prog, _report) =
        compiled.build_tape_layout(&std::collections::HashSet::new(), None, false);
    match prog {
        Ok(prog) => emit_program(&prog, compiled),
        Err(declined) => {
            let (rule, reason) = declined
                .declined_rules()
                .into_iter()
                .next()
                .expect("a declined tape names at least one rule");
            Err(XlaEmitError::new(
                rule,
                format!(
                    "rule is not on the tape (the lowering bailed: {reason}); a \
                     compiled program cannot re-enter the interpreter, so the whole \
                     model is refused"
                ),
            ))
        }
    }
}

/// Emit an already-built program. Split out so a test can drive a program it
/// built itself.
pub(crate) fn emit_program(
    prog: &TapeProgram,
    compiled: &ArrayCompiled,
) -> Result<EmittedRhs, XlaEmitError> {
    debug_assert!(!prog.col_major, "the emitter lowers row-major programs");
    let builder = XlaBuilder::new("earthsci_rhs");
    let mut em = Emitter::new(&builder, prog, compiled)?;
    let du = em.run()?;
    let forcing = em.forcing_param.take().map(|_| ForcingFeed {
        buffer: Rc::clone(&compiled.forcing),
        declared: compiled.declared_names.clone(),
        layout: em.forcing_layout.clone(),
        len: em.forcing_len,
    });
    let computation = du
        .build()
        .map_err(|e| XlaEmitError::new("<program>", format!("XlaBuilder::build failed: {e}")))?;
    Ok(EmittedRhs {
        computation,
        n_states: compiled.n_states,
        // The DECLARED buffer length, which is what a caller has to feed:
        // `Emitter::new` widens a zero-parameter model to `f64[1]`.
        params_len: prog.params_len.max(1),
        n_instrs: prog.instrs.len(),
        opcode_counts: opcode_counts(prog),
        forcing,
    })
}

/// Instructions per opcode, in first-seen program order.
fn opcode_counts(prog: &TapeProgram) -> Vec<(&'static str, usize)> {
    let mut counts: Vec<(&'static str, usize)> = Vec::new();
    for i in &prog.instrs {
        let op = i.opcode();
        match counts.iter_mut().find(|(o, _)| *o == op) {
            Some((_, n)) => *n += 1,
            None => counts.push((op, 1)),
        }
    }
    counts
}

// ---------------------------------------------------------------------------
// The emitter.
// ---------------------------------------------------------------------------

type R<T> = Result<T, XlaEmitError>;

struct Emitter<'a> {
    b: &'a XlaBuilder,
    prog: &'a TapeProgram,
    /// `u`, `p`, `t` parameters.
    u: XlaOp,
    p: XlaOp,
    t: XlaOp,
    /// Per-slot SSA value; `None` until the defining instruction is emitted.
    slots: Vec<Option<XlaOp>>,
    /// Per `prog.state_vars` index: the state array as the tape reads it
    /// (row-major over its shape; rank 0 for a 0-d variable).
    state: Vec<XlaOp>,
    /// `dy` accumulator, one ROW-MAJOR block per state variable, in
    /// `var_order`. `None` = nothing written yet (stays zero).
    blocks: Vec<Option<XlaOp>>,
    /// State variables ordered by `flat_offset`, with their layout.
    var_order: Vec<(String, VarShape)>,
    /// Variable name → index into `var_order` / `blocks`.
    var_ix: HashMap<String, usize>,
    /// Rule name attached to whatever instruction is being lowered, for the
    /// error's `rule` field.
    cur_rule: String,
    /// The fourth parameter `f`, present when the program loads a forcing.
    forcing_param: Option<XlaOp>,
    /// Per `prog.forcings` index: its offset in `f`.
    forcing_offsets: Vec<Option<usize>>,
    /// The loaded forcings in `f` order, and `f`'s length.
    forcing_layout: Vec<(ForcingRef, usize)>,
    forcing_len: usize,
}

impl<'a> Emitter<'a> {
    fn new(b: &'a XlaBuilder, prog: &'a TapeProgram, compiled: &'a ArrayCompiled) -> R<Self> {
        let n_states = compiled.n_states;
        let pm = prog.params_len.max(1);
        let u = b
            .parameter(0, f64::TY, &[n_states as i64], "u")
            .map_err(|e| XlaEmitError::new("<program>", format!("parameter u: {e}")))?;
        let p = b
            .parameter(1, f64::TY, &[pm as i64], "p")
            .map_err(|e| XlaEmitError::new("<program>", format!("parameter p: {e}")))?;
        let t = b
            .parameter(2, f64::TY, &[], "t")
            .map_err(|e| XlaEmitError::new("<program>", format!("parameter t: {e}")))?;

        // The flat layout is the model's, not the tape's: a variable the tape
        // never touches still owns its slice of `du`.
        let mut var_order: Vec<(String, VarShape)> = compiled
            .var_shapes
            .iter()
            .map(|(k, v)| (k.clone(), v.clone()))
            .collect();
        var_order.sort_by_key(|(_, v)| v.flat_offset);
        let mut next = 0usize;
        for (name, vs) in &var_order {
            if vs.flat_offset != next {
                return Err(XlaEmitError::new(
                    "<program>",
                    format!(
                        "flat state layout is not contiguous: variable {name} starts at \
                         {} but {next} was expected",
                        vs.flat_offset
                    ),
                ));
            }
            next += vs.shape.iter().product::<usize>().max(1);
        }
        if next != n_states {
            return Err(XlaEmitError::new(
                "<program>",
                format!("flat state layout covers {next} slots but the model declares {n_states}"),
            ));
        }
        let var_ix: HashMap<String, usize> = var_order
            .iter()
            .enumerate()
            .map(|(i, (n, _))| (n.clone(), i))
            .collect();

        // `f` holds only the forcings an instruction loads, in the order the
        // program first loads them.
        let mut forcing_offsets: Vec<Option<usize>> = vec![None; prog.forcings.len()];
        let mut forcing_layout: Vec<(ForcingRef, usize)> = Vec::new();
        let mut forcing_len = 0usize;
        for ins in &prog.instrs {
            if let Instr::LoadForcing { forcing, .. } = ins {
                let ix = *forcing as usize;
                if forcing_offsets[ix].is_none() {
                    let fr = prog.forcings[ix].clone();
                    forcing_offsets[ix] = Some(forcing_len);
                    forcing_len += super::exec::forcing_len(&fr);
                    forcing_layout.push((fr, forcing_offsets[ix].expect("just set")));
                }
            }
        }
        let forcing_param = if forcing_layout.is_empty() {
            None
        } else {
            Some(
                b.parameter(3, f64::TY, &[forcing_len.max(1) as i64], "f")
                    .map_err(|e| XlaEmitError::new("<program>", format!("parameter f: {e}")))?,
            )
        };

        let mut em = Emitter {
            b,
            prog,
            u,
            p,
            t,
            slots: vec![None; prog.slots.len()],
            state: Vec::new(),
            blocks: vec![None; var_order.len()],
            var_order,
            var_ix,
            cur_rule: "<program>".to_string(),
            forcing_param,
            forcing_offsets,
            forcing_layout,
            forcing_len,
        };
        em.state = em.build_state_views()?;
        Ok(em)
    }

    /// One view per `prog.state_vars` entry: the variable's block of `u`,
    /// relabelled from column-major to row-major (see the module docs).
    fn build_state_views(&self) -> R<Vec<XlaOp>> {
        let mut out = Vec::with_capacity(self.prog.state_vars.len());
        for sv in &self.prog.state_vars {
            let off = sv.flat_offset as i64;
            let n: usize = sv.shape.iter().product::<usize>().max(1);
            let flat = self
                .u
                .slice_in_dim(off, off + n as i64, 1, 0)
                .map_err(|e| self.err(format!("state slice for {}: {e}", sv.name)))?;
            let op = if sv.shape.is_empty() {
                flat.reshape(&[])
                    .map_err(|e| self.err(format!("state reshape for {}: {e}", sv.name)))?
            } else {
                let rev: Vec<i64> = sv.shape.iter().rev().map(|&d| d as i64).collect();
                let perm: Vec<i64> = (0..sv.shape.len() as i64).rev().collect();
                flat.reshape(&rev)
                    .and_then(|a| a.transpose(&perm))
                    .map_err(|e| self.err(format!("state relabel for {}: {e}", sv.name)))?
            };
            out.push(op);
        }
        Ok(out)
    }

    fn err(&self, reason: impl Into<String>) -> XlaEmitError {
        XlaEmitError::new(self.cur_rule.clone(), reason)
    }

    /// Lower the whole program and return the root `du`.
    fn run(&mut self) -> R<XlaOp> {
        self.emit_range(0, self.prog.instrs.len())?;
        self.assemble_du()
    }

    // -- small builder helpers ---------------------------------------------

    fn dims(&self, op: &XlaOp) -> R<Vec<usize>> {
        op.dims().map_err(|e| self.err(format!("dims: {e}")))
    }

    fn c(&self, v: f64) -> R<XlaOp> {
        self.b.c0(v).map_err(|e| self.err(format!("constant: {e}")))
    }

    fn ci(&self, v: i64) -> R<XlaOp> {
        self.b
            .c0(v)
            .map_err(|e| self.err(format!("s64 constant: {e}")))
    }

    fn zeros(&self, dims: &[usize]) -> R<XlaOp> {
        let z = self.c(0.0)?;
        if dims.is_empty() {
            return Ok(z);
        }
        let d: Vec<i64> = dims.iter().map(|&x| x as i64).collect();
        z.broadcast(&d).map_err(|e| self.err(format!("zeros: {e}")))
    }

    /// A scalar constant with `like`'s shape (so every binary op below has two
    /// operands of exactly the same shape and never leans on XLA's implicit
    /// rank-0 broadcast).
    ///
    /// The constant is built on `like`'s OWN builder, not on `self.b`. That is
    /// what lets [`Self::bin`] and [`Self::un`] be reused inside a `reduce`
    /// sub-computation, which lives on a separate [`XlaBuilder`]; mixing ops
    /// from two builders is an error XLA reports far from its cause.
    fn splat_like(&self, like: &XlaOp, v: f64) -> R<XlaOp> {
        let d = self.dims(like)?;
        let c = like
            .builder()
            .c0(v)
            .map_err(|e| self.err(format!("constant: {e}")))?;
        if d.is_empty() {
            return Ok(c);
        }
        let d: Vec<i64> = d.iter().map(|&x| x as i64).collect();
        c.broadcast(&d).map_err(|e| self.err(format!("splat: {e}")))
    }

    /// Broadcast a rank-0 value to `dims`; pass anything else through after
    /// checking it already has that shape.
    fn to_shape(&self, op: &XlaOp, dims: &[usize]) -> R<XlaOp> {
        let have = self.dims(op)?;
        if have == dims {
            return Ok(op.clone());
        }
        if have.is_empty() {
            let d: Vec<i64> = dims.iter().map(|&x| x as i64).collect();
            return op
                .broadcast(&d)
                .map_err(|e| self.err(format!("broadcast to {dims:?}: {e}")));
        }
        Err(self.err(format!(
            "operand shape {have:?} does not match the instruction's box {dims:?}"
        )))
    }

    fn pred_to_f64(&self, op: XlaOp) -> R<XlaOp> {
        op.convert(PrimitiveType::F64)
            .map_err(|e| self.err(format!("pred -> f64: {e}")))
    }

    fn wrap(&self, r: Result<XlaOp, xla::Error>, what: &str) -> R<XlaOp> {
        r.map_err(|e| self.err(format!("{what}: {e}")))
    }

    // -- operands ----------------------------------------------------------

    /// Resolve an operand to its natural value (rank 0 for a scalar).
    fn operand(&self, op: &Operand) -> R<XlaOp> {
        match op {
            Operand::Lit(v) => self.c(*v),
            Operand::Time => Ok(self.t.clone()),
            Operand::Param(ix) => {
                let i = *ix as i64;
                self.wrap(
                    self.p
                        .slice_in_dim(i, i + 1, 1, 0)
                        .and_then(|s| s.reshape(&[])),
                    &format!("parameter {ix}"),
                )
            }
            Operand::Slot(s) => self.slots[*s as usize]
                .clone()
                .ok_or_else(|| self.err(format!("read of slot {s} before its definition"))),
            Operand::State(ix) => Ok(self.state[*ix as usize].clone()),
        }
    }

    fn src(&self, s: &SrcRef) -> R<XlaOp> {
        match s {
            SrcRef::Slot(id) => self.slots[*id as usize]
                .clone()
                .ok_or_else(|| self.err(format!("gather source slot {id} is undefined"))),
            SrcRef::State(ix) => Ok(self.state[*ix as usize].clone()),
        }
    }

    /// The output box of a slot, as a `dims` slice (empty for a scalar).
    fn out_dims(&self, slot: SlotId) -> Vec<usize> {
        let d = &self.prog.slots[slot as usize];
        if d.scalar {
            Vec::new()
        } else {
            d.shape.to_vec()
        }
    }

    // -- kernels -----------------------------------------------------------

    /// `binary_kernel_of(op)` in XLA. Operands already share a shape.
    fn bin(&self, op: BinCode, a: &XlaOp, b: &XlaOp) -> R<XlaOp> {
        let cmp = |r: Result<XlaOp, xla::Error>, what: &str| -> R<XlaOp> {
            self.pred_to_f64(self.wrap(r, what)?)
        };
        match op {
            BinCode::Add => self.wrap(a.add_(b), "add"),
            BinCode::Sub => self.wrap(a.sub_(b), "sub"),
            BinCode::Mul => self.wrap(a.mul_(b), "mul"),
            BinCode::Div => self.wrap(a.div_(b), "div"),
            BinCode::Pow => self.wrap(a.pow(b), "pow"),
            BinCode::Atan2 => self.wrap(a.atan2(b), "atan2"),
            // `f64::min`/`max` return the OTHER operand when one is NaN; XLA's
            // Min/Max propagate it. Re-select so the two agree on NaN, which
            // costs four ops and removes a whole class of silent divergence.
            BinCode::Min => self.nan_aware(a, b, true),
            BinCode::Max => self.nan_aware(a, b, false),
            BinCode::Eq => cmp(a.eq(b), "eq"),
            BinCode::Ne => cmp(a.ne(b), "ne"),
            BinCode::Lt => cmp(a.lt(b), "lt"),
            BinCode::Le => cmp(a.le(b), "le"),
            BinCode::Gt => cmp(a.gt(b), "gt"),
            BinCode::Ge => cmp(a.ge(b), "ge"),
            BinCode::And => {
                let za = self.splat_like(a, 0.0)?;
                let x = self.wrap(a.ne(&za), "and: a != 0")?;
                let y = self.wrap(b.ne(&za), "and: b != 0")?;
                self.pred_to_f64(self.wrap(x.and(&y), "and")?)
            }
            BinCode::Or => {
                let za = self.splat_like(a, 0.0)?;
                let x = self.wrap(a.ne(&za), "or: a != 0")?;
                let y = self.wrap(b.ne(&za), "or: b != 0")?;
                self.pred_to_f64(self.wrap(x.or(&y), "or")?)
            }
            BinCode::Unknown => Err(self.err(
                "binary operator resolved to BinCode::Unknown (the interpreter's NaN \
                 sentinel); refusing rather than emitting a silent NaN",
            )),
        }
    }

    fn nan_aware(&self, a: &XlaOp, b: &XlaOp, is_min: bool) -> R<XlaOp> {
        let base = if is_min {
            self.wrap(a.min(b), "min")?
        } else {
            self.wrap(a.max(b), "max")?
        };
        let a_nan = self.wrap(a.ne(a), "a is NaN")?;
        let b_nan = self.wrap(b.ne(b), "b is NaN")?;
        let r = self.wrap(b_nan.select(a, &base), "min/max: b NaN")?;
        self.wrap(a_nan.select(b, &r), "min/max: a NaN")
    }

    /// `unary_kernel_of(op)` in XLA. XLA's op set has no `tan`, `asin`,
    /// `acos`, `atan`, `log10` or the hyperbolics other than `tanh`, so those
    /// are synthesized from the ops it does have. Each identity is chosen to
    /// keep the small-argument and infinite cases right, not merely the
    /// algebra: `sinh`/`asinh`/`acosh` switch between a `log1p`/`expm1` form
    /// near the cancelling region and a plain form outside it.
    ///
    /// KNOWN LIMIT: the plain arms of `asinh`/`acosh` square their argument,
    /// so they overflow to infinity above about 1e154 where the true value is
    /// finite (~355). No fixture in the corpus reaches that; a model that does
    /// needs the log-shifted identity instead.
    fn un(&self, op: UnCode, a: &XlaOp) -> R<XlaOp> {
        let one = self.splat_like(a, 1.0)?;
        let two = self.splat_like(a, 2.0)?;
        let half = self.splat_like(a, 0.5)?;
        let zero = self.splat_like(a, 0.0)?;
        match op {
            UnCode::Exp => self.wrap(a.exp(), "exp"),
            UnCode::Ln => self.wrap(a.log(), "log"),
            UnCode::Log10 => {
                let l = self.wrap(a.log(), "log10: log")?;
                let c = self.splat_like(a, std::f64::consts::LN_10)?;
                self.wrap(l.div_(&c), "log10: div")
            }
            UnCode::Sqrt => self.wrap(a.sqrt(), "sqrt"),
            UnCode::Abs => self.wrap(a.abs(), "abs"),
            UnCode::Sign => {
                // The interpreter's kernel is `x>0 ? 1 : x<0 ? -1 : 0`, which
                // maps NaN and -0.0 to +0.0; XLA's `sign` maps NaN to NaN and
                // -0.0 to -0.0. Spell out the interpreter's form.
                let neg = self.splat_like(a, -1.0)?;
                let pos = self.wrap(a.gt(&zero), "sign: >0")?;
                let negp = self.wrap(a.lt(&zero), "sign: <0")?;
                let inner = self.wrap(negp.select(&neg, &zero), "sign: inner")?;
                self.wrap(pos.select(&one, &inner), "sign: outer")
            }
            UnCode::Floor => self.wrap(a.floor(), "floor"),
            UnCode::Ceil => self.wrap(a.ceil(), "ceil"),
            UnCode::Sin => self.wrap(a.sin(), "sin"),
            UnCode::Cos => self.wrap(a.cos(), "cos"),
            UnCode::Tan => {
                let s = self.wrap(a.sin(), "tan: sin")?;
                let c = self.wrap(a.cos(), "tan: cos")?;
                self.wrap(s.div_(&c), "tan: div")
            }
            UnCode::Asin => {
                // atan2(x, sqrt(1 - x^2)): exact at |x| = 1 (atan2(±1, 0)) and
                // NaN outside the domain, both matching `f64::asin`.
                let x2 = self.wrap(a.mul_(a), "asin: x^2")?;
                let r = self.wrap(one.sub_(&x2), "asin: 1-x^2")?;
                let s = self.wrap(r.sqrt(), "asin: sqrt")?;
                self.wrap(a.atan2(&s), "asin: atan2")
            }
            UnCode::Acos => {
                let x2 = self.wrap(a.mul_(a), "acos: x^2")?;
                let r = self.wrap(one.sub_(&x2), "acos: 1-x^2")?;
                let s = self.wrap(r.sqrt(), "acos: sqrt")?;
                self.wrap(s.atan2(a), "acos: atan2")
            }
            UnCode::Atan => self.wrap(a.atan2(&one), "atan"),
            UnCode::Sinh => {
                // |x| <= 0.5: (expm1(x) + expm1(x)/(expm1(x)+1))/2, which has
                // no cancellation as x -> 0. Outside: (exp(x) - exp(-x))/2,
                // which is also what keeps ±inf mapping to ±inf.
                let e = self.wrap(a.expm1(), "sinh: expm1")?;
                let d = self.wrap(e.add_(&one), "sinh: expm1+1")?;
                let q = self.wrap(e.div_(&d), "sinh: ratio")?;
                let small = self.wrap(
                    e.add_(&q).and_then(|s| s.mul_(&half)),
                    "sinh: small-argument form",
                )?;
                let ex = self.wrap(a.exp(), "sinh: exp")?;
                let en = self.wrap(a.neg().and_then(|n| n.exp()), "sinh: exp(-x)")?;
                let big =
                    self.wrap(ex.sub_(&en).and_then(|s| s.mul_(&half)), "sinh: plain form")?;
                let thresh = self.splat_like(a, 0.5)?;
                let absx = self.wrap(a.abs(), "sinh: |x|")?;
                let use_big = self.wrap(absx.gt(&thresh), "sinh: select")?;
                self.wrap(use_big.select(&big, &small), "sinh: merge")
            }
            UnCode::Cosh => {
                let ex = self.wrap(a.exp(), "cosh: exp")?;
                let en = self.wrap(a.neg().and_then(|n| n.exp()), "cosh: exp(-x)")?;
                self.wrap(ex.add_(&en).and_then(|s| s.mul_(&half)), "cosh")
            }
            UnCode::Tanh => self.wrap(a.tanh(), "tanh"),
            UnCode::Asinh => {
                let absx = self.wrap(a.abs(), "asinh: |x|")?;
                let x2 = self.wrap(absx.mul_(&absx), "asinh: x^2")?;
                let root = self.wrap(x2.add_(&one).and_then(|s| s.sqrt()), "asinh: sqrt")?;
                // small: log1p(a + a^2/(1 + sqrt(1+a^2)))
                let den = self.wrap(one.add_(&root), "asinh: 1+sqrt")?;
                let frac = self.wrap(x2.div_(&den), "asinh: ratio")?;
                let small = self.wrap(
                    absx.add_(&frac).and_then(|s| s.log1p()),
                    "asinh: small-argument form",
                )?;
                let big = self.wrap(absx.add_(&root).and_then(|s| s.log()), "asinh: plain form")?;
                let use_big = self.wrap(absx.gt(&two), "asinh: select")?;
                let mag = self.wrap(use_big.select(&big, &small), "asinh: merge")?;
                let negm = self.wrap(mag.neg(), "asinh: negate")?;
                let isneg = self.wrap(a.lt(&zero), "asinh: sign")?;
                self.wrap(isneg.select(&negm, &mag), "asinh: apply sign")
            }
            UnCode::Acosh => {
                // x < 1 must be NaN: both arms take a square root of a
                // negative number there, so both deliver it.
                let s = self.wrap(a.sub_(&one), "acosh: x-1")?;
                let sp2 = self.wrap(s.add_(&two), "acosh: x+1")?;
                let root = self.wrap(s.mul_(&sp2).and_then(|q| q.sqrt()), "acosh: sqrt")?;
                let small = self.wrap(
                    s.add_(&root).and_then(|q| q.log1p()),
                    "acosh: small-argument form",
                )?;
                let x2 = self.wrap(a.mul_(a), "acosh: x^2")?;
                let root2 = self.wrap(x2.sub_(&one).and_then(|q| q.sqrt()), "acosh: sqrt2")?;
                let big = self.wrap(a.add_(&root2).and_then(|q| q.log()), "acosh: plain form")?;
                let use_big = self.wrap(a.gt(&two), "acosh: select")?;
                self.wrap(use_big.select(&big, &small), "acosh: merge")
            }
            UnCode::Atanh => {
                // 0.5 * log1p(2x / (1 - x)): x = ±1 give ±inf and |x| > 1 give
                // NaN, matching `f64::atanh`.
                let num = self.wrap(two.mul_(a), "atanh: 2x")?;
                let den = self.wrap(one.sub_(a), "atanh: 1-x")?;
                let q = self.wrap(num.div_(&den), "atanh: ratio")?;
                self.wrap(q.log1p().and_then(|l| l.mul_(&half)), "atanh: 0.5*log1p")
            }
            UnCode::Not => {
                let p = self.wrap(a.eq(&zero), "not")?;
                self.pred_to_f64(p)
            }
            UnCode::Unknown => Err(self.err(
                "unary operator resolved to UnCode::Unknown (the interpreter's NaN \
                 sentinel); refusing rather than emitting a silent NaN",
            )),
        }
    }

    // -- the instruction loop ----------------------------------------------

    fn emit_range(&mut self, start: usize, end: usize) -> R<()> {
        let mut pc = start;
        while pc < end {
            if let Some(ord) = self.prog.provenance.get(pc) {
                if let Some(info) = self.prog.rules.get(*ord as usize) {
                    self.cur_rule = info.name.clone();
                }
            }
            if let Instr::JmpIfZero {
                cond,
                n_true,
                n_false,
            } = &self.prog.instrs[pc]
            {
                let (nt, nf) = (*n_true as usize, *n_false as usize);
                let cond = *cond;
                let t_start = pc + 1;
                let f_start = t_start + nt;
                let f_end = f_start + nf;
                if f_end > end {
                    return Err(self.err(format!(
                        "JmpIfZero at {pc} spans past the end of its section ({f_end} > {end})"
                    )));
                }
                self.emit_branch(&cond, t_start, f_start, f_end)?;
                pc = f_end;
                continue;
            }
            if let Instr::Sweep { .. } = &self.prog.instrs[pc] {
                return Err(self.err(RECURRENCE_REFUSAL));
            }
            self.emit_one(pc)?;
            pc += 1;
        }
        Ok(())
    }

    /// Lower a `JmpIfZero` pair as `select`.
    ///
    /// The tape short-circuits: the untaken branch NEVER executes. A single
    /// XLA program has no such thing, so both branches are emitted and the
    /// live-out slots selected. That is value-equivalent because every
    /// instruction the tape can put in a branch is a pure floating-point map —
    /// an untaken `log(-1)` or `1/0` produces a NaN or an infinity that the
    /// `select` then discards, and IEEE-754 arithmetic raises no trap that
    /// could escape. (`XlaOp::conditional` would preserve the short circuit
    /// literally; it is not used because it would need a separate builder and
    /// an explicit tuple of every live value crossing the branch, for no
    /// difference in the numbers.)
    ///
    /// A branch carrying a side effect — a `DyWrite` or an `Export` — is
    /// refused rather than guessed at.
    fn emit_branch(&mut self, cond: &Operand, t0: usize, f0: usize, f1: usize) -> R<()> {
        for pc in t0..f1 {
            match &self.prog.instrs[pc] {
                Instr::DyWrite { .. } | Instr::Export { .. } => {
                    return Err(self.err(format!(
                        "instruction {} inside a JmpIfZero branch has an effect beyond \
                         its slot; this emitter lowers branches as `select` over both \
                         arms and cannot order such an effect",
                        self.prog.instrs[pc].opcode()
                    )));
                }
                _ => {}
            }
        }
        let c = self.operand(cond)?;
        if !self.dims(&c)?.is_empty() {
            return Err(self.err("JmpIfZero condition is not a scalar"));
        }
        let cz = self.c(0.0)?;
        let pred0 = self.wrap(c.ne(&cz), "JmpIfZero: cond != 0")?;

        let before = self.slots.clone();
        let blocks_before = self.blocks.clone();
        self.emit_range(t0, f0)?;
        let after_t = std::mem::replace(&mut self.slots, before.clone());
        let blocks_t = std::mem::replace(&mut self.blocks, blocks_before.clone());
        self.emit_range(f0, f1)?;
        let after_f = std::mem::take(&mut self.slots);
        let blocks_f = std::mem::take(&mut self.blocks);
        if blocks_t.len() != blocks_f.len() {
            return Err(self.err("JmpIfZero branches disagree on the dy blocks"));
        }
        self.blocks = blocks_before;

        let mut merged = Vec::with_capacity(after_t.len());
        for i in 0..after_t.len() {
            let v = match (&after_t[i], &after_f[i]) {
                (Some(tv), Some(fv)) => {
                    if before[i].is_some() && std::ptr::eq(tv, fv) {
                        // Untouched by both arms.
                        Some(tv.clone())
                    } else {
                        let dims = self.dims(tv)?;
                        let fd = self.dims(fv)?;
                        if dims != fd {
                            return Err(self.err(format!(
                                "JmpIfZero phi slot {i} has shape {dims:?} on one arm and \
                                 {fd:?} on the other"
                            )));
                        }
                        let pred = if dims.is_empty() {
                            pred0.clone()
                        } else {
                            let d: Vec<i64> = dims.iter().map(|&x| x as i64).collect();
                            self.wrap(pred0.broadcast(&d), "JmpIfZero: broadcast cond")?
                        };
                        Some(self.wrap(pred.select(tv, fv), "JmpIfZero: select")?)
                    }
                }
                // Defined on one arm only: it can only be read inside that
                // arm (the tape joins its branches on a shared phi slot), so
                // keeping it is safe and keeps nested branches working.
                (Some(v), None) | (None, Some(v)) => Some(v.clone()),
                (None, None) => None,
            };
            merged.push(v);
        }
        self.slots = merged;
        Ok(())
    }

    fn emit_one(&mut self, pc: usize) -> R<()> {
        let instr = self.prog.instrs[pc].clone();
        match &instr {
            Instr::Bin { op, a, b, out } => {
                let dims = self.out_dims(*out);
                let av = self.operand(a)?;
                let bv = self.operand(b)?;
                let av = self.to_shape(&av, &dims)?;
                let bv = self.to_shape(&bv, &dims)?;
                let v = self.bin(*op, &av, &bv)?;
                self.define(*out, v);
            }
            Instr::Un { op, a, out } => {
                let dims = self.out_dims(*out);
                let av = self.operand(a)?;
                let av = self.to_shape(&av, &dims)?;
                let v = self.un(*op, &av)?;
                self.define(*out, v);
            }
            Instr::Neg { a, out } => {
                let dims = self.out_dims(*out);
                let av = self.operand(a)?;
                let av = self.to_shape(&av, &dims)?;
                // `vec_negate`, i.e. `-x`, NOT `0 - x` (they differ on ±0).
                let v = self.wrap(av.neg(), "neg")?;
                self.define(*out, v);
            }
            Instr::Select { cond, a, b, out } => {
                let dims = self.out_dims(*out);
                let cv = self.operand(cond)?;
                let av = self.operand(a)?;
                let bv = self.operand(b)?;
                let cv = self.to_shape(&cv, &dims)?;
                let av = self.to_shape(&av, &dims)?;
                let bv = self.to_shape(&bv, &dims)?;
                let z = self.splat_like(&cv, 0.0)?;
                let pred = self.wrap(cv.ne(&z), "select: cond != 0")?;
                let v = self.wrap(pred.select(&av, &bv), "select")?;
                self.define(*out, v);
            }
            Instr::Gather { src, plan, out } => {
                let s = self.src(src)?;
                let plan = &self.prog.plans[*plan as usize];
                let have = self.dims(&s)?;
                if have != plan.src_shape.as_slice() {
                    return Err(self.err(format!(
                        "gather source has shape {have:?} but the plan expects {:?}",
                        &plan.src_shape[..]
                    )));
                }
                let v = self.emit_gather(plan, &s)?;
                self.define(*out, v);
            }
            Instr::LoadElem { src, idx, out } => {
                let s = self.src(src)?;
                let mut cur = s;
                for (d, &i) in idx.iter().enumerate() {
                    cur = self.wrap(
                        cur.slice_in_dim(i as i64, i as i64 + 1, 1, d as i64),
                        "load_elem: slice",
                    )?;
                }
                let v = self.wrap(cur.reshape(&[]), "load_elem: reshape")?;
                self.define(*out, v);
            }
            Instr::Ramp { axis, lo, out } => {
                let dims = self.out_dims(*out);
                if dims.is_empty() {
                    return Err(self.err("Ramp into a scalar slot"));
                }
                let d: Vec<i64> = dims.iter().map(|&x| x as i64).collect();
                // Integer ramp then one convert, so the value is exactly
                // `(lo + p) as f64` rather than an f64 add.
                let io = self.wrap(
                    self.b.iota(ElementType::S64, &d, *axis as i64),
                    "ramp: iota",
                )?;
                let loc = self.ci(*lo)?;
                let loc = self.wrap(loc.broadcast(&d), "ramp: broadcast lo")?;
                let v = self.wrap(
                    io.add_(&loc).and_then(|s| s.convert(PrimitiveType::F64)),
                    "ramp: add + convert",
                )?;
                self.define(*out, v);
            }
            Instr::Fill { v, out } => {
                let dims = self.out_dims(*out);
                let vv = self.operand(v)?;
                if !self.dims(&vv)?.is_empty() {
                    return Err(self.err("Fill with an array operand"));
                }
                let vv = self.to_shape(&vv, &dims)?;
                self.define(*out, vv);
            }
            Instr::Copy { a, out } => {
                let dims = self.out_dims(*out);
                let av = self.operand(a)?;
                let av = self.to_shape(&av, &dims)?;
                self.define(*out, av);
            }
            Instr::Region {
                base,
                src,
                region,
                out,
            } => {
                let spec = &self.prog.regions[*region as usize];
                let basev = self.slots[*base as usize]
                    .clone()
                    .ok_or_else(|| self.err("Region base slot is undefined"))?;
                let sv = self.operand(src)?;
                let upd = self.to_shape(&sv, &spec.shape)?;
                let starts: Vec<XlaOp> = spec
                    .dest_lo
                    .iter()
                    .map(|&x| self.ci(x as i64))
                    .collect::<R<Vec<_>>>()?;
                let v = self.wrap(
                    basev.dynamic_update_slice(&upd, &starts),
                    "region: dynamic_update_slice",
                )?;
                self.define(*out, v);
            }
            Instr::Assemble { table, out } => {
                let dims = self.out_dims(*out);
                let mut cur = self.zeros(&dims)?;
                for (src, region) in &self.prog.assemblies[*table as usize].parts {
                    let spec = &self.prog.regions[*region as usize];
                    let sv = self.operand(src)?;
                    let upd = self.to_shape(&sv, &spec.shape)?;
                    let starts: Vec<XlaOp> = spec
                        .dest_lo
                        .iter()
                        .map(|&x| self.ci(x as i64))
                        .collect::<R<Vec<_>>>()?;
                    cur = self.wrap(
                        cur.dynamic_update_slice(&upd, &starts),
                        "assemble: dynamic_update_slice",
                    )?;
                }
                self.define(*out, cur);
            }
            // An export publishes an observed for the passes the tape serves
            // outside the right-hand side; the emitted program's only output
            // is `du`, and every in-program read is by slot.
            Instr::Export { .. } => {}
            Instr::DyWrite { write } => {
                let w = self.prog.dy_writes[*write as usize].clone();
                self.emit_dy_write(&w)?;
            }
            Instr::ConstArray { data, out } => {
                let payload = &self.prog.const_data[*data as usize];
                let dims = self.out_dims(*out);
                if payload.shape.as_slice() != dims.as_slice() {
                    return Err(self.err(format!(
                        "ConstArray literal has shape {:?} but its slot's box is {dims:?}",
                        &payload.shape[..]
                    )));
                }
                // Row-major elements, which is the layout every slot uses, so
                // the literal goes in as a flat r1 constant and is reshaped.
                let flat = self.wrap(
                    self.b.constant_r1(&payload.values[..]),
                    "const array: literal",
                )?;
                let v = if dims.is_empty() {
                    self.wrap(flat.reshape(&[]), "const array: reshape to scalar")?
                } else {
                    let d: Vec<i64> = dims.iter().map(|&x| x as i64).collect();
                    self.wrap(flat.reshape(&d), "const array: reshape")?
                };
                self.define(*out, v);
            }
            Instr::LoadForcing { forcing, out } => {
                let fr = &self.prog.forcings[*forcing as usize];
                let off = self.forcing_offsets[*forcing as usize]
                    .ok_or_else(|| self.err("forcing load with no place in `f`"))?;
                let n = super::exec::forcing_len(fr);
                let f = self
                    .forcing_param
                    .clone()
                    .ok_or_else(|| self.err("forcing load in a program without `f`"))?;
                let flat = self.wrap(
                    f.slice_in_dim(off as i64, (off + n) as i64, 1, 0),
                    "forcing: slice",
                )?;
                let dims = self.out_dims(*out);
                let d: Vec<i64> = dims.iter().map(|&x| x as i64).collect();
                let v = self.wrap(flat.reshape(&d), "forcing: reshape")?;
                self.define(*out, v);
            }
            // `reverse` fixes a fold order, and XLA's `reduce` pins none (module
            // docs): either direction is measured against the tolerance class.
            Instr::Reduce {
                op,
                init,
                src,
                axes,
                reverse: _,
                src_shape,
                out,
            } => {
                let s = self.src(src)?;
                let have = self.dims(&s)?;
                if have != src_shape.as_slice() {
                    return Err(self.err(format!(
                        "reduce source has shape {have:?} but the instruction expects {:?}",
                        &src_shape[..]
                    )));
                }
                let comp = self.reduce_computation(*op)?;
                let init_op = self.c(*init)?;
                let dims: Vec<i64> = axes.iter().map(|&a| a as i64).collect();
                let v = self.wrap(s.reduce(init_op, comp, &dims, false), "reduce")?;
                // `keep_dims = false` drops exactly `axes`, which is how the
                // tape defines the output box; check rather than trust.
                let got = self.dims(&v)?;
                let want = self.out_dims(*out);
                if got != want {
                    return Err(self.err(format!(
                        "reduce produced shape {got:?} but its slot's box is {want:?}"
                    )));
                }
                self.define(*out, v);
            }
            Instr::Scan {
                op,
                init,
                src,
                axis,
                inclusive,
                reverse,
                src_shape,
                out,
            } => {
                let s = self.src(src)?;
                let have = self.dims(&s)?;
                if have != src_shape.as_slice() {
                    return Err(self.err(format!(
                        "scan source has shape {have:?} but the instruction expects {:?}",
                        &src_shape[..]
                    )));
                }
                let v = self.emit_scan_loop(
                    *op,
                    *init,
                    &s,
                    &have,
                    *axis as usize,
                    *inclusive,
                    *reverse,
                )?;
                self.define(*out, v);
            }
            Instr::TableGather { src, table, out } => {
                let s = self.src(src)?;
                let tbl = &self.prog.gather_tables[*table as usize];
                let have = self.dims(&s)?;
                if have != tbl.src_shape.as_slice() {
                    return Err(self.err(format!(
                        "table gather source has shape {have:?} but its table expects {:?}",
                        &tbl.src_shape[..]
                    )));
                }
                let v = self.emit_table_gather(&s, &have, &tbl.pos)?;
                self.define(*out, v);
            }
            Instr::SegReduce {
                op,
                init,
                src,
                mask,
                table,
                out,
            } => {
                let terms = self.operand(&Operand::Slot(*src))?;
                let mask = match mask {
                    Some(m) => Some(self.operand(&Operand::Slot(*m))?),
                    None => None,
                };
                let rows = &self.prog.seg_tables[*table as usize].rows;
                let v = self.emit_seg_reduce(*op, *init, &terms, mask.as_ref(), rows)?;
                let want = self.out_dims(*out);
                let d: Vec<i64> = want.iter().map(|&x| x as i64).collect();
                let v = self.wrap(v.reshape(&d), "segmented reduce: output box")?;
                self.define(*out, v);
            }
            Instr::Calendar { func, .. } => {
                return Err(self.err(format!(
                    "`{}` in a Float32 document has no XLA lowering (the calendar is \
                     evaluated by the closed-function registry in binary64)",
                    CALENDAR_FNS[*func as usize]
                )));
            }
            Instr::Interp { table, x, y, out } => {
                let tbl = &self.prog.interp_tables[*table as usize];
                if let Some(id) = &tbl.strict {
                    // Like `Instr::Fault`: a computation has no channel to
                    // raise the out-of-range error through, and returning the
                    // clamp or a bare `NaN` without it is the wrong answer
                    // §9.5.3a forbids. Refused by name instead.
                    return Err(self.err(format!(
                        "a lookup into table `{id}`, which declares `out_of_bounds: \"error\"` \
                         (esm-spec §9.5.1): a compiled program cannot raise \
                         `table_lookup_out_of_bounds` at run time"
                    )));
                }
                let dims = self.out_dims(*out);
                let xv = self.operand(x)?;
                let xv = self.to_shape(&xv, &dims)?;
                let yv = match y {
                    Some(y) => {
                        let v = self.operand(y)?;
                        Some(self.to_shape(&v, &dims)?)
                    }
                    None => None,
                };
                let v = self.emit_interp(tbl, &dims, &xv, yv.as_ref())?;
                self.define(*out, v);
            }
            Instr::PolyArea { a, b, geom, out } => {
                let v = self.emit_poly_area(a, b, *geom, *out)?;
                self.define(*out, v);
            }
            Instr::IndexGather {
                src,
                idx,
                spec,
                out,
            } => {
                let v = self.emit_index_gather(src, idx, *spec)?;
                self.define(*out, v);
            }
            Instr::Reshape { src, out } => {
                let s = self.src(src)?;
                let dims = self.out_dims(*out);
                let have = self.dims(&s)?;
                if have.iter().product::<usize>() != dims.iter().product::<usize>() {
                    return Err(self.err(format!(
                        "reshape source has shape {have:?}, which does not hold the \
                         {dims:?} its slot's box needs"
                    )));
                }
                // XLA's reshape is row-major, which is the tape's definition.
                let d: Vec<i64> = dims.iter().map(|&x| x as i64).collect();
                let v = self.wrap(s.reshape(&d), "reshape")?;
                self.define(*out, v);
            }
            Instr::Fault { fault } => {
                // A computation has no channel to latch a fault through: its
                // one output is `du`. The model is refused with the fault it
                // would raise rather than compiled to return the `NaN` the
                // interpreter substitutes WITHOUT the error that goes with it.
                return Err(self.err(format!(
                    "the program raises a fail-closed evaluation fault, which a compiled \
                     program cannot report: {}",
                    self.prog.faults[*fault as usize]
                )));
            }
            Instr::FaultScan { .. } => {
                // As for `Instr::Fault`: the program may raise a fail-closed
                // fault, depending on which cells it evaluates, and has no
                // channel to report it through.
                return Err(self.err(
                    "the program may raise a fail-closed evaluation fault (a read the \
                     interpreter evaluates only where a condition holds), which a compiled \
                     program cannot report",
                ));
            }
            Instr::JmpIfZero { .. } => {
                return Err(self.err("JmpIfZero reached emit_one (handled by emit_range)"));
            }
            Instr::Sweep { .. } => return Err(self.err(RECURRENCE_REFUSAL)),
            Instr::DataGather { faults, .. } => {
                return Err(self.err(if *faults {
                    "a computed-subscript const-array gather (`DataGather`) can raise a \
                     fail-closed fault, which a compiled program cannot report"
                } else {
                    "a computed-subscript gather (`DataGather`) has no XLA lowering yet"
                }));
            }
            Instr::ScalarRead { .. } => {
                // Only a recurrence body holds one, and the sweep is refused
                // before its body is reached; this arm keeps the refusal named
                // should another construct start using it.
                return Err(self.err(
                    "a run-time-subscript read (`ScalarRead`) can raise a fail-closed \
                     fault, which a compiled program cannot report",
                ));
            }
            Instr::Fused { .. } => {
                return Err(self.err(
                    "Instr::Fused reached the emitter; the tape must be built with the \
                     fusion pass OFF for emission (build_tape_opts(.., None)) — a fused \
                     group is a scalar micro-program, which is the shape XLA's own \
                     fusion chooses for itself",
                ));
            }
            // ---- WHERE NEW TAPE INSTRUCTIONS PLUG IN --------------------
            // `ConstArray` and `Reduce` arrived this way and are lowered
            // above. Anything this emitter has not learned lands here and is
            // REFUSED by name, so the gap shows up as a named refusal in the
            // conformance report instead of as wrong numbers. Add the arm
            // above, not a special case here.
            #[allow(unreachable_patterns)]
            other => {
                return Err(self.err(format!(
                    "tape instruction `{}` has no lowering: the tape has grown and the \
                     XLA emitter must follow",
                    other.opcode()
                )));
            }
        }
        Ok(())
    }

    fn define(&mut self, slot: SlotId, v: XlaOp) {
        self.slots[slot as usize] = Some(v);
    }

    // -- interp.* (esm-spec §9.2) ------------------------------------------

    /// Lower one [`Instr::Interp`]: the §9.2 blend over a compile-time table.
    ///
    /// The table and axes go in as `r1` constants and the cell index comes out
    /// of a COMPARE-AND-REDUCE — `count(axis[k] <= q)` over the lower-corner
    /// candidates — which is the branch-free form of the registry's scan and
    /// the shape XLA vectorizes. The corners are then read with `take`
    /// (a `gather` along the table's one axis), and the blend is written in
    /// §9.2's pinned `a + w * (b - a)` order.
    ///
    /// NUMERICAL STANDING, like every other arm here: equal to the
    /// interpreter within the tier's tolerance classes, not bit for bit. The
    /// index arithmetic and the clamps are exact, and the blend is the same
    /// three operations in the same order — but nothing stops XLA contracting
    /// `t_lo + w * (t_hi - t_lo)` into an FMA, which §9.2 explicitly allows a
    /// binding to do and pins a relaxed tolerance for.
    ///
    /// COST: the index reduce materializes one predicate per output element
    /// per axis knot. That is the same work the registry's scan does, laid out
    /// wide instead of sequentially; a binary search over the constant axis
    /// would be `log N` selects instead, and is the optimization to reach for
    /// if a wide box over a long axis ever shows up.
    fn emit_interp(
        &self,
        tbl: &InterpTable,
        dims: &[usize],
        x: &XlaOp,
        y: Option<&XlaOp>,
    ) -> R<XlaOp> {
        match tbl.kind {
            InterpKind::Linear => {
                let blend = self.interp_blend(&tbl.table, &tbl.axis_x, dims, x)?;
                // The two extrapolate-flat clamps, in the registry's order
                // (below first). They cannot both hold: a validated axis is
                // strictly increasing.
                let n = tbl.axis_x.len();
                let lo_knot = self.splat(tbl.axis_x[0], dims)?;
                let hi_knot = self.splat(tbl.axis_x[n - 1], dims)?;
                let lo_val = self.splat(tbl.table[0], dims)?;
                let hi_val = self.splat(tbl.table[n - 1], dims)?;
                let above = self.wrap(x.ge(&hi_knot), "interp.linear: x >= axis[N]")?;
                let inner =
                    self.wrap(above.select(&hi_val, &blend), "interp.linear: high clamp")?;
                let below = self.wrap(x.le(&lo_knot), "interp.linear: x <= axis[1]")?;
                self.wrap(below.select(&lo_val, &inner), "interp.linear: low clamp")
            }
            InterpKind::Bilinear => {
                let y = y.ok_or_else(|| self.err("interp.bilinear without a `y` operand"))?;
                let ny = tbl.axis_y.len();
                // Per-axis clamp of the QUERY (not of the result), which is
                // what makes the cell search below well defined.
                let xq = self.interp_clamp_query(&tbl.axis_x, dims, x)?;
                let yq = self.interp_clamp_query(&tbl.axis_y, dims, y)?;
                let i = self.interp_cell(&tbl.axis_x, dims, &xq)?;
                let j = self.interp_cell(&tbl.axis_y, dims, &yq)?;
                let ax_lo = self.interp_take_at(&tbl.axis_x, &i)?;
                let ax_hi = self.interp_take_next(&tbl.axis_x, &i)?;
                let ay_lo = self.interp_take_at(&tbl.axis_y, &j)?;
                let ay_hi = self.interp_take_next(&tbl.axis_y, &j)?;
                let wx = self.interp_weight(&xq, &ax_lo, &ax_hi)?;
                let wy = self.interp_weight(&yq, &ay_lo, &ay_hi)?;
                // The flat row-major address of each of the four corners,
                // `i * ny + j`, so all four are one `take` from the same r1
                // constant the tape already holds flat.
                let nyc = self.ci32(ny as i32, dims)?;
                let one = self.ci32(1, dims)?;
                let irow = self.wrap(i.mul_(&nyc), "interp.bilinear: i * ny")?;
                let irow_n = self.wrap(irow.add_(&nyc), "interp.bilinear: (i+1) * ny")?;
                let jp1 = self.wrap(j.add_(&one), "interp.bilinear: j + 1")?;
                let c00 = self.interp_take_at(&tbl.table, &self.wrap(irow.add_(&j), "c00")?)?;
                let c10 = self.interp_take_at(&tbl.table, &self.wrap(irow_n.add_(&j), "c10")?)?;
                let c01 = self.interp_take_at(&tbl.table, &self.wrap(irow.add_(&jp1), "c01")?)?;
                let c11 = self.interp_take_at(&tbl.table, &self.wrap(irow_n.add_(&jp1), "c11")?)?;
                let row_j = self.interp_lerp(&c00, &c10, &wx)?;
                let row_jp1 = self.interp_lerp(&c01, &c11, &wx)?;
                self.interp_lerp(&row_j, &row_jp1, &wy)
            }
            InterpKind::SearchSorted => {
                // The 1-based index of the first entry >= x. With `xs`
                // non-decreasing the entries BELOW x are a prefix, so that
                // index is `count(xs[k] < x) + 1` — and `count` is 0 for a NaN
                // query, where the registry answers N+1 instead.
                let xs = &tbl.axis_x;
                let n = xs.len();
                let below = self.interp_count(xs, dims, x, CountCmp::AxisLtQuery)?;
                let one = self.ci32(1, dims)?;
                let idx = self.wrap(below.add_(&one), "searchsorted: count + 1")?;
                let idx = self.wrap(
                    idx.convert(PrimitiveType::F64),
                    "searchsorted: index -> f64",
                )?;
                if n == 0 {
                    // Degenerate table: the registry returns 1 for every
                    // query, NaN included.
                    return Ok(idx);
                }
                let past_end = self.splat((n as i64 + 1) as f64, dims)?;
                let is_nan = self.wrap(x.ne(x), "searchsorted: x is NaN")?;
                self.wrap(is_nan.select(&past_end, &idx), "searchsorted: NaN query")
            }
        }
    }

    /// The in-cell part of `interp.linear` (no clamps): locate the cell, read
    /// its two corners, blend.
    fn interp_blend(&self, table: &[f64], axis: &[f64], dims: &[usize], q: &XlaOp) -> R<XlaOp> {
        let i = self.interp_cell(axis, dims, q)?;
        let a_lo = self.interp_take_at(axis, &i)?;
        let a_hi = self.interp_take_next(axis, &i)?;
        let t_lo = self.interp_take_at(table, &i)?;
        let t_hi = self.interp_take_next(table, &i)?;
        let w = self.interp_weight(q, &a_lo, &a_hi)?;
        self.interp_lerp(&t_lo, &t_hi, &w)
    }

    /// `axis[0] <= q <= axis[N-1]`, extrapolate-flat (`interp.bilinear`
    /// step 1). A NaN query falls through both comparisons and stays NaN,
    /// exactly as the registry's `if / else if` chain leaves it.
    fn interp_clamp_query(&self, axis: &[f64], dims: &[usize], q: &XlaOp) -> R<XlaOp> {
        let n = axis.len();
        let lo = self.splat(axis[0], dims)?;
        let hi = self.splat(axis[n - 1], dims)?;
        let above = self.wrap(q.ge(&hi), "interp: q >= axis[N]")?;
        let inner = self.wrap(above.select(&hi, q), "interp: clamp high")?;
        let below = self.wrap(q.le(&lo), "interp: q <= axis[1]")?;
        self.wrap(below.select(&lo, &inner), "interp: clamp low")
    }

    /// The 0-based lower-corner index: `count(axis[k] <= q)` over
    /// `k in [0, N-2]`, minus one, clamped into `[0, N-2]`.
    ///
    /// The clamp is what reproduces the registry's NaN behaviour: every
    /// comparison is false for NaN, so the count is zero and the index would
    /// be -1, where both the 1-D scan and `locate_cell` leave `i = 0` and let
    /// the NaN weight carry the result.
    fn interp_cell(&self, axis: &[f64], dims: &[usize], q: &XlaOp) -> R<XlaOp> {
        let n = axis.len();
        let count = self.interp_count(&axis[..n - 1], dims, q, CountCmp::AxisLeQuery)?;
        let one = self.ci32(1, dims)?;
        let cand = self.wrap(count.sub_(&one), "interp: count - 1")?;
        let lo = self.ci32(0, dims)?;
        let hi = self.ci32(n as i32 - 2, dims)?;
        self.wrap(cand.clamp(&lo, &hi), "interp: clamp cell index")
    }

    /// `count(cmp(knots[k], q))` per output element, as an `s32` of `dims`.
    fn interp_count(&self, knots: &[f64], dims: &[usize], q: &XlaOp, cmp: CountCmp) -> R<XlaOp> {
        if knots.is_empty() {
            return self.ci32(0, dims);
        }
        // `q` broadcast over a trailing knot axis, against the knot vector
        // broadcast over the output box.
        let mut wide: Vec<i64> = dims.iter().map(|&d| d as i64).collect();
        wide.push(knots.len() as i64);
        let keep: Vec<i64> = (0..dims.len() as i64).collect();
        let qb = self.wrap(q.broadcast_in_dim(&wide, &keep), "interp: broadcast query")?;
        let kv = self.wrap(self.b.constant_r1(knots), "interp: knot vector")?;
        let kb = self.wrap(
            kv.broadcast_in_dim(&wide, &[dims.len() as i64]),
            "interp: broadcast knots",
        )?;
        let pred = match cmp {
            CountCmp::AxisLeQuery => self.wrap(kb.le(&qb), "interp: axis <= q")?,
            CountCmp::AxisLtQuery => self.wrap(kb.lt(&qb), "interp: axis < q")?,
        };
        let ones = self.wrap(pred.convert(PrimitiveType::S32), "interp: pred -> s32")?;
        self.wrap(
            ones.reduce_sum(&[dims.len() as i64], false),
            "interp: count knots",
        )
    }

    /// `values[idx]`, with `idx` an `s32` tensor over the output box.
    fn interp_take_at(&self, values: &[f64], idx: &XlaOp) -> R<XlaOp> {
        let v = self.wrap(self.b.constant_r1(values), "interp: table constant")?;
        self.wrap(v.take(idx, 0), "interp: table gather")
    }

    /// `values[idx + 1]` — the upper corner of the located cell.
    fn interp_take_next(&self, values: &[f64], idx: &XlaOp) -> R<XlaOp> {
        let dims = self.dims(idx)?;
        let one = self.ci32(1, &dims)?;
        let next = self.wrap(idx.add_(&one), "interp: idx + 1")?;
        self.interp_take_at(values, &next)
    }

    /// `w = (q - lo) / (hi - lo)`.
    fn interp_weight(&self, q: &XlaOp, lo: &XlaOp, hi: &XlaOp) -> R<XlaOp> {
        let num = self.wrap(q.sub_(lo), "interp: q - axis[i]")?;
        let den = self.wrap(hi.sub_(lo), "interp: axis[i+1] - axis[i]")?;
        self.wrap(num.div_(&den), "interp: weight")
    }

    /// §9.2's pinned blend `a + w * (b - a)`.
    fn interp_lerp(&self, a: &XlaOp, b: &XlaOp, w: &XlaOp) -> R<XlaOp> {
        let d = self.wrap(b.sub_(a), "interp: b - a")?;
        let s = self.wrap(w.mul_(&d), "interp: w * (b - a)")?;
        self.wrap(a.add_(&s), "interp: a + w * (b - a)")
    }

    /// An `f64` constant with the output box's shape.
    fn splat(&self, v: f64, dims: &[usize]) -> R<XlaOp> {
        let c = self.c(v)?;
        self.to_shape(&c, dims)
    }

    /// An `s32` constant with the output box's shape (the index arithmetic
    /// above stays integral so the cell address is exact).
    fn ci32(&self, v: i32, dims: &[usize]) -> R<XlaOp> {
        let c = self
            .b
            .c0(v)
            .map_err(|e| self.err(format!("s32 constant: {e}")))?;
        if dims.is_empty() {
            return Ok(c);
        }
        let d: Vec<i64> = dims.iter().map(|&x| x as i64).collect();
        self.wrap(c.broadcast(&d), "interp: broadcast s32 constant")
    }

    // -- geometry ----------------------------------------------------------

    /// Lower one [`Instr::PolyArea`] by evaluating it HERE, at emit time,
    /// into a constant.
    ///
    /// A polygon clip has no expression in XLA's operation set, so the
    /// emitter can take the instruction only when both ring tables are
    /// compile-time constants: slots an [`Instr::ConstArray`] defines. It then
    /// runs the slab executor's own [`super::geom::run_poly_area`] over those
    /// literals (broad phase included) and emits the areas as a literal, so
    /// the compiled program reads the same `f64`s the slab executor computes.
    /// A table that depends on a parameter, the state or an observed is a
    /// named refusal.
    fn emit_poly_area(&self, a: &SrcRef, b: &SrcRef, geom: u32, out: SlotId) -> R<XlaOp> {
        let spec = &self.prog.geoms[geom as usize];
        let ta = self.const_table(a)?;
        let tb = self.const_table(b)?;
        let (sa, sb) = (rm_strides(&ta.shape), rm_strides(&tb.shape));
        let dims = self.out_dims(out);
        let mut areas = vec![0.0f64; dims.iter().product::<usize>().max(1)];
        unsafe {
            super::geom::run_poly_area(
                spec,
                &super::geom::RingTable {
                    ptr: ta.values.as_ptr(),
                    shape: &ta.shape,
                    strides: &sa,
                },
                &super::geom::RingTable {
                    ptr: tb.values.as_ptr(),
                    shape: &tb.shape,
                    strides: &sb,
                },
                areas.as_mut_ptr(),
            );
        }
        let flat = self.wrap(self.b.constant_r1(&areas[..]), "polygon area: literal")?;
        let d: Vec<i64> = dims.iter().map(|&x| x as i64).collect();
        self.wrap(flat.reshape(&d), "polygon area: reshape")
    }

    /// The literal an [`Instr::ConstArray`] stores into ring-table slot `s`,
    /// or the refusal saying the table is not a compile-time constant.
    fn const_table(&self, s: &SrcRef) -> R<&ConstArrayData> {
        let not_const = || {
            self.err(
                "polygon_intersection_area: a ring table that is not a compile-time \
                 constant (it depends on a parameter, the state or an observed) has no \
                 lowering: XLA has no polygon clip",
            )
        };
        let SrcRef::Slot(id) = s else {
            return Err(not_const());
        };
        self.prog
            .instrs
            .iter()
            .find_map(|i| match i {
                Instr::ConstArray { data, out } if out == id => {
                    Some(&self.prog.const_data[*data as usize])
                }
                _ => None,
            })
            .ok_or_else(not_const)
    }

    /// Lower one [`Instr::IndexGather`]: the subscript is rounded and
    /// range-checked exactly as [`data_subscript`] does it (NaN reads as 0,
    /// an out-of-range position as the zero ghost), each output element's
    /// row-major source offset is built in `s64`, and one `take` over the
    /// flattened source reads them.
    fn emit_index_gather(&self, src: &SrcRef, idx: &Operand, spec: u32) -> R<XlaOp> {
        let spec = &self.prog.index_gathers[spec as usize];
        let s = self.src(src)?;
        let have = self.dims(&s)?;
        if have != spec.src_shape.as_slice() {
            return Err(self.err(format!(
                "indexed gather source has shape {have:?} but the spec expects {:?}",
                &spec.src_shape[..]
            )));
        }
        let dims: Vec<usize> = spec.shape.to_vec();
        let d64: Vec<i64> = dims.iter().map(|&x| x as i64).collect();
        let total: i64 = spec.src_shape.iter().product::<usize>() as i64;
        let flat = self.wrap(s.reshape(&[total]), "indexed gather: flatten source")?;
        let strides = rm_strides(&spec.src_shape);
        let data_d = spec.data_axis();
        let n = spec.src_shape[data_d] as f64;

        let iv = self.operand(idx)?;
        let iv = self.to_shape(&iv, &dims)?;
        let r = self.wrap(iv.round(), "indexed gather: round")?;
        let nan = self.wrap(r.ne(&r), "indexed gather: isnan")?;
        let zero = self.splat(0.0, &dims)?;
        let r = self.wrap(nan.select(&zero, &r), "indexed gather: NaN reads as 0")?;
        // Pinned to [0, n + 1] before the integer convert, so every value
        // converts exactly and the out-of-range ones stay out of range.
        let top = self.splat(n + 1.0, &dims)?;
        let r = self.wrap(
            r.max(&zero).and_then(|x| x.min(&top)),
            "indexed gather: pin",
        )?;
        let one_based = self.wrap(r.convert(PrimitiveType::S64), "indexed gather: convert")?;
        let splat_i = |v: i64| -> R<XlaOp> {
            let c = self.ci(v)?;
            self.to_shape(&c, &dims)
        };
        let (one, zero_i, n_i, last_i) = (
            splat_i(1)?,
            splat_i(0)?,
            splat_i(n as i64)?,
            splat_i(n as i64 - 1)?,
        );
        let valid = self.wrap(
            one_based
                .ge(&one)
                .and_then(|lo| one_based.le(&n_i).and_then(|hi| lo.and(&hi))),
            "indexed gather: in range",
        )?;
        let pos = self.wrap(
            one_based
                .sub_(&one)
                .and_then(|p| p.max(&zero_i))
                .and_then(|p| p.min(&last_i)),
            "indexed gather: position",
        )?;
        let stride_data = splat_i(strides[data_d])?;
        let mut off = self.wrap(pos.mul_(&stride_data), "indexed gather: data offset")?;
        for (d, ax) in spec.axes.iter().enumerate() {
            let term = match *ax {
                GatherAxis::Data => continue,
                GatherAxis::Fixed(i) => splat_i(i as i64 * strides[d])?,
                GatherAxis::Affine { axis, off: k } => {
                    let io = self.wrap(
                        self.b.iota(ElementType::S64, &d64, axis as i64),
                        "indexed gather: iota",
                    )?;
                    let (k, st) = (splat_i(k)?, splat_i(strides[d])?);
                    self.wrap(
                        io.add_(&k).and_then(|x| x.mul_(&st)),
                        "indexed gather: affine offset",
                    )?
                }
            };
            off = self.wrap(off.add_(&term), "indexed gather: offset")?;
        }
        let g = self.wrap(flat.take(&off, 0), "indexed gather: take")?;
        self.wrap(valid.select(&g, &zero), "indexed gather: ghost")
    }

    // -- gather ------------------------------------------------------------

    /// Lower one precompiled [`GatherPlan`] — a transliteration of
    /// `refexec::exec_gather`'s four stages into XLA:
    ///
    /// 1. fixed source axes -> `slice_in_dim` + `reshape` (drop the axis);
    /// 2. the mapped-axis permutation -> `transpose`;
    /// 3. broadcast axes -> one `broadcast_in_dim`;
    /// 4. the per-axis copy-segment schedule -> nested `slice_in_dim` +
    ///    `concat_in_dim`, with a zero constant for every uncovered (ghost)
    ///    stretch. The schedule is a CARTESIAN product of per-axis segments,
    ///    which is exactly what that nesting expresses; contiguous plans (the
    ///    common shifted-stencil read) collapse to a single slice with no
    ///    concatenation at all.
    ///
    /// The alternative — an XLA `gather` over an explicit index array — would
    /// materialize one s64 index per output element for what is, in every plan
    /// the corpus produces, pure affine data movement.
    fn emit_gather(&self, plan: &GatherPlan, src: &XlaOp) -> R<XlaOp> {
        let mut cur = src.clone();
        // 1. fixed axes, descending so lower axis numbers stay valid.
        for &(d, i0) in &plan.fixed_desc {
            let dims = self.dims(&cur)?;
            if d >= dims.len() {
                return Err(self.err("gather plan fixes an axis the source does not have"));
            }
            let mut nd: Vec<i64> = dims.iter().map(|&x| x as i64).collect();
            nd.remove(d);
            cur = self.wrap(
                cur.slice_in_dim(i0 as i64, i0 as i64 + 1, 1, d as i64)
                    .and_then(|s| s.reshape(&nd)),
                "gather: fixed axis",
            )?;
        }
        // 2. permutation of the remaining source axes.
        if plan.perm.len() > 1 && !plan.perm.iter().enumerate().all(|(i, &p)| i == p) {
            let perm: Vec<i64> = plan.perm.iter().map(|&x| x as i64).collect();
            cur = self.wrap(cur.transpose(&perm), "gather: transpose")?;
        }
        // 3. broadcast axes.
        let out_ndim = plan.shape.len();
        let mapped_axes: Vec<i64> = (0..out_ndim)
            .filter(|&a| plan.mapped[a])
            .map(|a| a as i64)
            .collect();
        let cur_dims = self.dims(&cur)?;
        if cur_dims.len() != mapped_axes.len() {
            return Err(self.err(format!(
                "gather plan maps {} source axes but {} output axes are marked mapped",
                cur_dims.len(),
                mapped_axes.len()
            )));
        }
        if out_ndim > 0 {
            let mut bshape: Vec<i64> = Vec::with_capacity(out_ndim);
            let mut k = 0usize;
            for a in 0..out_ndim {
                if plan.mapped[a] {
                    bshape.push(cur_dims[k] as i64);
                    k += 1;
                } else {
                    bshape.push(plan.shape[a] as i64);
                }
            }
            if bshape.len() != cur_dims.len()
                || !mapped_axes
                    .iter()
                    .enumerate()
                    .all(|(i, &a)| a as usize == i)
            {
                cur = self.wrap(
                    cur.broadcast_in_dim(&bshape, &mapped_axes),
                    "gather: broadcast_in_dim",
                )?;
            }
        }
        // 4. the copy-segment schedule.
        self.emit_segments(&cur, plan, 0)
    }

    fn emit_segments(&self, cur: &XlaOp, plan: &GatherPlan, d: usize) -> R<XlaOp> {
        if d == plan.shape.len() {
            return Ok(cur.clone());
        }
        let extent = plan.shape[d];
        let mut segs: Vec<(usize, usize, usize)> = plan.segs[d].to_vec();
        segs.sort_by_key(|s| s.0);
        // Fast path: one segment covering the whole axis with no offset.
        if segs.len() == 1 && segs[0] == (0, extent, 0) {
            let cd = self.dims(cur)?;
            if cd[d] == extent {
                return self.emit_segments(cur, plan, d + 1);
            }
        }
        let cd = self.dims(cur)?;
        let piece_zeros = |len: usize| -> R<XlaOp> {
            let mut shape: Vec<usize> = cd.clone();
            shape[d] = len;
            for (e, s) in shape.iter_mut().enumerate().skip(d + 1) {
                *s = plan.shape[e];
            }
            self.zeros(&shape)
        };
        let mut pieces: Vec<XlaOp> = Vec::new();
        let mut next = 0usize;
        for (o, l, s) in segs {
            if o < next {
                return Err(self.err("gather plan segments overlap along an axis"));
            }
            if o > next {
                pieces.push(piece_zeros(o - next)?);
            }
            let sl = self.wrap(
                cur.slice_in_dim(s as i64, (s + l) as i64, 1, d as i64),
                "gather: segment slice",
            )?;
            pieces.push(self.emit_segments(&sl, plan, d + 1)?);
            next = o + l;
        }
        if next > extent {
            return Err(self.err("gather plan segments run past the output extent"));
        }
        if next < extent {
            pieces.push(piece_zeros(extent - next)?);
        }
        match pieces.len() {
            0 => Err(self.err("gather plan has no segments and no extent along an axis")),
            1 => Ok(pieces.remove(0)),
            _ => {
                let head = pieces.remove(0);
                self.wrap(
                    head.concat_in_dim(&pieces, d as i64),
                    "gather: concat_in_dim",
                )
            }
        }
    }

    // -- dy ----------------------------------------------------------------

    /// Apply one [`DyWrite`] to the accumulator block of its variable.
    ///
    /// `refexec` writes straight into the flat, COLUMN-MAJOR `dy`. Here each
    /// variable keeps a ROW-MAJOR block of its own logical shape and the whole
    /// set is relabelled to column-major once, in [`Self::assemble_du`], so a
    /// write is a plain `dynamic_update_slice` at `dest_lo` rather than a
    /// strided scatter.
    fn emit_dy_write(&mut self, w: &DyWrite) -> R<()> {
        if w.scatter.is_some() {
            return Err(self.err(
                "a state derivative whose left-hand side is not a constant shift of its \
                 output indices (a per-cell dy scatter) has no XLA lowering",
            ));
        }
        // `DyWrite::var` is only meaningful for the ARRAY form. A scalar rule
        // (`RhsRule::Scalar` / `IndexedScalar`) carries `var: 0` as a dummy
        // and addresses `dy` by the absolute flat slot in `scalar_flat`, so
        // its owning variable has to be found from that slot instead — using
        // `var` there would write into whichever variable happens to sit
        // first in the layout.
        let vi = match w.scalar_flat {
            Some(flat) => self.var_owning_slot(flat)?,
            None => {
                let sv = &self.prog.state_vars[w.var as usize];
                *self.var_ix.get(&sv.name).ok_or_else(|| {
                    self.err(format!("dy write names unknown variable {}", sv.name))
                })?
            }
        };
        let (var_name, vs) = &self.var_order[vi];
        let var_name = var_name.clone();
        let shape: Vec<usize> = vs.shape.clone();
        let var_offset = vs.flat_offset;
        let val = self.slots[w.slot as usize]
            .clone()
            .ok_or_else(|| self.err("dy-write slot is undefined"))?;
        let val_dims = self.dims(&val)?;

        let cur = match &self.blocks[vi] {
            Some(b) => b.clone(),
            None => self.zeros(&shape)?,
        };

        let next = match w.scalar_flat {
            Some(flat) => {
                // A scalar rule writing ONE flat slot: decode its
                // column-major position inside the variable's box.
                if !val_dims.is_empty() {
                    return Err(self.err("scalar dy write with an array value"));
                }
                if shape.is_empty() {
                    val
                } else {
                    let mut rem = flat - var_offset;
                    let mut multi = vec![0usize; shape.len()];
                    for d in 0..shape.len() {
                        multi[d] = rem % shape[d];
                        rem /= shape[d];
                    }
                    if rem != 0 {
                        return Err(self.err("scalar dy write lands past its variable block"));
                    }
                    let ones: Vec<i64> = shape.iter().map(|_| 1i64).collect();
                    let upd = self.wrap(val.reshape(&ones), "dy write: reshape scalar")?;
                    let starts: Vec<XlaOp> = multi
                        .iter()
                        .map(|&x| self.ci(x as i64))
                        .collect::<R<Vec<_>>>()?;
                    self.wrap(
                        cur.dynamic_update_slice(&upd, &starts),
                        "dy write: dynamic_update_slice",
                    )?
                }
            }
            None => {
                if shape.is_empty() {
                    if !val_dims.is_empty() {
                        return Err(self.err("array dy write into a 0-d variable"));
                    }
                    val
                } else {
                    if val_dims.len() != shape.len() {
                        return Err(self.err(format!(
                            "dy write value has rank {} but variable {var_name} has rank {}",
                            val_dims.len(),
                            shape.len()
                        )));
                    }
                    let starts: Vec<XlaOp> = w
                        .dest_lo
                        .iter()
                        .map(|&x| self.ci(x as i64))
                        .collect::<R<Vec<_>>>()?;
                    self.wrap(
                        cur.dynamic_update_slice(&val, &starts),
                        "dy write: dynamic_update_slice",
                    )?
                }
            }
        };
        self.blocks[vi] = Some(next);
        Ok(())
    }

    /// The scalar `(a, b) -> kernel(op)(a, b)` sub-computation an
    /// [`Instr::Reduce`] folds with, on its own [`XlaBuilder`] (XLA requires
    /// the reducer to be a separate computation).
    fn reduce_computation(&self, op: BinCode) -> R<XlaComputation> {
        let rb = XlaBuilder::new("reduce_kernel");
        let a = rb
            .parameter(0, f64::TY, &[], "a")
            .map_err(|e| self.err(format!("reduce kernel parameter a: {e}")))?;
        let b = rb
            .parameter(1, f64::TY, &[], "b")
            .map_err(|e| self.err(format!("reduce kernel parameter b: {e}")))?;
        let body = self.bin(op, &a, &b)?;
        body.build()
            .map_err(|e| self.err(format!("reduce kernel build: {e}")))
    }

    /// Lower one [`Instr::TableGather`]: the source flattened row-major, one
    /// `take` at the build-time positions, and the ghost positions selected
    /// to `+0.0`.
    fn emit_table_gather(&self, src: &XlaOp, dims: &[usize], pos: &[u32]) -> R<XlaOp> {
        let n: usize = dims.iter().product::<usize>().max(1);
        let flat = self.wrap(src.reshape(&[n as i64]), "table gather: flatten")?;
        let at: Vec<i64> = pos
            .iter()
            .map(|&p| if p == GATHER_GHOST { 0 } else { p as i64 })
            .collect();
        let at = self.wrap(self.b.constant_r1(&at[..]), "table gather: positions")?;
        let taken = self.wrap(flat.take(&at, 0), "table gather: take")?;
        if !pos.contains(&GATHER_GHOST) {
            return Ok(taken);
        }
        let live: Vec<f64> = pos
            .iter()
            .map(|&p| if p == GATHER_GHOST { 0.0 } else { 1.0 })
            .collect();
        let live = self.wrap(self.b.constant_r1(&live[..]), "table gather: live mask")?;
        let z = self.splat_like(&taken, 0.0)?;
        let keep = self.wrap(live.ne(&z), "table gather: live != 0")?;
        self.wrap(keep.select(&taken, &z), "table gather: ghost")
    }

    /// Lower one [`Instr::SegReduce`] as a scatter-combine of each term into
    /// its output cell, excluded terms first replaced by the identity. Like
    /// [`Instr::Reduce`] this is equal to the interpreter numerically, not
    /// bit for bit: XLA does not order the updates a scatter combines.
    /// Returns the flat `n_out` result.
    fn emit_seg_reduce(
        &self,
        op: BinCode,
        init: f64,
        terms: &XlaOp,
        mask: Option<&XlaOp>,
        rows: &[u32],
    ) -> R<XlaOp> {
        let n_out = rows.len() - 1;
        let m = rows[n_out] as usize;
        let base = {
            let c = self.c(init)?;
            self.wrap(c.broadcast(&[n_out as i64]), "segmented reduce: identity")?
        };
        if m == 0 {
            return Ok(base);
        }
        let terms = match mask {
            None => terms.clone(),
            Some(mk) => {
                let z = self.splat_like(mk, 0.0)?;
                let keep = self.wrap(mk.ne(&z), "segmented reduce: mask != 0")?;
                let id = self.splat_like(terms, init)?;
                self.wrap(keep.select(terms, &id), "segmented reduce: mask")?
            }
        };
        let mut cell: Vec<i64> = Vec::with_capacity(m);
        for c in 0..n_out {
            cell.extend(std::iter::repeat_n(
                c as i64,
                (rows[c + 1] - rows[c]) as usize,
            ));
        }
        let cell = self.wrap(
            self.b
                .constant_r1(&cell[..])
                .and_then(|c| c.reshape(&[m as i64, 1])),
            "segmented reduce: cell ids",
        )?;
        let comp = self.reduce_computation(op)?;
        self.wrap(
            base.scatter(&cell, &terms, &comp, &[], &[0], &[0], 1),
            "segmented reduce: scatter",
        )
    }

    /// Lower one [`Instr::Scan`] as ONE `while` loop over the scanned axis,
    /// so the program's size does not grow with the scan's length.
    ///
    /// The loop carries `(k, src, acc, out)`: step `k` slices plane `k` of
    /// `src`, combines it into the running `acc` (a plane-shaped value), and
    /// writes the inclusive or exclusive partial into plane `k` of `out`.
    /// Planes are folded strictly in ascending order, one elementwise combine
    /// per step, so every position folds its terms in the tape's own
    /// association: this lowering adds no reordering of its own.
    fn emit_scan_loop(
        &self,
        op: BinCode,
        init: f64,
        src: &XlaOp,
        dims: &[usize],
        ax: usize,
        inclusive: bool,
        reverse: bool,
    ) -> R<XlaOp> {
        use xla::Shape;
        let full: Vec<i64> = dims.iter().map(|&x| x as i64).collect();
        let mut plane = full.clone();
        plane[ax] = 1;
        let carried = Shape::tuple(vec![
            Shape::array::<i64>(vec![]),
            Shape::array::<f64>(full.clone()),
            Shape::array::<f64>(plane.clone()),
            Shape::array::<f64>(full.clone()),
        ]);
        let n = full[ax];

        let cb = XlaBuilder::new("scan_cond");
        let st = self.wrap(cb.parameter_s(0, &carried, "st"), "scan cond: parameter")?;
        let k = self.wrap(st.get_tuple_element(0), "scan cond: k")?;
        let bound = self.wrap(cb.c0(n), "scan cond: bound")?;
        let cond = self.wrap(k.lt(&bound), "scan cond: k < n")?;
        let cond = cond
            .build()
            .map_err(|e| self.err(format!("scan cond build: {e}")))?;

        let bb = XlaBuilder::new("scan_body");
        let st = self.wrap(bb.parameter_s(0, &carried, "st"), "scan body: parameter")?;
        let part = |i: i64| self.wrap(st.get_tuple_element(i), "scan body: unpack");
        let (k, s, acc, out) = (part(0)?, part(1)?, part(2)?, part(3)?);
        let zero = self.wrap(bb.c0(0i64), "scan body: zero")?;
        // Step `k` visits plane `k`, or plane `n - 1 - k` for a reverse scan.
        let pos = if reverse {
            let last = self.wrap(bb.c0(n - 1), "scan body: last")?;
            self.wrap(last.sub_(&k), "scan body: n - 1 - k")?
        } else {
            k.clone()
        };
        let starts: Vec<XlaOp> = (0..full.len())
            .map(|d| if d == ax { pos.clone() } else { zero.clone() })
            .collect();
        let pl = self.wrap(s.dynamic_slice(&starts, &plane), "scan body: plane")?;
        let next = self.bin(op, &acc, &pl)?;
        let written = if inclusive { &next } else { &acc };
        let out = self.wrap(
            out.dynamic_update_slice(written, &starts),
            "scan body: write plane",
        )?;
        let one = self.wrap(bb.c0(1i64), "scan body: one")?;
        let k1 = self.wrap(k.add_(&one), "scan body: k + 1")?;
        let body = self.wrap(bb.tuple(&[k1, s, next, out]), "scan body: pack")?;
        let body = body
            .build()
            .map_err(|e| self.err(format!("scan body build: {e}")))?;

        let acc0 = {
            let c = self.c(init)?;
            self.wrap(c.broadcast(&plane), "scan: initial accumulator")?
        };
        let start = self.wrap(
            self.b
                .tuple(&[self.ci(0)?, src.clone(), acc0, self.zeros(dims)?]),
            "scan: initial state",
        )?;
        let res = self.wrap(XlaOp::while_(cond, body, start), "scan: while")?;
        self.wrap(res.get_tuple_element(3), "scan: result")
    }

    /// The variable whose flat block contains `flat`.
    fn var_owning_slot(&self, flat: usize) -> R<usize> {
        for (i, (_, vs)) in self.var_order.iter().enumerate() {
            let n: usize = vs.shape.iter().product::<usize>().max(1);
            if flat >= vs.flat_offset && flat < vs.flat_offset + n {
                return Ok(i);
            }
        }
        Err(self.err(format!(
            "scalar dy write addresses flat slot {flat}, which lies in no variable block"
        )))
    }

    /// Concatenate the per-variable blocks into the flat, column-major `du`.
    fn assemble_du(&mut self) -> R<XlaOp> {
        self.cur_rule = "<program>".to_string();
        let mut flat: Vec<XlaOp> = Vec::with_capacity(self.var_order.len());
        for (i, (_, vs)) in self.var_order.iter().enumerate() {
            let n: usize = vs.shape.iter().product::<usize>().max(1);
            let blk = match &self.blocks[i] {
                Some(b) => b.clone(),
                None => self.zeros(&vs.shape)?,
            };
            let f = if vs.shape.is_empty() {
                self.wrap(blk.reshape(&[1]), "du: reshape scalar block")?
            } else {
                // Row-major -> column-major is the reversing transpose.
                let perm: Vec<i64> = (0..vs.shape.len() as i64).rev().collect();
                self.wrap(
                    blk.transpose(&perm).and_then(|b| b.reshape(&[n as i64])),
                    "du: relabel block",
                )?
            };
            flat.push(f);
        }
        match flat.len() {
            0 => self.zeros(&[0]),
            1 => Ok(flat.remove(0)),
            _ => {
                let head = flat.remove(0);
                self.wrap(head.concat_in_dim(&flat, 0), "du: concat")
            }
        }
    }
}