Training & autograd

Reverse-mode autograd over the graph, the Trainer loop (with the optimizer step on the engine), and layer-streamed training for deep stacks.

autograd

Reverse-mode autograd over the aneforge graph; forward and backward both run on the ANE. See docs/developer/autograd.md.

CEHandle

A softmax-cross-entropy objective carrying logits and one-hot target; the logit grad is the analytic (softmax-target)/N form.

Source code in aneforge/autograd.py

class CEHandle:
  """A softmax-cross-entropy objective carrying logits and one-hot target; the logit grad is the analytic (softmax-target)/N form."""
  def __init__(self, logits: Tensor, target: Tensor):
    if len(logits.shape) != 2 or logits.shape != target.shape:
      raise ValueError(f"softmax_cross_entropy expects 2-D logits/target [N,K]; "
                       f"got {logits.shape}, {target.shape}")
    self.logits, self.target, self.n = logits, target, int(logits.shape[0])

  def seed(self, loss_scale: float) -> Tensor:
    """dL/dlogits * loss_scale = (softmax(logits) - target) * (loss_scale / n)."""
    return (self.logits.softmax(-1) - self.target) * (float(loss_scale) / self.n)

seed

seed(loss_scale: float) -> Tensor

dL/dlogits * loss_scale = (softmax(logits) - target) * (loss_scale / n).

Source code in aneforge/autograd.py

def seed(self, loss_scale: float) -> Tensor:
  """dL/dlogits * loss_scale = (softmax(logits) - target) * (loss_scale / n)."""
  return (self.logits.softmax(-1) - self.target) * (float(loss_scale) / self.n)

SGD

Host fp32 SGD over the parameters' master values (loss-scaled grads divided out before the step).

Source code in aneforge/autograd.py

class SGD:
  """Host fp32 SGD over the parameters' master values (loss-scaled grads divided out before the step)."""
  def __init__(self, params, lr: float, loss_scale: float = 1.0):
    self.params, self.lr, self.scale = list(params), float(lr), float(loss_scale)
    self._nonfinite_skips = 0

  def step(self, grads):
    if not _check_finite_grads(self, grads): return
    for p, g in zip(self.params, grads):
      p.attrs["value"] = p.attrs["value"] - self.lr * (g.astype(np.float32) / self.scale)

Adam

Host fp32 Adam over the parameters' master values (loss-scaled grads divided out before the moment update).

Source code in aneforge/autograd.py

class Adam:
  """Host fp32 Adam over the parameters' master values (loss-scaled grads divided out before the moment update)."""
  def __init__(self, params, lr: float = 1e-3, betas=(0.9, 0.999),
               eps: float = 1e-8, loss_scale: float = 1.0):
    self.params = list(params)
    self.lr, (self.b1, self.b2), self.eps, self.scale = float(lr), betas, float(eps), float(loss_scale)
    self.m = [np.zeros(p.attrs["value"].shape, np.float32) for p in self.params]
    self.v = [np.zeros_like(x) for x in self.m]
    self.t = 0
    self._nonfinite_skips = 0

  def step(self, grads):
    if not _check_finite_grads(self, grads): return
    self.t += 1
    bc1, bc2 = 1.0 - self.b1 ** self.t, 1.0 - self.b2 ** self.t
    for i, (p, g) in enumerate(zip(self.params, grads)):
      g = g.astype(np.float32) / self.scale
      self.m[i] = self.b1 * self.m[i] + (1 - self.b1) * g
      self.v[i] = self.b2 * self.v[i] + (1 - self.b2) * (g * g)
      mhat, vhat = self.m[i] / bc1, self.v[i] / bc2
      p.attrs["value"] = p.attrs["value"] - self.lr * mhat / (np.sqrt(vhat) + self.eps)

Trainer

Compiles a forward program once plus one backward program per parameter; step evals them and applies the optimizer (host-side by default, or on-ANE with device_optimizer=True).

Source code in aneforge/autograd.py

