def scatter_l

bendtorch@1.3.0.0 · tensor.bend:189

import Base

189# g[m(j)] = g[m(j)] + x[j] for k positions from j; rx holds x[j], rg holds g[m(j)]
190def scatter_l(k: Nat, +j: U32, +m: IMap, rx: Array<F32> & F32, rg: Array<F32> & F32) -> P2:
191  match k rx rg:
192    case 0n Tuple{x, u} Tuple{g, w}:
193      P2{x, g}
194    case 1n+p Tuple{x, +u} Tuple{g, +w}:
195      scatter_l(p, (j + 1 : U32), m, Array.get(F32, x, (j + 1 : U32)), Array.get(F32, Array.set(F32, g, imap(m, j), (w + u : F32)), imap(m, (j + 1 : U32))))