src/tensor/contracts/dot.bend

bendtorch-core@0.1.0.0 · 28 lines · 6 definitions

1import Base
2import bendtorch-core@0.1.0.0/src/tensor/core.bend as Core
3import bendtorch-core@0.1.0.0/src/tensor/contracts/fold.bend as Fold
4
5# Ordered-fold family, dot instance: output cell (row, column) of a row-major
6# product of a rows x inner and an inner x columns matrix folds
7# left[row, k] * right[k, column] for k = 0, 1, ..., inner - 1.
8
9type Cell is Data:
10  Cell{row: Nat, column: Nat, inner: Nat, columns: Nat}
11
12def multiply_add(operation: Unit, accumulator: F32, left: F32, right: F32) -> F32:
13  F32.add(accumulator, F32.mul(left, right))
14
15def left_index(cell: Cell, k: Nat) -> Nat:
16  Cell{row, column, inner, columns} = cell
17  Nat.add(Nat.mul(row, inner), k)
18
19def right_index(cell: Cell, k: Nat) -> Nat:
20  Cell{row, column, inner, columns} = cell
21  Nat.add(Nat.mul(k, columns), column)
22
23def inner(cell: Cell) -> Nat:
24  Cell{row, column, inner, columns} = cell
25  inner
26
27def dot(+cell: Cell, left: Core.Dynamic, right: Core.Dynamic) -> (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, F32>:
28  Fold.pair_run(~Unit, ~multiply_add, ~Cell, ~left_index, ~right_index, Unit{}, cell, inner(cell), left, right)