class Trainer:
  """Compiles a forward program once plus one backward program per parameter; `step` evals them and applies the optimizer (host-side by default, or on-ANE with `device_optimizer=True`)."""
  def __init__(self, objective, params, lr: float, loss_scale: float = 1.0,
               data_inputs: dict | None = None, optimizer: str = "sgd",
               betas=(0.9, 0.999), eps: float = 1e-8, device_optimizer: bool = False,
               resident_state: bool = False):
    from . import _compile as _c
    self.params = list(params)
    loss_scale = _guard_a13_conv_loss_scale(self.params, float(loss_scale))  # before scale is consumed
    self.data = dict(data_inputs or {})       # {input Tensor: numpy value}
    self.scale = float(loss_scale)
    self.lr = float(lr)
    self.b1, self.b2 = float(betas[0]), float(betas[1])
    self.eps = float(eps)
    self.optimizer = optimizer
    # resident_state forces device_optimizer on (it is the on-device update).
    self.resident_state = bool(resident_state)
    self.device_optimizer = bool(device_optimizer) or self.resident_state
    self.opt = (Adam(self.params, lr, betas, eps=eps, loss_scale=loss_scale)
                if optimizer == "adam" else SGD(self.params, lr, loss_scale))
    if isinstance(objective, CEHandle):
      self.ce = objective
      grads = backward_from(objective.seed(loss_scale), objective.logits, self.params)
      fwd_out = objective.logits                 # forward program -> logits
    else:
      self.ce = None
      grads = backward(objective, self.params, loss_scale=loss_scale)
      fwd_out = objective                        # forward program -> loss scalar
    # one backward program per param, in natural shape (one wide row trips the wide-row wall).
    self._fwd = _c.compile(fwd_out, _check_precision=False)
    if self.resident_state:
      self._build_resident(_c, grads)            # _fwd kept only for checkpoint accuracy
    else:
      self._bwd = [_c.compile(grads[p], _check_precision=False) for p in self.params]
      if self.device_optimizer: self._build_device_optimizer(_c)

  def _build_device_optimizer(self, _c):
    """Compile a per-param update program so optimizer arithmetic runs on the ANE (SGD -> w'; Adam -> stack(w', m', v'))."""
    self._upd, self._upd_g, self._upd_lr = [], [], []
    if self.optimizer == "adam":
      self._m = [np.zeros(p.shape, np.float16) for p in self.params]
      self._v = [np.zeros(p.shape, np.float16) for p in self.params]
      self._upd_m, self._upd_v = [], []
      self._t = 0
    for p in self.params:
      g_in = graph.input(p.shape)
      lr_in = graph.input((1, 1))
      self._upd_g.append(g_in); self._upd_lr.append(lr_in)
      if self.optimizer == "adam":
        m_in = graph.input(p.shape); v_in = graph.input(p.shape)
        self._upd_m.append(m_in); self._upd_v.append(v_in)
        w2, m2, v2 = _adam_update(p, m_in, v_in, g_in, lr_in, self.b1, self.b2, self.eps)
        out = _stack3(w2, m2, v2)
      else:
        out = _sgd_update(p, g_in, lr_in)
      self._upd.append(_c.compile(out, _check_precision=False))

  def _build_resident(self, _c, grads):
    """Assemble the whole step as one fused program with optimizer state resident on-device (each updated-state output aliased onto its input port via share_buffer)."""
    lr_in = graph.input((1, 1))
    self._res_lr = lr_in
    self._t = 0
    outs, alias, self._res_state = [], [], []
    for p in self.params:
      g = grads[p]
      entry = {"p": p}
      if self.optimizer == "adam":
        m_in, v_in = graph.input(p.shape), graph.input(p.shape)
        w2, m2, v2 = _adam_update(p, m_in, v_in, g, lr_in, self.b1, self.b2, self.eps)
        outs += [w2, m2, v2]
        alias += [(w2, p), (m2, m_in), (v2, v_in)]
        entry.update(m_in=m_in, v_in=v_in, w_out=w2)
      else:
        w2 = _sgd_update(p, g, lr_in)
        outs += [w2]
        alias += [(w2, p)]
        entry.update(w_out=w2)
      self._res_state.append(entry)

    mm = _c.compile_multi(outs)
    self._res = mm
    prog = mm.prog
    self._res_in_name = dict(mm.input_ports)
    self._res_out_name = dict(mm.output_ports)
    # alias each updated-state output onto its input port, then seed once.
    for out_t, in_t in alias:
      prog.share_buffer(0, self._res_out_name[out_t], 0, self._res_in_name[in_t])
    for entry in self._res_state:
      p = entry["p"]
      prog.set_input(self._res_in_name[p], p.attrs["value"].astype(np.float16))
      if self.optimizer == "adam":
        prog.set_input(self._res_in_name[entry["m_in"]], np.zeros(p.shape, np.float16))
        prog.set_input(self._res_in_name[entry["v_in"]], np.zeros(p.shape, np.float16))
    # the remaining inputs (neither state nor lr) are the data ports (x, target)
    state_ids = set()
    for entry in self._res_state:
      state_ids.add(id(entry["p"]))
      if self.optimizer == "adam":
        state_ids.update({id(entry["m_in"]), id(entry["v_in"])})
    self._res_data_inputs = [t for t, _ in mm.input_ports
                             if id(t) not in state_ids and t is not lr_in]
    self._res_dirty = False

  def _resident_step(self) -> None:
    """One resident step: feed only the minibatch + lr_t, execute (state stays on-device)."""
    if getattr(self, "_ds", None) is not None:
      xin, X, tin, Y = self._ds
      idx = self._next_batch()
      self.data[xin] = X[idx]; self.data[tin] = Y[idx]
    prog = self._res.prog
    for t in self._res_data_inputs:
      prog.set_input(self._res_in_name[t], np.asarray(self.data[t], np.float16))
    if self.optimizer == "adam":
      self._t += 1
      lr_t = self.lr * math.sqrt(1.0 - self.b2 ** self._t) / (1.0 - self.b1 ** self._t)
    else:
      lr_t = self.lr / self.scale
    prog.set_input(self._res_in_name[self._res_lr], np.full((1, 1), lr_t, np.float16))
    prog.execute()
    self._res_dirty = True

  def _sync_params_from_device(self) -> None:
    """Checkpoint read: copy resident params off-device into the host masters; moments stay resident."""
    prog = self._res.prog
    for i, entry in enumerate(self._res_state):
      w = prog.read_output(self._res_out_name[entry["w_out"]]).astype(np.float32)
      if not np.isfinite(w).all():
        import warnings
        warnings.warn(
          f"aneforge.Trainer: resident param {i} (shape {tuple(entry['p'].shape)}) read "
          f"back non-finite values (inf/nan) at checkpoint - the on-device update was "
          f"poisoned by an overflowed fp16 gradient; if you see inf/nan weight-grads, "
          f"lower loss_scale.",
          stacklevel=3)
      entry["p"].attrs["value"] = w.reshape(entry["p"].shape)
    self._res_dirty = False

  def _feed(self, model, override: dict | None = None):
    """Map each compiled input Tensor to a fp16 array: trainable -> master value; in `override` -> override[t]; else the baked attrs['value']."""
    lk = self.data if override is None else override
    return [
      t.attrs["value"].astype(np.float16) if t.attrs.get("trainable")
      else np.asarray(lk[t] if t in lk else t.attrs["value"]).astype(np.float16)
      for t in model._input_tensors
    ]

  def _feed_update(self, net, p, extra):
    """Feed a per-param update program (delegates to _feed with `extra`)."""
    return self._feed(net, extra)

  def set_dataset(self, x_input, X_full, target_input, Y_onehot, seed: int = 0):
    """Provide the full dataset for mini-batch sampling (`x_input`/`target_input` are the batch-B graph placeholders)."""
    self._ds = (x_input, np.asarray(X_full, np.float32), target_input, np.asarray(Y_onehot, np.float32))
    self._ds_B = int(x_input.shape[0])
    self._ds_rng = np.random.default_rng(seed)
    self._ds_perm = self._ds_rng.permutation(len(self._ds[1]))
    self._ds_pos = 0

  def _next_batch(self):
    B = self._ds_B
    if self._ds_pos + B > len(self._ds_perm):       # reshuffle each epoch
      self._ds_perm = self._ds_rng.permutation(len(self._ds[1]))
      self._ds_pos = 0
    idx = self._ds_perm[self._ds_pos:self._ds_pos + B]
    self._ds_pos += B
    return idx

  def step(self) -> None:
    if self.resident_state:
      self._resident_step()
      return
    if getattr(self, "_ds", None) is not None:
      xin, X, tin, Y = self._ds
      idx = self._next_batch()
      self.data[xin] = X[idx]
      self.data[tin] = Y[idx]
    grads = [np.asarray(net(*self._feed(net))).reshape(p.shape)
             for net, p in zip(self._bwd, self.params)]
    if not self.device_optimizer:
      self.opt.step(grads)
      return
    # On-ANE optimizer: update arithmetic runs as graph ops; host only computes lr_t.
    if self.optimizer == "adam":
      self._device_adam_step(grads)
    else:
      self._device_sgd_step(grads)

  def _device_sgd_step(self, grads):
    lr_t = self.lr / self.scale          # fold grad-unscale into lr_t
    lr_arr = np.full((1, 1), lr_t, np.float16)
    for i, p in enumerate(self.params):
      net, g_in, lr_in = self._upd[i], self._upd_g[i], self._upd_lr[i]
      extra = {g_in: grads[i], lr_in: lr_arr}
      w2 = np.asarray(net(*self._feed_update(net, p, extra))).reshape(p.shape)
      p.attrs["value"] = w2.astype(np.float32)

  def _device_adam_step(self, grads):
    self._t += 1
    # loss-scale is NOT divided out: it cancels in Adam's m/sqrt(v) ratio.
    lr_t = self.lr * math.sqrt(1.0 - self.b2 ** self._t) / (1.0 - self.b1 ** self._t)
    lr_arr = np.full((1, 1), lr_t, np.float16)
    for i, p in enumerate(self.params):
      net = self._upd[i]
      extra = {self._upd_g[i]: grads[i], self._upd_m[i]: self._m[i],
               self._upd_v[i]: self._v[i], self._upd_lr[i]: lr_arr}
      out = np.asarray(net(*self._feed_update(net, p, extra)))
      w2, m2, v2 = _split3(out, p.shape)
      p.attrs["value"] = w2.astype(np.float32)
      self._m[i] = m2.astype(np.float16)       # fp16 optimizer state held host-side
      self._v[i] = v2.astype(np.float16)

  def accuracy(self, X, y_labels) -> float:
    """Argmax accuracy over X (any length) via the batch-B forward program, chunking X into B-row pieces."""
    assert self.ce is not None, "accuracy() is for classification objectives"
    if self.resident_state and getattr(self, "_res_dirty", False):
      self._sync_params_from_device()
    X = np.asarray(X, np.float32); y = np.asarray(y_labels)
    # the feature input is the non-trainable, non-target input whose per-sample shape matches X.
    feat_shape = X.shape[1:]
    assert isinstance(self._fwd, _Model)  # Trainer always compiles non-sdpa graphs
    xin = next(t for t in self._fwd._input_tensors
               if not t.attrs.get("trainable") and t is not getattr(self.ce, "target", None)
               and tuple(t.shape[1:]) == tuple(feat_shape))
    B = xin.shape[0]
    saved = self.data.get(xin)
    preds = []
    for s in range(0, X.shape[0], B):
      chunk = X[s:s + B]
      m = chunk.shape[0]
      if m < B:
        pad = np.zeros((B - m,) + tuple(feat_shape), np.float32)
        chunk = np.concatenate([chunk, pad], axis=0)
      self.data[xin] = chunk
      logits = np.asarray(self._fwd(*self._feed(self._fwd)))
      preds.append(logits[:m].argmax(1))
    if saved is not None: self.data[xin] = saved
    return float((np.concatenate(preds) == y).mean())

  def loss(self) -> float:
    if self.resident_state and getattr(self, "_res_dirty", False):
      self._sync_params_from_device()
    out = np.asarray(self._fwd(*self._feed(self._fwd)))
    if self.ce is None: return float(out.reshape(-1)[0])
    logits = out.reshape(self.ce.logits.shape)
    t = np.asarray(self.data[self.ce.target])
    z = logits - logits.max(1, keepdims=True)
    logsm = z - np.log(np.exp(z).sum(1, keepdims=True))
    return float(-(t * logsm).sum(1).mean())

  def release(self) -> None:
    self._fwd.release()
    if getattr(self, "_res", None) is not None: self._res.release()
    for net in getattr(self, "_bwd", []): net.release()
    for net in getattr(self, "_upd", []): net.release()

