def batched_record

bendtorch-core@0.1.0.0 · src/autograd/tensor_backward.bend:1214

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

1214# torch.matmul on the tape: two rank-2 operands record `Matmul`, the
1215# product `matmul` records, and every other pair records `BatchedMatmul`.
1216def batched_record(+a: Nat, +b: Nat, +left_shape: Shape.Shape, +right_shape: Shape.Shape, entries: List<&1, Entry>,
1217  left: Core.Dynamic, right: Core.Dynamic, matrices: Bool) -> Tape & Result<&2, &2, Error, Nat>:
1218  match matrices:
1219    case True{}:
1220      multiplied_record(Matmul{a, b, left_shape, right_shape}, entries, Matrix.matmul(left, right))
1221    case False{}:
1222      multiplied_record(BatchedMatmul{a, b, left_shape, right_shape}, entries, Batched.matmul(left, right))