law gather_build_cells

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

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

237# For a balanced input array A, a balanced index array I and every split,
238# the slot of the forward build that holds output cell (o, k, r), below
239# the count, holds input cell (o, j, r), where j is the entry the cell
240# reads: the word PyTorch's index_select and gather put there.
241law gather_build_cells:
242  for +depth: Nat
243  for +splits: Nat
244  for +kind: Gather.Lookup
245  for +outer: Nat
246  for +rows: Nat
247  for +inner: Nat
248  for +count: Nat
249  for +cells: Nat
250  for +base: Nat
251  for +da: Nat
252  for -A: Array<F32>
253  for a: Array<F32>
254  for same: {a == A : Array<F32>}
255  for valid: Model.Balanced(F32, da, A)
256  for +di: Nat
257  for -I: Array<U32>
258  for index: Array<U32>
259  for same_index: {index == I : Array<U32>}
260  for index_valid: Model.Balanced(U32, di, I)
261  for +path: Model.Path(depth)
262  for +o: Nat
263  for +k: Nat
264  for +r: Nat
265  for k_below: {Nat.is_lt(k, count) == True{} : Bool}
266  for r_below: {Nat.is_lt(r, inner) == True{} : Bool}
267  for at: {Nat.add(Model.offset(depth, path), base) == position(count, inner, o, k, r) : Nat}
268  for in_count: {Nat.is_lt(Nat.add(Model.offset(depth, path), base), cells) == True{} : Bool}
269  for +j: U32
270  for read: {Model.observe(U32, di, I, Refinement.route(di, U32.from_nat(entry_at(kind, count, inner, o, k, r)))) == Some{j} : Maybe<&2, U32>}
271  Holds<{Model.observe(F32, depth, Gather.build(depth, splits, Gather.Plan{kind, outer, rows, inner, count}, cells, base, a, index), path) ==
272    Model.observe(F32, da, A, Refinement.route(da, U32.from_nat(position(rows, inner, o, U32.to_nat(j), r)))) : Maybe<&2, F32>}>