set_dataset

set_dataset(x_input, X_full, target_input, Y_onehot, seed: int = 0)

Provide the full dataset for mini-batch sampling (x_input/target_input are the batch-B graph placeholders).

Source code in aneforge/autograd.py

def set_dataset(self, x_input, X_full, target_input, Y_onehot, seed: int = 0):
  """Provide the full dataset for mini-batch sampling (`x_input`/`target_input` are the batch-B graph placeholders)."""
  self._ds = (x_input, np.asarray(X_full, np.float32), target_input, np.asarray(Y_onehot, np.float32))
  self._ds_B = int(x_input.shape[0])
  self._ds_rng = np.random.default_rng(seed)
  self._ds_perm = self._ds_rng.permutation(len(self._ds[1]))
  self._ds_pos = 0

accuracy

accuracy(X, y_labels) -> float

Argmax accuracy over X (any length) via the batch-B forward program, chunking X into B-row pieces.

Source code in aneforge/autograd.py

def accuracy(self, X, y_labels) -> float:
  """Argmax accuracy over X (any length) via the batch-B forward program, chunking X into B-row pieces."""
  assert self.ce is not None, "accuracy() is for classification objectives"
  if self.resident_state and getattr(self, "_res_dirty", False):
    self._sync_params_from_device()
  X = np.asarray(X, np.float32); y = np.asarray(y_labels)
  # the feature input is the non-trainable, non-target input whose per-sample shape matches X.
  feat_shape = X.shape[1:]
  assert isinstance(self._fwd, _Model)  # Trainer always compiles non-sdpa graphs
  xin = next(t for t in self._fwd._input_tensors
             if not t.attrs.get("trainable") and t is not getattr(self.ce, "target", None)
             and tuple(t.shape[1:]) == tuple(feat_shape))
  B = xin.shape[0]
  saved = self.data.get(xin)
  preds = []
  for s in range(0, X.shape[0], B):
    chunk = X[s:s + B]
    m = chunk.shape[0]
    if m < B:
      pad = np.zeros((B - m,) + tuple(feat_shape), np.float32)
      chunk = np.concatenate([chunk, pad], axis=0)
    self.data[xin] = chunk
    logits = np.asarray(self._fwd(*self._feed(self._fwd)))
    preds.append(logits[:m].argmax(1))
  if saved is not None: self.data[xin] = saved
  return float((np.concatenate(preds) == y).mean())

