def operand_offset

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

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

192# The row-major offset of that term in an operand of batch extents
193# [e1, e2] read at batch (b1, b2) of the output: an axis of extent 1 is
194# broadcast and read at index 0 (`Shape.broadcast_index`, the rule of the
195# elementwise operations), any other at the output's own index.
196def operand_offset(+order: Matrix.Order, +inner: Nat, +outers: Nat, +e1: Nat, +e2: Nat, +b1: Nat, +b2: Nat, +outer: Nat, +k: Nat) -> Nat:
197  Nat.add(Nat.mul(Nat.add(Nat.mul(Shape.broadcast_index(e1, b1), e2), Shape.broadcast_index(e2, b2)), Nat.mul(outers, inner)),
198    matrix_offset(order, inner, outers, outer, k))