def left_start

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

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

260# The first term's offsets of cell `position`: batch b = position / (rows *
261# columns) splits into (b / o2, b % o2), and each operand reads its own
262# batch, a stretched axis at 0, then the cell's row and column at term 0.
263def left_start(+layout: Layout, +position: Nat) -> Nat:
264  Layout{+lo, ro, +inner, +rows, +columns, +l1, +l2, r1, r2, +o2} = layout
265  +batch = Nat.div(position, Nat.mul(rows, columns))
266  Nat.add(Nat.mul(batch_number(l1, l2, Nat.div(batch, o2), Nat.mod(batch, o2)), Nat.mul(rows, inner)),
267    Matrix.read_offset(lo, inner, rows, Nat.div(Nat.mod(position, Nat.mul(rows, columns)), columns), 0n))