UnrolledTrainer

Train with K Adam steps unrolled into one fused ANE program (each step() runs K fwd->bwd->update in one dispatch); resident=True keeps state on-device.

Source code in aneforge/autograd.py

class UnrolledTrainer:
  """Train with K Adam steps unrolled into one fused ANE program (each `step()` runs K fwd->bwd->update in one dispatch); `resident=True` keeps state on-device."""
  def __init__(self, params, forward, kind, x_inputs, t_inputs, dataset, lr,
               loss_scale: float = 1.0, betas=(0.9, 0.999), eps: float = 1e-8,
               seed: int = 0, resident: bool = True):
    from . import _compile as _c
    if kind not in ("ce", "mse"):
      raise ValueError("kind must be 'ce' or 'mse'")
    self.params = list(params)
    self.forward = forward
    self.kind = kind
    self.K = len(x_inputs)
    self.lr = float(lr); self.scale = float(loss_scale)
    self.b1, self.b2 = float(betas[0]), float(betas[1]); self.eps = float(eps)
    self.X = np.asarray(dataset[0], np.float32)
    self.Y = np.asarray(dataset[1], np.float32)
    self.B = int(x_inputs[0].shape[0])
    self.m = [np.zeros(p.shape, np.float16) for p in self.params]
    self.v = [np.zeros(p.shape, np.float16) for p in self.params]
    self.t = 0
    self.rng = np.random.default_rng(seed)
    self._perm = self.rng.permutation(len(self.X)); self._pos = 0

    # unrolled graph: thread (P, m, v) through K Adam steps
    m_in = [graph.input(p.shape) for p in self.params]
    v_in = [graph.input(p.shape) for p in self.params]
    lr_ins = [graph.input((1, 1)) for _ in range(self.K)]
    P, M, V = list(self.params), list(m_in), list(v_in)
    for k in range(self.K):
      out = forward(P, x_inputs[k])
      if kind == "ce":
        g = backward_from(softmax_cross_entropy(out, t_inputs[k]).seed(self.scale), out, P)
      else:
        g = backward(mse(out, t_inputs[k]), P, loss_scale=self.scale)
      P, M, V = adam_step(P, M, V, g, lr_ins[k], (self.b1, self.b2), self.eps)
    self._net = _c.compile_multi([*P, *M, *V])
    self._oname = dict(self._net.output_ports)
    self._P_out, self._M_out, self._V_out = P, M, V
    self._m_in, self._v_in, self._lr_ins = m_in, v_in, lr_ins
    # map each data input tensor -> (step k, 'x'|'t')
    self._data_map = {}
    for k in range(self.K):
      self._data_map[id(x_inputs[k])] = (k, "x")
      self._data_map[id(t_inputs[k])] = (k, "t")

    # separate single-batch forward program for checkpoint predict, with its OWN
    # weight leaves (compile mutates Tensor names; sharing params would clobber ports).
    ev_w = [graph.input(p.shape) for p in self.params]
    for ew, p in zip(ev_w, self.params):
      if "conv_shape" in p.attrs: ew.attrs["conv_shape"] = p.attrs["conv_shape"]
    xe = graph.input(x_inputs[0].shape)
    self._ev_w = ev_w
    self._eval = _c.compile(forward(ev_w, xe), _check_precision=False)

    self.resident = bool(resident)
    if self.resident:
      # alias each final state output onto its initial input port (state lives on-device); seed once.
      prog = self._net.prog
      inm = {id(t): n for t, n in self._net.input_ports}
      self._res_inm = inm
      self._res_lr_names = [inm[id(t)] for t in lr_ins]
      self._res_data = ([(inm[id(x_inputs[k])], k, "x") for k in range(self.K)] +
                        [(inm[id(t_inputs[k])], k, "t") for k in range(self.K)])
      pairs = (list(zip(self._P_out, self.params)) +
               list(zip(self._M_out, m_in)) + list(zip(self._V_out, v_in)))
      for out_t, in_t in pairs:
        prog.share_buffer(0, self._oname[out_t], 0, inm[id(in_t)])
      for i, p in enumerate(self.params):
        prog.set_input(inm[id(p)], p.attrs["value"].astype(np.float16))
        prog.set_input(inm[id(m_in[i])], np.zeros(p.shape, np.float16))
        prog.set_input(inm[id(v_in[i])], np.zeros(p.shape, np.float16))
      self._res_dirty = False

  def _next(self):
    if self._pos + self.B > len(self._perm):
      self._perm = self.rng.permutation(len(self.X)); self._pos = 0
    idx = self._perm[self._pos:self._pos + self.B]; self._pos += self.B
    return idx

  def step(self) -> None:
    """Run K training steps on the ANE in one dispatch (resident: feed only the K minibatches + per-step lr; else shuttle params/m/v)."""
    batches = [self._next() for _ in range(self.K)]
    if self.resident:
      prog = self._net.prog
      for name, k, which in self._res_data:
        idx = batches[k]
        prog.set_input(name, (self.X[idx] if which == "x" else self.Y[idx]).astype(np.float16))
      for k, name in enumerate(self._res_lr_names):
        gt = self.t + k + 1
        prog.set_input(name, np.full((1, 1), self.lr * math.sqrt(1.0 - self.b2 ** gt) /
                                     (1.0 - self.b1 ** gt), np.float16))
      prog.execute()
      self.t += self.K
      self._res_dirty = True
      return
    vals = []
    for t in self._net.input_tensors:
      if t.attrs.get("trainable"):
        vals.append(t.attrs["value"].astype(np.float16))
      elif t in self._m_in:
        vals.append(self.m[self._m_in.index(t)])
      elif t in self._v_in:
        vals.append(self.v[self._v_in.index(t)])
      elif t in self._lr_ins:
        gt = self.t + self._lr_ins.index(t) + 1          # global step for bias correction
        lr_t = self.lr * math.sqrt(1.0 - self.b2 ** gt) / (1.0 - self.b1 ** gt)
        vals.append(np.full((1, 1), lr_t, np.float16))
      else:
        k, which = self._data_map[id(t)]
        idx = batches[k]
        vals.append((self.X[idx] if which == "x" else self.Y[idx]).astype(np.float16))
    out = self._net(*vals)
    for i, p in enumerate(self.params):
      p.attrs["value"] = out[self._oname[self._P_out[i]]].reshape(p.shape)
      self.m[i] = out[self._oname[self._M_out[i]]].astype(np.float16).reshape(p.shape)
      self.v[i] = out[self._oname[self._V_out[i]]].astype(np.float16).reshape(p.shape)
    self.t += self.K

  def _sync_from_device(self) -> None:
    """Checkpoint read: copy the resident params off-device into the host masters."""
    prog = self._net.prog
    for i, p in enumerate(self.params):
      w = prog.read_output(self._oname[self._P_out[i]]).astype(np.float32)
      p.attrs["value"] = w.reshape(p.shape)
    self._res_dirty = False

  def predict(self, X) -> np.ndarray:
    """Run the trained weights forward on the ANE in B-sized chunks; returns logits ('ce') or prediction ('mse')."""
    if self.resident and getattr(self, "_res_dirty", False):
      self._sync_from_device()
    X = np.asarray(X, np.float32)
    feeds_w = [p.attrs["value"].astype(np.float16) for p in self.params]
    outs = []
    for s in range(0, len(X), self.B):
      chunk = X[s:s + self.B]
      pad = self.B - len(chunk)
      if pad:
        chunk = np.concatenate([chunk, np.zeros((pad, *chunk.shape[1:]), np.float32)])
      # eval inputs are [weight leaves..., xe]: feed masters then the chunk.
      assert isinstance(self._eval, _Model)  # UnrolledTrainer eval never uses sdpa nodes
      args = []
      for t in self._eval._input_tensors:
        args.append(feeds_w[self._ev_w.index(t)] if t in self._ev_w
                    else chunk.astype(np.float16))
      o = np.asarray(self._eval(*args), np.float32)
      outs.append(o[:len(chunk)] if pad else o)
    return np.concatenate(outs)

  def release(self) -> None:
    self._net.release(); self._eval.release()

