def dot_broadcast

bendtorch-core@0.1.0.0 · src/shape/projection_proofs.bend:105

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

105def dot_broadcast(+dims: List<&2, Nat>, +values: List<&2, Nat>) -> {dot(strides_of(dims), values) == Shape.offset(dims, Shape.broadcast_indices(dims, values)) : Nat}:
106  match dims values:
107    case +head <> +tail +value <> +rest:
108      %Equal.sym(Nat, dot(strides_of(tail), rest), Shape.offset(tail, Shape.broadcast_indices(tail, rest)), dot_broadcast(tail, rest)) : {Nat.add(Nat.mul(Shape.broadcast_index(head, Shape.numel(tail)), value), _) == Nat.add(Nat.mul(Shape.broadcast_index(head, value), Shape.numel(tail)), Shape.offset(tail, Shape.broadcast_indices(tail, rest))) : Nat}
109      Equal.cong(Nat, Nat, term => Nat.add(term, Shape.offset(tail, Shape.broadcast_indices(tail, rest))), Nat.mul(Shape.broadcast_index(head, Shape.numel(tail)), value), Nat.mul(Shape.broadcast_index(head, value), Shape.numel(tail)), axis_term(head, Shape.numel(tail), value))
110    case Nil{} values:
111      {==}
112    case head <> tail Nil{}:
113      {==}