def lf_mse

bendtorch-gpu@1.3.0.0 · main.bend:1196

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

1196# x1 = the target; y written in place with the gradient, ps[t] = leaf t's block sum (one leaf a block)
1197def lf_mse(+t: U32, +l: U32, +e: EMS, x1: Array<F32>, x2: Array<F32>, x3: Array<F32>, x4: Array<F32>, o: TG.O2) -> TG.O2:
1198  match e o:
1199    case EMS{+nb, +p, +bs, +kk} TG.O2{y, ps}: ms_if(U32.is_lt(t, nb), t, p, bs, kk, x1, y, ps)