BendtorchFast.Embedding: def of bendtorch-fast@1.3.0.0
Import: import bendtorch-fast@1.3.0.0/main.bend as BendtorchFast
Lines 342-345 of bendtorch-fast@1.3.0.0/main.bend: https://jend.leeeet.dev/read.txt?path=bendtorch-fast@1.3.0.0/main.bend&start_line=342&end_line=345

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

# main.bend's Embedding with its optimizer step in blocks (mat_step): the table is SmolLM2's largest matrix (49152 x 576,
# tied to the output head), and main.bend's step updates it on one thread
def Embedding(+n: Nat, +v: Nat, +d: Nat) -> T.Layer<TA.Mat<v, d>, List<&2, Nat>, TA.Mat<n, d>, List<&2, Nat>>:
  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))}

Uses (13):
  https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=List
  https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=Nat
  https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=Nat.mul
  https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=U32.from_nat
  https://jend.leeeet.dev/definition.txt?path=bend-ml-tensor-array@0.1.19.0/main.bend&name=Mat
  https://jend.leeeet.dev/definition.txt?path=bend-ml-tensor-array@0.1.19.0/main.bend&name=Mat.zeros
  https://jend.leeeet.dev/definition.txt?path=bendtorch-fast@1.3.0.0/main.bend&name=mat_step
  https://jend.leeeet.dev/definition.txt?path=bendtorch@1.3.0.0/main.bend&name=Layer
  https://jend.leeeet.dev/definition.txt?path=bendtorch@1.3.0.0/main.bend&name=Mat.dump
  https://jend.leeeet.dev/definition.txt?path=bendtorch@1.3.0.0/main.bend&name=Mat.restore
  https://jend.leeeet.dev/definition.txt?path=bendtorch@1.3.0.0/main.bend&name=Mat.store
  https://jend.leeeet.dev/definition.txt?path=bendtorch@1.3.0.0/main.bend&name=emb_bwd
  https://jend.leeeet.dev/definition.txt?path=bendtorch@1.3.0.0/main.bend&name=emb_f