def ph_leafs

bendtorch-fast@1.3.0.0 · main.bend:142

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

142def ph_leafs(+r: Nat, +c: Nat, +o: T.Opt, wa: Array<F32>, ga: Array<F32>, t: T.PTree) -> T.MatStep<r, c>:
143  match o t:
144    case T.Adam{lr, b1, b2, eps, wd, dec, step, bc2, c1, c2} T.PLeaf{+n, ma, va}:
145      ph_fin(r, c, n, TF.par4(~PE, ~F32, ~F32, ~U32, ~U32, ~ph_leaf, PE{o, n}, ew_depth(Nat.mul(r, c)), ph_nw(r, c), wa, ga, ma, va))
146    case T.Adam{lr, b1, b2, eps, wd, dec, step, bc2, c1, c2} T.PLeafV{+n, ma, va}:
147      phv_fin(r, c, n, TF.par4(~PE, ~F32, ~F32, ~U32, ~F32, ~phv_leaf, PE{o, n}, ew_depth(Nat.mul(r, c)), ph_nw(r, c), wa, ga, ma, va))
148    case _ t:
149      T.MatStep{TA.Mat{wa}, TA.Mat{ga}, T.Packed{t}}