def cell_position

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

import Base

173# Which words each cell reads, written from torch.matmul's definition and
174# the shape module's broadcast, with no call into the batched code.
175#
176# Cell (b1, b2, row, column) of an output whose batch has a second extent
177# of o2 and whose matrices are rows x columns sits at its row-major
178# position, the cells of one batch after those of the batch before it.
179def cell_position(+o2: Nat, +rows: Nat, +columns: Nat, +b1: Nat, +b2: Nat, +row: Nat, +column: Nat) -> Nat:
180  Nat.add(Nat.mul(Nat.add(Nat.mul(b1, o2), b2), Nat.mul(rows, columns)), Nat.add(Nat.mul(row, columns), column))