step

step() -> None

Run K training steps on the ANE in one dispatch (resident: feed only the K minibatches + per-step lr; else shuttle params/m/v).

Source code in aneforge/autograd.py

def step(self) -> None:
  """Run K training steps on the ANE in one dispatch (resident: feed only the K minibatches + per-step lr; else shuttle params/m/v)."""
  batches = [self._next() for _ in range(self.K)]
  if self.resident:
    prog = self._net.prog
    for name, k, which in self._res_data:
      idx = batches[k]
      prog.set_input(name, (self.X[idx] if which == "x" else self.Y[idx]).astype(np.float16))
    for k, name in enumerate(self._res_lr_names):
      gt = self.t + k + 1
      prog.set_input(name, np.full((1, 1), self.lr * math.sqrt(1.0 - self.b2 ** gt) /
                                   (1.0 - self.b1 ** gt), np.float16))
    prog.execute()
    self.t += self.K
    self._res_dirty = True
    return
  vals = []
  for t in self._net.input_tensors:
    if t.attrs.get("trainable"):
      vals.append(t.attrs["value"].astype(np.float16))
    elif t in self._m_in:
      vals.append(self.m[self._m_in.index(t)])
    elif t in self._v_in:
      vals.append(self.v[self._v_in.index(t)])
    elif t in self._lr_ins:
      gt = self.t + self._lr_ins.index(t) + 1          # global step for bias correction
      lr_t = self.lr * math.sqrt(1.0 - self.b2 ** gt) / (1.0 - self.b1 ** gt)
      vals.append(np.full((1, 1), lr_t, np.float16))
    else:
      k, which = self._data_map[id(t)]
      idx = batches[k]
      vals.append((self.X[idx] if which == "x" else self.Y[idx]).astype(np.float16))
  out = self._net(*vals)
  for i, p in enumerate(self.params):
    p.attrs["value"] = out[self._oname[self._P_out[i]]].reshape(p.shape)
    self.m[i] = out[self._oname[self._M_out[i]]].astype(np.float16).reshape(p.shape)
    self.v[i] = out[self._oname[self._V_out[i]]].astype(np.float16).reshape(p.shape)
  self.t += self.K

