def batch_number_spec

bendtorch-core@0.1.0.0 · src/tensor/batched_matmul_PROOF.bend:358

import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
import bendtorch-core@0.1.0.0/src/tensor/batched_matmul.bend as Batched

358# The batch number of the build reads each operand's batch as the shape
359# module broadcasts it.
360def batch_number_spec(+first: Nat, +second: Nat, +b1: Nat, +b2: Nat) ->
361  {Batched.batch_number(first, second, b1, b2) == Nat.add(Nat.mul(Shape.broadcast_index(first, b1), second), Shape.broadcast_index(second, b2)) : Nat}:
362  %Equal.sym(Nat, Batched.stretch(b1, Nat.is_eq(first, 1n)), Shape.broadcast_index(first, b1), stretch_spec(first, b1)) :
363    {Nat.add(Nat.mul(_, second), Batched.stretch(b2, Nat.is_eq(second, 1n))) ==
364      Nat.add(Nat.mul(Shape.broadcast_index(first, b1), second), Shape.broadcast_index(second, b2)) : Nat}
365  %Equal.sym(Nat, Batched.stretch(b2, Nat.is_eq(second, 1n)), Shape.broadcast_index(second, b2), stretch_spec(second, b2)) :
366    {Nat.add(Nat.mul(Shape.broadcast_index(first, b1), second), _) ==
367      Nat.add(Nat.mul(Shape.broadcast_index(first, b1), second), Shape.broadcast_index(second, b2)) : Nat}
368  {==}