def tree_node_eq

bendtorch-core@0.1.0.0 · src/autograd/reduce_replayed_proof.bend:776

import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
import bendtorch-core@0.1.0.0/src/storage/model.bend as Model
import bendtorch-core@0.1.0.0/src/storage/LAWS.bend as StorageLaws
import bendtorch-core@0.1.0.0/src/storage/api.bend as Api
import bendtorch-core@0.1.0.0/src/tensor/parallel.bend as Parallel
import bendtorch-core@0.1.0.0/src/autograd/reduce.bend as Reduce

776def tree_node_eq(+plan: Shape.BroadcastPlan, +rows: List<&2, Shape.Row>,
777  +rows_exact: {rows == Shape.index_table(Shape.extents(Shape.left_shape(plan)), Shape.extents(Shape.right_shape(plan)), Shape.extents(Shape.output_shape(plan))) : List<&2, Shape.Row>},
778  +shape: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>, +h: List<&2, Api.Effect>, +d: Nat, +f: F32, +count: Nat, +lower: Nat, +base: Nat,
779  low: Array<F32> & (Array<F32> & Unit),
780  facts: TreeEq(plan, rows, rows_exact, shape, fill, ops, h, d, f, count, lower, base, low),
781  next: TreeEq(plan, rows, rows_exact, shape, fill, ops, h, d, f, count, lower, Nat.add(Model.capacity(lower), base),
782    Parallel.tree(~Reduce.Fibers, ~Array<F32>, ~Unit, ~Reduce.cell_at, Reduce.fibers(Reduce.Fibers{rows, Shape.extents(Shape.output_shape(plan))}), lower, count, Nat.add(Model.capacity(lower), base), (Api.arr(h, d, f), Unit{})))) ->
783  TreeEq(plan, rows, rows_exact, shape, fill, ops, h, d, f, count, 1n+lower, base,
784    Parallel.node(~Array<F32>, ~Unit, l => r => Parallel.tree(~Reduce.Fibers, ~Array<F32>, ~Unit, ~Reduce.cell_at, Reduce.fibers(Reduce.Fibers{rows, Shape.extents(Shape.output_shape(plan))}), lower, count, Nat.add(Model.capacity(lower), base), (l, r)), low)):
785  (a, state) = low
786  (array, none) = state
787  (ea, ek) = facts
788  %Equal.sym(Array<F32> & Unit, (array, none), (Api.arr(h, d, f), Unit{}), ek) :
789    TreeEq(plan, rows, rows_exact, shape, fill, ops, h, d, f, count, 1n+lower, base,
790      Parallel.join(~(Array<F32> & Unit), a, Parallel.tree(~Reduce.Fibers, ~Array<F32>, ~Unit, ~Reduce.cell_at, Reduce.fibers(Reduce.Fibers{rows, Shape.extents(Shape.output_shape(plan))}), lower, count, Nat.add(Model.capacity(lower), base), _)))
791  tree_join_eq(plan, rows, rows_exact, shape, fill, ops, h, d, f, count, lower, base, a,
792    Parallel.tree(~Reduce.Fibers, ~Array<F32>, ~Unit, ~Reduce.cell_at, Reduce.fibers(Reduce.Fibers{rows, Shape.extents(Shape.output_shape(plan))}), lower, count, Nat.add(Model.capacity(lower), base), (Api.arr(h, d, f), Unit{})), ea, next)