def batched

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

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

141def batched(+rank: Nat, +rows: U32, +inner: U32, +columns: U32, +l1: U32, +l2: U32, +r1: U32, +r2: U32,
142  +lv: Bool, +rv: Bool, first: Result<&2, &2, Matrix.Error, U32>,
143  second: Result<&2, &2, Matrix.Error, U32>) -> Result<&2, &2, Matrix.Error, Plan>:
144  match first second:
145    case Fail{error} second:
146      Fail{error}
147    case Done{o1} Fail{error}:
148      Fail{error}
149    case Done{+o1} Done{+o2}:
150      planned_shape(batch_dims(rank, o1, o2), rows, inner, columns, l1, l2, r1, r2, o1, o2, lv, rv,
151        List.append(&2, U32, batch_dims(rank, o1, o2), matrix_dims(rows, lv, columns, rv)),
152        Shape.construct(List.append(&2, U32, batch_dims(rank, o1, o2), matrix_dims(rows, lv, columns, rv))))