def Embedding

bendtorch-fast@1.3.0.0 · main.bend:342

import Base
import bend-ml-tensor-array@0.1.18.0/main.bend as TA
import bendtorch@1.3.0.0/main.bend as T

342# main.bend's Embedding with its optimizer step in blocks (mat_step): the table is SmolLM2's largest matrix (49152 x 576,
343# tied to the output head), and main.bend's step updates it on one thread
344def Embedding(+n: Nat, +v: Nat, +d: Nat) -> T.Layer<TA.Mat<v, d>, List<&2, Nat>, TA.Mat<n, d>, List<&2, Nat>>:
345  T.Layer{p => x => T.emb_f(n, v, d, x, p), p => s => dy => g => T.emb_bwd(n, v, d, p, s, dy, g), op => p => g => st => mat_step(v, d, op, p, g, st), u => TA.Mat.zeros(v, d), p => T.Mat.dump(v, d, p), p => xs => off => T.Mat.restore(v, d, p, xs, off), p => xs => off => T.Mat.store(v, d, p, xs, off), u => U32.from_nat(Nat.mul(v, d))}