///|
/// What one forward pass produces for one question.
pub struct Outputs {
  /// Raw scorer logits, one per option slot. Always at least two long: a
  /// one-option choice is padded with a `-1e4` slot, matching the reference
  /// runtime, because the action head's top-two feature needs two slots.
  logits : Array[Double]
  /// How many of those slots are real options.
  options : Int
  /// Raw action-head logits.
  action : Array[Double]
} derive(Debug)

///|
/// Value written into padded option slots, matching the reference runtime.
/// After the softmax shift this is indistinguishable from `-inf`.
const MASKED_LOGIT : Double = -1.0e4

///|
/// Value used to mask out-of-window attention scores. Large enough that
/// `exp(x - rowmax)` underflows to exactly zero, so it behaves as `-inf`
/// without risking a NaN if a row were ever fully masked.
const MASKED_SCORE : Float = -1.0e30

///|
/// Scratch buffers for a single sequence.
///
/// One question is run per forward pass rather than a padded batch. The
/// reference runtime pads a batch and then masks the padding out of every
/// attention row, so results for real tokens are identical either way — and
/// without padding there is no attention mask to carry around, and peak memory
/// stays proportional to the actual sequence length.
priv struct Scratch {
  hidden : FixedArray[Float]
  normed : FixedArray[Float]
  qkv : FixedArray[Float]
  q : FixedArray[Float]
  k : FixedArray[Float]
  v : FixedArray[Float]
  scores : FixedArray[Float]
  context : FixedArray[Float]
  packed : FixedArray[Float]
  projected : FixedArray[Float]
  wide : FixedArray[Float]
  gated : FixedArray[Float]
}

///|
fn Scratch::new(model : Model, length : Int) -> Scratch {
  let hidden = model.encoder.hidden_size
  let inner = model.encoder.intermediate_size
  // The decision head's feed-forward is 4x hidden, wider than the encoder's
  // 2x intermediate on these checkpoints, so `wide` is sized for whichever is
  // larger and shared between them.
  let widest = @cmp.maximum(2 * inner, 4 * hidden)
  let heads = @cmp.maximum(model.encoder.num_attention_heads, model.head_heads)
  {
    hidden: @tensor.zeros(length * hidden),
    normed: @tensor.zeros(length * hidden),
    qkv: @tensor.zeros(length * 3 * hidden),
    q: @tensor.zeros(length * hidden),
    k: @tensor.zeros(length * hidden),
    v: @tensor.zeros(length * hidden),
    scores: @tensor.zeros(heads * length * length),
    context: @tensor.zeros(length * hidden),
    packed: @tensor.zeros(length * hidden),
    projected: @tensor.zeros(length * hidden),
    wide: @tensor.zeros(length * widest),
    gated: @tensor.zeros(length * @cmp.maximum(inner, 4 * hidden)),
  }
}

///|
/// Run one question through the encoder, the decision head, the scorer and the
/// action head.
///
/// `markers` holds the position of each option's `[MASK]` token in `ids`, and
/// `qtype` selects the learned type embedding (0 choice, 1 score, 2 noul).
///
/// Passing `trace` records each stage's activations under a descriptive key;
/// it exists for the parity tests and costs a copy per stage when supplied.
pub fn Model::forward(
  self : Model,
  ids : Array[Int],
  markers : Array[Int],
  qtype : Int,
  trace? : Map[String, Array[Double]],
) -> Outputs raise @error.LayaError {
  self.check_inputs(ids, markers, qtype)
  let options = markers.length()
  let slots = @cmp.maximum(options, 2)
  let (logits, cls) = match self.metal {
    Some(state) => self.forward_metal(state, ids, markers, qtype, trace)
    None => self.forward_cpu(ids, markers, qtype, trace)
  }
  let action = self.action_logits(cls, logits, options, slots)
  { logits, options, action, }
}

///|
/// Reject inputs neither backend can run.
fn Model::check_inputs(
  self : Model,
  ids : Array[Int],
  markers : Array[Int],
  qtype : Int,
) -> Unit raise @error.LayaError {
  let length = ids.length()
  @error.expect(length > 0, "cannot run an empty sequence")
  @error.expect(
    length <= self.encoder.max_position_embeddings,
    "sequence of \{length} tokens exceeds the encoder's \{self.encoder.max_position_embeddings}-position limit",
  )
  @error.expect(
    qtype >= 0 && qtype < 3,
    "question type must be 0, 1 or 2, got \{qtype}",
  )
  @error.expect(markers.length() > 0, "a question needs at least one option")
  for marker in markers {
    @error.expect(
      marker >= 0 && marker < length,
      "option marker \{marker} falls outside a \{length}-token sequence",
    )
  }
}

///|
/// Run one sequence on the CPU, returning the option logits and the pooled CLS
/// row the action head needs.
fn Model::forward_cpu(
  self : Model,
  ids : Array[Int],
  markers : Array[Int],
  qtype : Int,
  trace : Map[String, Array[Double]]?,
) -> (Array[Double], FixedArray[Float]) raise @error.LayaError {
  let length = ids.length()
  let hidden = self.encoder.hidden_size
  let scratch = Scratch::new(self, length)

  // Embeddings: a row gather straight out of the mapped checkpoint, then a
  // norm. The embedding matrix is never converted to float32 wholesale.
  let rows = FixedArray::make(length, 0)
  for index, id in ids {
    @error.expect(
      id >= 0 && id < self.encoder.vocab_size,
      "token id \{id} is outside the \{self.encoder.vocab_size}-entry vocabulary",
    )
    rows[index] = id
  }
  self.file.gather_rows(self.tok_embeddings, rows, scratch.normed)
  @tensor.layer_norm(
    scratch.normed,
    scratch.hidden,
    rows=length,
    cols=hidden,
    self.embed_norm.values,
    None,
    eps=self.encoder.norm_eps,
  )
  record(trace, "embeddings", scratch.hidden, length * hidden)
  for index, layer in self.layers {
    self.encoder_layer(scratch, layer, length)
    record(trace, "encoder_layer_\{index}", scratch.hidden, length * hidden)
  }
  @tensor.layer_norm(
    scratch.hidden,
    scratch.normed,
    rows=length,
    cols=hidden,
    self.final_norm.values,
    None,
    eps=self.encoder.norm_eps,
  )
  scratch.normed.blit_to(scratch.hidden, len=length * hidden)
  record(trace, "encoder_out", scratch.hidden, length * hidden)

  // The question type is a single learned vector added to every position.
  for position in 0..