def pair_join_eq

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

import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
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

794# The depth-1 tree: the leaf of cell `base` and the leaf of the cell after
795# it, which reads the state the first hands back.
796def pair_join_eq(+plan: Shape.BroadcastPlan, +rows: List<&2, Shape.Row>,
797  +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>},
798  +shape: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>, +h: List<&2, Api.Effect>, +d: Nat, +f: F32, +count: Nat, +base: Nat, x: F32, high: F32 & Array<F32> & Unit,
799  ex: {ALeaf{x} == out_spec(plan, Shape.Indexed{rows, rows_exact}, shape, fill, ops, count, 0n, base) : Array<F32>},
800  leaf_facts: TreeEq(plan, rows, rows_exact, shape, fill, ops, h, d, f, count, 0n, Nat.add(1n, base), Parallel.leaf(~(Array<F32> & Unit), high))) ->
801  TreeEq(plan, rows, rows_exact, shape, fill, ops, h, d, f, count, 1n, base, Parallel.pair_join(~(Array<F32> & Unit), x, high)):
802  (y, inputs) = high
803  (ey, kept) = leaf_facts
804  (%Equal.sym(Array<F32>, ALeaf{x}, out_spec(plan, Shape.Indexed{rows, rows_exact}, shape, fill, ops, count, 0n, base), ex) : {ANode{_, ALeaf{y}} == ANode{out_spec(plan, Shape.Indexed{rows, rows_exact}, shape, fill, ops, count, 0n, base), out_spec(plan, Shape.Indexed{rows, rows_exact}, shape, fill, ops, count, 0n, Nat.add(1n, base))} : Array<F32>}
805   %Equal.sym(Array<F32>, ALeaf{y}, out_spec(plan, Shape.Indexed{rows, rows_exact}, shape, fill, ops, count, 0n, Nat.add(1n, base)), ey) : {ANode{out_spec(plan, Shape.Indexed{rows, rows_exact}, shape, fill, ops, count, 0n, base), _} == ANode{out_spec(plan, Shape.Indexed{rows, rows_exact}, shape, fill, ops, count, 0n, base), out_spec(plan, Shape.Indexed{rows, rows_exact}, shape, fill, ops, count, 0n, Nat.add(1n, base))} : Array<F32>}
806   {==}, kept)