def matrix_offset

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

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

182# Term k of row (or column) `outer` of a matrix of `outers` rows (or
183# columns) and `inner` terms: `[outers, inner]` read by rows (Kept), or
184# `[inner, outers]` read by columns (Swapped).
185def matrix_offset(+order: Matrix.Order, +inner: Nat, +outers: Nat, +outer: Nat, +k: Nat) -> Nat:
186  match order:
187    case Matrix.Kept{}:
188      Nat.add(Nat.mul(outer, inner), k)
189    case Matrix.Swapped{}:
190      Nat.add(Nat.mul(k, outers), outer)