law gather_scatter_cells

bendtorch-core@0.1.0.0 · src/tensor/gather_LAWS.bend:305

import Base
import bendtorch-core@0.1.0.0/src/storage/model.bend as Model
import bendtorch-core@0.1.0.0/src/tensor/gather.bend as Gather

305# For a balanced gradient array G, a balanced index array I and every split,
306# the slot of the backward build that holds input cell (o, i, r), below the
307# count, holds the gradient words of the cells (o, k, r) whose entry is i,
308# added to +0.0 in increasing k, and +0.0 where no entry is i: PyTorch's
309# one-thread index_add and scatter_add.
310law gather_scatter_cells:
311  for +depth: Nat
312  for +splits: Nat
313  for +kind: Gather.Lookup
314  for +outer: Nat
315  for +rows: Nat
316  for +inner: Nat
317  for +count: Nat
318  for +cells: Nat
319  for +base: Nat
320  for +dg: Nat
321  for -G: Array<F32>
322  for g: Array<F32>
323  for same: {g == G : Array<F32>}
324  for valid: Model.Balanced(F32, dg, G)
325  for +di: Nat
326  for -I: Array<U32>
327  for index: Array<U32>
328  for same_index: {index == I : Array<U32>}
329  for index_valid: Model.Balanced(U32, di, I)
330  for +path: Model.Path(depth)
331  for +o: Nat
332  for +i: Nat
333  for +r: Nat
334  for i_below: {Nat.is_lt(i, rows) == True{} : Bool}
335  for r_below: {Nat.is_lt(r, inner) == True{} : Bool}
336  for at: {Nat.add(Model.offset(depth, path), base) == position(rows, inner, o, i, r) : Nat}
337  for in_count: {Nat.is_lt(Nat.add(Model.offset(depth, path), base), cells) == True{} : Bool}
338  Summed<depth, Gather.scatter_build(depth, splits, Gather.Plan{kind, outer, rows, inner, count}, cells, base, g, index), path,
339    dg, G, di, I, kind, count, inner, o, i, r>