predict

predict(X) -> np.ndarray

Run the trained weights forward on the ANE in B-sized chunks; returns logits ('ce') or prediction ('mse').

Source code in aneforge/autograd.py

def predict(self, X) -> np.ndarray:
  """Run the trained weights forward on the ANE in B-sized chunks; returns logits ('ce') or prediction ('mse')."""
  if self.resident and getattr(self, "_res_dirty", False):
    self._sync_from_device()
  X = np.asarray(X, np.float32)
  feeds_w = [p.attrs["value"].astype(np.float16) for p in self.params]
  outs = []
  for s in range(0, len(X), self.B):
    chunk = X[s:s + self.B]
    pad = self.B - len(chunk)
    if pad:
      chunk = np.concatenate([chunk, np.zeros((pad, *chunk.shape[1:]), np.float32)])
    # eval inputs are [weight leaves..., xe]: feed masters then the chunk.
    assert isinstance(self._eval, _Model)  # UnrolledTrainer eval never uses sdpa nodes
    args = []
    for t in self._eval._input_tensors:
      args.append(feeds_w[self._ev_w.index(t)] if t in self._ev_w
                  else chunk.astype(np.float16))
    o = np.asarray(self._eval(*args), np.float32)
    outs.append(o[:len(chunk)] if pad else o)
  return np.concatenate(outs)

vjp

vjp(*names: str)

Register a vjp rule fn(node, g) -> list[grad|None] (one per node.srcs).

Source code in aneforge/autograd.py

def vjp(*names: str):
  """Register a vjp rule `fn(node, g) -> list[grad|None]` (one per node.srcs)."""
  def reg(fn):
    for n in names: VJP[n] = fn
    return fn
  return reg

parameter

parameter(init) -> Tensor

A trainable leaf: a graph input tagged trainable, holding an fp32 master value in attrs['value'].

Source code in aneforge/autograd.py

def parameter(init) -> Tensor:
  """A trainable leaf: a graph input tagged trainable, holding an fp32 master value in attrs['value']."""
  init = np.asarray(init, dtype=np.float32)
  t = graph.input(init.shape)
  t.attrs["trainable"] = True
  t.attrs["value"] = init
  return t

backward

backward(loss: Tensor, params, loss_scale: float = 1.0, stop=None) -> dict

Reverse-mode grads of scalar loss wrt each Tensor in params; stop is the detach frontier (defaults to params).

Source code in aneforge/autograd.py

def backward(loss: Tensor, params, loss_scale: float = 1.0, stop=None) -> dict:
  """Reverse-mode grads of scalar `loss` wrt each Tensor in `params`; `stop` is the detach frontier (defaults to `params`)."""
  stop_ids = {id(t) for t in (params if stop is None else stop)}
  order = _topo(loss, stop_ids)
  return _reverse(order, {id(loss): _const_like(loss, float(loss_scale))}, params, stop_ids)

backward_from

backward_from(grad_root, root, params, stop=None) -> dict

Reverse-mode from an explicit gradient grad_root at root (e.g. logits) rather than a scalar loss seed.

Source code in aneforge/autograd.py

def backward_from(grad_root, root, params, stop=None) -> dict:
  """Reverse-mode from an explicit gradient `grad_root` at `root` (e.g. logits) rather than a scalar loss seed."""
  stop_ids = {id(t) for t in (params if stop is None else stop)}
  return _reverse(_topo(root, stop_ids), {id(root): grad_root}, params, stop_ids)

conv_param

conv_param(weight_init) -> Tensor

A trainable conv weight parameter; weight_init is [Cout, Cin, kH, kW] (PyTorch), stored as the flat patch matrix [CinkHkW, Cout].

Source code in aneforge/autograd.py

def conv_param(weight_init) -> Tensor:
  """A trainable conv weight parameter; `weight_init` is [Cout, Cin, kH, kW] (PyTorch), stored as the flat patch matrix [Cin*kH*kW, Cout]."""
  W = np.asarray(weight_init, dtype=np.float32)
  Cout, Cin, kH, kW = W.shape
  flat = W.reshape(Cout, Cin * kH * kW).T.copy()        # [Cin*kH*kW, Cout]
  p = parameter(flat)
  p.attrs["conv_shape"] = (Cout, Cin, kH, kW)
  return p

