def zip

bendtorch-fast@1.3.0.0 · main.bend:204

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

204# a[i] = f(a[i], b[i]); b comes back unchanged
205def zip(~f: @+a: F32 -> @+b: F32 -> F32, +r: Nat, +c: Nat, a: TA.Mat<r, c>, b: TA.Mat<r, c>) -> TA.MPair<r, c>:
206  match a b:
207    case TA.Mat{x} TA.Mat{y}:
208      zip_fin(r, c, TF.par2(~Unit, ~F32, ~F32, ~(e => i0 => k => q => ew2_leaf(~f, e, i0, k, q)), Unit{}, ew_depth(Nat.mul(r, c)), Nat.mul(r, c), x, y))