def torch_output

bendtorch-core@0.1.0.0 · src/tensor/batched_matmul_LAWS.bend:356

import Base

356# The output: the broadcast batch, then the rows of the promoted left
357# operand unless it was a vector, then the columns of the promoted right
358# operand unless it was a vector.
359def torch_output(+left: List<&2, Nat>, +right: List<&2, Nat>) -> List<&2, Nat>:
360  List.append(&2, Nat, broadcast_batch(batch_of(promote_left(left)), batch_of(promote_right(right))),
361    List.append(&2, Nat, kept(vector(left), rows_of(promote_left(left))), kept(vector(right), columns_of(promote_right(right)))))