def stretched

bendtorch-core@0.1.0.0 · src/tensor/batched_matmul.bend:75

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

75# One batch axis: equal extents, or an extent of 1 that stretches to the
76# other one.
77def stretched(+l: U32, +r: U32, same: Bool, left_one: Bool, right_one: Bool) -> Result<&2, &2, Matrix.Error, U32>:
78  match same left_one right_one:
79    case True{} _ _:
80      Done{l}
81    case False{} True{} _:
82      Done{r}
83    case False{} False{} True{}:
84      Done{l}
85    case False{} False{} False{}:
86      Fail{Matrix.ShapeError{Shape.IncompatibleBroadcast{}}}