def mulb_loop

bendtorch-gpu@1.3.0.0 · main.bend:120

import Base
import bend-ml-tensor-gpu@0.1.5.0/main.bend as TG

120# x1 = dy; the saved a and b, written in place: a becomes db = dy * a and b becomes da = dy * b, in one bang
121def mulb_loop(k: Nat, +i: U32, +l: U32, rd: Array<F32> & F32, ra: Array<F32> & F32, rb: Array<F32> & F32) -> TG.O2:
122  match k rd ra rb:
123    case 0n Tuple{d, +x} Tuple{a, +y} Tuple{b, +z}: TG.O2{a, b}
124    case 1n+q Tuple{d, +x} Tuple{a, +y} Tuple{b, +z}:
125      +u = (i + l : U32)
126      mulb_loop(q, u, l, Array.get(F32, d, u), Array.get(F32, Array.set(F32, a, i, (x * y : F32)), u), Array.get(F32, Array.set(F32, b, i, (x * z : F32)), u))