conv2d

conv2d(x: Tensor, weight: Tensor, stride: int = 1, pad: int = 0) -> Tensor

A trainable stride-1 2-D conv built from primitives so weight is a real graph parameter; x [N,Cin,H,W] -> [N,Cout,Hout,Wout]. Train in mini-batches (compile time grows with N).

Source code in aneforge/autograd.py

def conv2d(x: Tensor, weight: Tensor, stride: int = 1, pad: int = 0) -> Tensor:
  """A trainable stride-1 2-D conv built from primitives so `weight` is a real graph parameter; `x` [N,Cin,H,W] -> [N,Cout,Hout,Wout]. Train in mini-batches (compile time grows with N)."""
  if stride != 1:
    raise NotImplementedError("conv2d (trainable) supports stride=1 only; "
                              "downsample with avg_pool/max_pool.")
  if pad < 0:
    raise ValueError(f"conv2d: pad must be >= 0, got {pad}")
  if "conv_shape" not in weight.attrs:
    raise ValueError("conv2d weight must come from af.conv_param([Cout,Cin,kH,kW])")
  N, Cin, H, W = x.shape
  Cout, Cin_w, kH, kW = weight.attrs["conv_shape"]
  if Cin_w != Cin:
    raise ValueError(f"conv2d: weight Cin {Cin_w} != input Cin {Cin}")
  if pad:
    # in-graph zero padding: concat baked zero borders onto H then W. const_array (not graph.input), so a
    # standalone compile(conv2d(...)) bakes them instead of exposing phantom zero-pad input ports.
    zh = graph._const(np.zeros((N, Cin, pad, W), np.float16))
    x = graph.concat([zh, x, zh], axis=2)            # [N, Cin, H+2pad, W]
    H = H + 2 * pad
    zw = graph._const(np.zeros((N, Cin, H, pad), np.float16))
    x = graph.concat([zw, x, zw], axis=3)            # [N, Cin, H+2pad, W+2pad]
    W = W + 2 * pad
  Hout, Wout = H - kH + 1, W - kW + 1
  L, K = Hout * Wout, Cin * kH * kW
  parts = []
  for u in range(kH):
    for v in range(kW):
      # patch index on axis 2 (not last): keeps the large grad off the width axis,
      # avoiding A13's x16 crop-DMA saturation (>4094) on nonzero last-axis offsets.
      parts.append(x.slice_by_size([0, 0, u, v], [N, Cin, Hout, Wout]).reshape(N, Cin, 1, L))
  patches = graph.concat(parts, axis=2).transpose([0, 3, 1, 2]).reshape(N, L, K)   # [N,L,K]
  y = patches @ weight.reshape(1, K, Cout)                  # broadcast bmm -> [N,L,Cout]
  return y.transpose([0, 2, 1]).reshape(N, Cout, Hout, Wout)

mse

mse(y: Tensor, target: Tensor) -> Tensor

Mean squared error over all axes (a scalar loss).

Source code in aneforge/autograd.py

def mse(y: Tensor, target: Tensor) -> Tensor:
  """Mean squared error over all axes (a scalar loss)."""
  diff = y - target
  return diff.square().mean(tuple(range(len(y.shape))))

adam_step

adam_step(params, m, v, grads: dict, lr_t, betas=(0.9, 0.999), eps: float = 1e-08)

One Adam update as graph ops over lists params/m/v, returning new (params, m, v); used to unroll K steps into one program.

Source code in aneforge/autograd.py

def adam_step(params, m, v, grads: dict, lr_t, betas=(0.9, 0.999), eps: float = 1e-8):
  """One Adam update as graph ops over lists `params`/`m`/`v`, returning new (params, m, v); used to unroll K steps into one program."""
  b1, b2 = betas
  nP, nM, nV = [], [], []
  for p, mi, vi in zip(params, m, v):
    w2, m2, v2 = _adam_update(p, mi, vi, grads[p], lr_t, b1, b2, eps)
    if "conv_shape" in p.attrs: w2.attrs["conv_shape"] = p.attrs["conv_shape"]
    nP.append(w2); nM.append(m2); nV.append(v2)
  return nP, nM, nV

Layer-streamed training

streaming

Layer-streamed (gradient-checkpointed) training for deep stacks of identical layers: compile one layer's forward/backward once and reuse per layer (depth-independent compile).

CheckpointedStack

A depth-independent compile for a stack of identical layers. layer_fn(params, x) builds one layer; example_params gives a layer's param shapes; io_shape is the inter-layer activation shape.

Source code in aneforge/streaming.py

