Bendtorch.adamh_leaf: def of bendtorch@1.3.0.0
Import: import bendtorch@1.3.0.0/main.bend as Bendtorch
Lines 394-402 of bendtorch@1.3.0.0/main.bend: https://jend.leeeet.dev/read.txt?path=bendtorch@1.3.0.0/main.bend&start_line=394&end_line=402

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

# one matrix's step on its leaf; anything else leaves the values untouched, as Opt.mat's fallback
def adamh_leaf(+r: Nat, +c: Nat, +o: Opt, wa: Array<F32>, ga: Array<F32>, t: PTree) -> MatStep<r, c>:
  match o t:
    case Adam{lr, b1, b2, eps, wd, dec, step, bc2, c1, c2} PLeaf{+n, ma, va}:
      ph_fin(r, c, n, adamh_l(Nat.mul(r, c), 0, lr, c1, b2, c2, eps, wd, dec, step, bc2, Array.get(F32, wa, 0), Array.get(F32, ga, 0), Array.get(U32, ma, 0), Array.get(U32, va, 0)))
    case Adam{lr, b1, b2, eps, wd, dec, step, bc2, c1, c2} PLeafV{+n, ma, va}:
      phv_fin(r, c, n, adamhv_l(Nat.mul(r, c), 0, lr, c1, b2, c2, eps, wd, dec, step, bc2, Array.get(F32, wa, 0), Array.get(F32, ga, 0), Array.get(U32, ma, 0), Array.get(F32, va, 0)))
    case _ t:
      MatStep{TA.Mat{wa}, TA.Mat{ga}, Packed{t}}

Uses (15):
  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=F32
  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
  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=bendtorch@1.3.0.0/main.bend&name=MatStep
  https://jend.leeeet.dev/definition.txt?path=bendtorch@1.3.0.0/main.bend&name=Opt
  https://jend.leeeet.dev/definition.txt?path=bendtorch@1.3.0.0/main.bend&name=OptState
  https://jend.leeeet.dev/definition.txt?path=bendtorch@1.3.0.0/main.bend&name=PTree
  https://jend.leeeet.dev/definition.txt?path=bendtorch@1.3.0.0/main.bend&name=adamh_l
  https://jend.leeeet.dev/definition.txt?path=bendtorch@1.3.0.0/main.bend&name=adamhv_l
  https://jend.leeeet.dev/definition.txt?path=bendtorch@1.3.0.0/main.bend&name=ph_fin
  https://jend.leeeet.dev/definition.txt?path=bendtorch@1.3.0.0/main.bend&name=phv_fin