def tree_eq

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

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/contracts/matmul_parallel_proof.bend as Par
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

821def tree_eq(+plan: Shape.BroadcastPlan, +rows: List<&2, Shape.Row>,
822  +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>},
823  +shape: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>, +h: List<&2, Api.Effect>, +d: Nat, +f: F32, +count: Nat,
824  +source: Par.Source<shape, fill, ops, h, d, f>, +facts: Facts<plan, shape, count>, +depth: Nat, +base: Nat) ->
825  TreeEq(plan, rows, rows_exact, shape, fill, ops, h, d, f, count, depth, base,
826    Parallel.tree(~Reduce.Fibers, ~Array<F32>, ~Unit, ~Reduce.cell_at, Reduce.fibers(Reduce.Fibers{rows, Shape.extents(Shape.output_shape(plan))}), depth, count, base, (Api.arr(h, d, f), Unit{}))):
827  match depth:
828    case 0n:
829      cell_eq(plan, rows, rows_exact, shape, fill, ops, h, d, f, count, source, facts, base, Nat.is_lt(base, count), {==})
830    case 1n+lower:
831      match lower:
832        case 0n:
833          pair_eq(plan, rows, rows_exact, shape, fill, ops, h, d, f, count, source, facts, base, Reduce.cell_at(Reduce.fibers(Reduce.Fibers{rows, Shape.extents(Shape.output_shape(plan))}), base, (Api.arr(h, d, f), Unit{}), Nat.is_lt(base, count)),
834            cell_eq(plan, rows, rows_exact, shape, fill, ops, h, d, f, count, source, facts, base, Nat.is_lt(base, count), {==}))
835        case 1n+below:
836          tree_node_eq(plan, rows, rows_exact, shape, fill, ops, h, d, f, count, lower, base,
837            Parallel.tree(~Reduce.Fibers, ~Array<F32>, ~Unit, ~Reduce.cell_at, Reduce.fibers(Reduce.Fibers{rows, Shape.extents(Shape.output_shape(plan))}), lower, count, base, (Api.arr(h, d, f), Unit{})),
838            tree_eq(plan, rows, rows_exact, shape, fill, ops, h, d, f, count, source, facts, lower, base),
839            tree_eq(plan, rows, rows_exact, shape, fill, ops, h, d, f, count, source, facts, lower, Nat.add(Model.capacity(lower), base)))