def axis_right

bendtorch-core@0.1.0.0 · src/shape/broadcast_proofs.bend:31

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

31def axis_right(+left: Nat, +right: Nat, +output: Nat, +axis: Shape.Axis<left, right, output>, +index: Shape.Fin(output)) -> {Shape.Fin.value(right, Shape.Axis.project_right(left, right, output, axis, index)) == Shape.broadcast_index(right, Shape.Fin.value(output, index)) : Nat}:
32  match axis:
33    case Shape.AxisEqual{left_eq, right_eq}:
34      %h@right_eq : {Shape.Fin.value(_, Shape.Axis.project_right(left, _, output, Shape.AxisEqual{left_eq, h}, index)) == Shape.broadcast_index(_, Shape.Fin.value(output, index)) : Nat}
35      fin_broadcast(output, index)
36    case Shape.AxisLeftOne{one, right_eq}:
37      %h@right_eq : {Shape.Fin.value(_, Shape.Axis.project_right(left, _, output, Shape.AxisLeftOne{one, h}, index)) == Shape.broadcast_index(_, Shape.Fin.value(output, index)) : Nat}
38      fin_broadcast(output, index)
39    case Shape.AxisRightOne{left_eq, one}:
40      %h@one : {Shape.Fin.value(_, Shape.Axis.project_right(left, _, output, Shape.AxisRightOne{left_eq, h}, index)) == Shape.broadcast_index(_, Shape.Fin.value(output, index)) : Nat}
41      {==}