def Embedding

bendtorch@1.3.0.0 · main.bend:2047

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

2047# a table of v rows of d numbers, looked up by n ids; an id has no gradient, so dX is the empty list
2048def Embedding(+n: Nat, +v: Nat, +d: Nat) -> Layer<TA.Mat<v, d>, List<&2, Nat>, TA.Mat<n, d>, List<&2, Nat>>:
2049  Layer{p => x => emb_f(n, v, d, x, p), p => s => dy => g => 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 => Mat.dump(v, d, p), p => xs => off => Mat.restore(v, d, p, xs, off), p => xs => off => Mat.store(v, d, p, xs, off), u => U32.from_nat(Nat.mul(v, d))}