def Mat.sum2

bendtorch-gpu@1.3.0.0 · main.bend:150

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

150# a + b, as main.bend's Mat.sum2 adds each cell, written over a (both are consumed)
151def Mat.sum2(+r: Nat, +c: Nat, a: TA.Mat<r, c>, b: TA.Mat<r, c>) -> TA.Mat<r, c>:
152  match a b:
153    case TA.Mat{x} TA.Mat{y}: TA.Mat{r1_a(TG.run1(~EN, ~lf_add, TG.depth(Nat.mul(r, c)), en(Nat.mul(r, c)), y, nn(Unit{}), nn(Unit{}), nn(Unit{}), x))}