class CheckpointedStack:
  """A depth-independent compile for a stack of identical layers. `layer_fn(params, x)` builds one layer; `example_params` gives a layer's param shapes; `io_shape` is the inter-layer activation shape."""

  def __init__(self, layer_fn, example_params, io_shape):
    self.io_shape = tuple(io_shape)
    self._nparam = len(example_params)

    # per-layer forward: y = layer_fn(params, x)
    self._x = _g.input(self.io_shape)
    self._p = [_ag.parameter(np.asarray(p, np.float32)) for p in example_params]
    y = layer_fn(self._p, self._x)
    if tuple(y.shape) != self.io_shape:
      raise ValueError(f"layer_fn output shape {y.shape} != io_shape {self.io_shape}")
    self._fwd = _compile(y)

    # per-layer backward: return param grads + input grad (recompute-in-backward)
    self._xb = _g.input(self.io_shape)
    self._pb = [_ag.parameter(np.asarray(p, np.float32)) for p in example_params]
    self._gout = _g.input(self.io_shape)
    yb = layer_fn(self._pb, self._xb)
    grads = _ag.backward_from(self._gout, yb, [*self._pb, self._xb])
    self._g_param = [grads[p] for p in self._pb]
    self._g_in = grads[self._xb]
    self._bwd = _compile_multi([*self._g_param, self._g_in])
    self._bwd_in = {id(t): n for t, n in self._bwd.input_ports}
    self._bwd_out = dict(self._bwd.output_ports)
    # baked-constant input ports (e.g. a causal mask) carried into the backward graph
    _fed = {id(self._xb), id(self._gout), *(id(p) for p in self._pb)}
    self._bwd_consts = [(t, n) for t, n in self._bwd.input_ports if id(t) not in _fed]

  def forward(self, layers_params, x0):
    """Run the stack; `layers_params[i]` is layer i's param arrays. Returns `(output, checkpoints)`, checkpoints[i] = layer i's input activation."""
    x = np.asarray(x0, np.float32)
    checkpoints = []
    for lp in layers_params:
      checkpoints.append(x)
      feed = {id(self._x): x.astype(_F16),
          **{id(t): np.asarray(v, _F16) for t, v in zip(self._p, lp)}}
      # other input ports are baked constants (e.g. a causal mask)
      assert isinstance(self._fwd, _Model)  # layer_fn never contains sdpa nodes
      vals = [feed[id(t)] if id(t) in feed else np.asarray(t.attrs["value"], _F16)
          for t in self._fwd._input_tensors]
      x = np.asarray(self._fwd(*vals), np.float32)
    return x, checkpoints

  def backward(self, layers_params, checkpoints, g_out):
    """Backprop the stack from `g_out`. Returns `(param_grads, g_in)`: param_grads[i] is layer i's grads, g_in the stack-input grad."""
    g = np.asarray(g_out, np.float32)
    param_grads: list[list[np.ndarray] | None] = [None] * len(layers_params)
    for i in range(len(layers_params) - 1, -1, -1):
      self._bwd.prog.set_input(self._bwd_in[id(self._gout)], g.astype(_F16))
      self._bwd.prog.set_input(self._bwd_in[id(self._xb)], checkpoints[i].astype(_F16))
      for t, v in zip(self._pb, layers_params[i]):
        self._bwd.prog.set_input(self._bwd_in[id(t)], np.asarray(v, _F16))
      for t, n in self._bwd_consts:                   # baked constants (e.g. mask)
        self._bwd.prog.set_input(n, np.asarray(t.attrs["value"], _F16))
      self._bwd.prog.execute()
      param_grads[i] = [np.asarray(self._bwd.prog.read_output(self._bwd_out[gp]), np.float32)
                for gp in self._g_param]
      g = np.asarray(self._bwd.prog.read_output(self._bwd_out[self._g_in]), np.float32)
    return param_grads, g

  def release(self):
    self._fwd.release()
    self._bwd.release()

forward

forward(layers_params, x0)

Run the stack; layers_params[i] is layer i's param arrays. Returns (output, checkpoints), checkpoints[i] = layer i's input activation.

Source code in aneforge/streaming.py

def forward(self, layers_params, x0):
  """Run the stack; `layers_params[i]` is layer i's param arrays. Returns `(output, checkpoints)`, checkpoints[i] = layer i's input activation."""
  x = np.asarray(x0, np.float32)
  checkpoints = []
  for lp in layers_params:
    checkpoints.append(x)
    feed = {id(self._x): x.astype(_F16),
        **{id(t): np.asarray(v, _F16) for t, v in zip(self._p, lp)}}
    # other input ports are baked constants (e.g. a causal mask)
    assert isinstance(self._fwd, _Model)  # layer_fn never contains sdpa nodes
    vals = [feed[id(t)] if id(t) in feed else np.asarray(t.attrs["value"], _F16)
        for t in self._fwd._input_tensors]
    x = np.asarray(self._fwd(*vals), np.float32)
  return x, checkpoints

backward

backward(layers_params, checkpoints, g_out)

Backprop the stack from g_out. Returns (param_grads, g_in): param_grads[i] is layer i's grads, g_in the stack-input grad.

Source code in aneforge/streaming.py

def backward(self, layers_params, checkpoints, g_out):
  """Backprop the stack from `g_out`. Returns `(param_grads, g_in)`: param_grads[i] is layer i's grads, g_in the stack-input grad."""
  g = np.asarray(g_out, np.float32)
  param_grads: list[list[np.ndarray] | None] = [None] * len(layers_params)
  for i in range(len(layers_params) - 1, -1, -1):
    self._bwd.prog.set_input(self._bwd_in[id(self._gout)], g.astype(_F16))
    self._bwd.prog.set_input(self._bwd_in[id(self._xb)], checkpoints[i].astype(_F16))
    for t, v in zip(self._pb, layers_params[i]):
      self._bwd.prog.set_input(self._bwd_in[id(t)], np.asarray(v, _F16))
    for t, n in self._bwd_consts:                   # baked constants (e.g. mask)
      self._bwd.prog.set_input(n, np.asarray(t.attrs["value"], _F16))
    self._bwd.prog.execute()
    param_grads[i] = [np.asarray(self._bwd.prog.read_output(self._bwd_out[gp]), np.float32)
              for gp in self._g_param]
    g = np.asarray(self._bwd.prog.read_output(self._bwd_out[self._g_in]), np.float32)
  return param_grads, g