def broadcast_checked_complete
bendtorch-core@0.1.0.0 · src/shape/plan_proofs.bend:178
import bendtorch-core@0.1.0.0/src/shape/plan_proofs.bend as PlanProofs
import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
import bendtorch-core@0.1.0.0/src/shape/nat_proofs.bend as Nats
import bendtorch-core@0.1.0.0/src/shape/unique_proofs.bend as Unique
import bendtorch-core@0.1.0.0/src/shape/LAWS.bend as Laws
178def broadcast_checked_complete(+left: Shape.Shape, +right: Shape.Shape, +output: List<&2, U32>, +alignment: Shape.Broadcast<Shape.extents(left), Shape.extents(right), Shape.naturals(output)>, +rank_bound: Shape.LE(Shape.rank_dimensions(output), 4n), +capacity_bound: Shape.LE(Shape.logical_numel(output), 1048576n)) -> {Laws.broadcast_built(Shape.broadcast_checked(left, right)) == True{} : Bool}:
179 match alignment:
180 case Shape.BAlign{+left_padding, +right_padding, +minimal, +aligned}:
181 +left_rank = Shape.rank_dimensions(Shape.dimensions(left))
182 +right_rank = Shape.rank_dimensions(Shape.dimensions(right))
183 +left_extra = Nat.sub(right_rank, left_rank)
184 +right_extra = Nat.sub(left_rank, right_rank)
185 +balanced = Equal.trans(Nat, Nat.add(left_padding, left_rank), Nat.add(left_padding, List.length(&2, Nat, Shape.extents(left))), Nat.add(right_padding, right_rank), Equal.cong(Nat, Nat, value => Nat.add(left_padding, value), left_rank, List.length(&2, Nat, Shape.extents(left)), Shape.rank_naturals(Shape.dimensions(left))), Equal.trans(Nat, Nat.add(left_padding, List.length(&2, Nat, Shape.extents(left))), Nat.add(right_padding, List.length(&2, Nat, Shape.extents(right))), Nat.add(right_padding, right_rank), Unique.balance(left_padding, right_padding, Shape.extents(left), Shape.extents(right), Shape.naturals(output), aligned), Equal.sym(Nat, Nat.add(right_padding, right_rank), Nat.add(right_padding, List.length(&2, Nat, Shape.extents(right))), Equal.cong(Nat, Nat, value => Nat.add(right_padding, value), right_rank, List.length(&2, Nat, Shape.extents(right)), Shape.rank_naturals(Shape.dimensions(right))))))
186 +computed = sub_balance(left_rank, right_rank)
187 +same_left = Unique.left_padding_equal(left_padding, right_padding, left_extra, right_extra, left_rank, right_rank, minimal, Shape.minimal(left_rank, right_rank), balanced, computed)
188 +same_right = Unique.cancel_right(right_padding, right_extra, right_rank, Equal.trans(Nat, Nat.add(right_padding, right_rank), Nat.add(left_padding, left_rank), Nat.add(right_extra, right_rank), Equal.sym(Nat, Nat.add(left_padding, left_rank), Nat.add(right_padding, right_rank), balanced), Equal.trans(Nat, Nat.add(left_padding, left_rank), Nat.add(left_extra, left_rank), Nat.add(right_extra, right_rank), Equal.cong(Nat, Nat, value => Nat.add(value, left_rank), left_padding, left_extra, same_left), computed)))
189 +retyped = Unique.retype(left_extra, left_padding, right_extra, right_padding, Shape.extents(left), Shape.extents(right), Shape.naturals(output), Equal.sym(Nat, left_padding, left_extra, same_left), Equal.sym(Nat, right_padding, right_extra, same_right), aligned)
190 +seen = align_complete(Shape.pad_words(left_extra, Shape.dimensions(left)), Shape.pad_words(right_extra, Shape.dimensions(right)), Shape.naturals(output), padded_words(left_extra, right_extra, Shape.dimensions(left), Shape.dimensions(right), Shape.naturals(output), retyped))
191 complete_align(left, right, left_extra, right_extra, Shape.minimal(left_rank, right_rank), Shape.align(Shape.pad_words(left_extra, Shape.dimensions(left)), Shape.pad_words(right_extra, Shape.dimensions(right))), Shape.naturals(output), seen, Nats.le_rewrite(List.length(&2, Nat, Shape.naturals(output)), Shape.rank_dimensions(output), 4n, Shape.rank_naturals(output), rank_bound), capacity_bound)