BendtorchGpu.mom_loop: def of bendtorch-gpu@1.3.0.0 Import: import bendtorch-gpu@1.3.0.0/main.bend as BendtorchGpu Lines 220-226 of bendtorch-gpu@1.3.0.0/main.bend: https://jend.leeeet.dev/read.txt?path=bendtorch-gpu@1.3.0.0/main.bend&start_line=220&end_line=226 import Base import bend-ml-tensor-gpu@0.1.5.0/main.bend as TG def mom_loop(k: Nat, +i: U32, +l: U32, +lr: F32, +mu: F32, rw: Array & F32, rg: Array & F32, rm: Array & F32) -> TG.O3: match k rw rg rm: case 0n Tuple{w, +x} Tuple{g, +y} Tuple{m, +b}: TG.O3{w, g, m} case 1n+q Tuple{w, +x} Tuple{g, +y} Tuple{m, +b}: +nb = ((b * mu : F32) + y : F32) +u = (i + l : U32) mom_loop(q, u, l, lr, mu, Array.get(F32, Array.set(F32, w, i, (x + (F32.neg(lr) * nb : F32) : F32)), u), Array.get(F32, Array.set(F32, g, i, 0.0), u), Array.get(F32, Array.set(F32, m, i, nb), u)) Uses (13): https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=Array https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=Array.get https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=Array.set https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=F32 https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=F32.add https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=F32.mul https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=F32.neg https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=Nat https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=Pair https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=Sigma https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=U32 https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=U32.add https://jend.leeeet.dev/definition.txt?path=bend-ml-tensor-gpu@0.1.5.1/main.bend&name=O3