type MatmulPairContract

bendtorch-core@0.1.0.0 · src/autograd/tensor_vjp_LAWS.bend:264

import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
import bendtorch-core@0.1.0.0/src/autograd/tensor_backward.bend as Backward
import bendtorch-core@0.1.0.0/src/autograd/tensor_backward_LAWS.bend as TensorBackward

264# What the Matmul VJP promises: the first contribution goes to node a with
265# left's dimensions and holds left_cell(row, column) at row * left_columns +
266# column for every row < left_rows and column < left_columns; the second goes
267# to node b with right's dimensions and holds right_cell(row, column) likewise.
268type MatmulPairContract<-a: Nat, -left: Shape.Shape, -left_rows: Nat, -left_columns: Nat, -left_cell: Nat -> Nat -> F32,
269  -b: Nat, -right: Shape.Shape, -right_rows: Nat, -right_columns: Nat, -right_cell: Nat -> Nat -> F32, -pass: Backward.Pass> is Type:
270  MatmulPairContract{
271    first_target: {TensorBackward.sole_target(first_of(pass)) == Some{a} : Maybe<&2, Nat>},
272    first_dimensions: {TensorBackward.sole_dimensions(first_of(pass)) == Some{Shape.dimensions(left)} : Maybe<&2, List<&2, U32>>},
273    first_value: @+row: Nat -> @row_below: {Nat.is_lt(row, left_rows) == True{} : Bool} ->
274      @+column: Nat -> @column_below: {Nat.is_lt(column, left_columns) == True{} : Bool} ->
275      {TensorBackward.sole_point(first_of(pass), U32.from_nat(Nat.add(Nat.mul(row, left_columns), column))) ==
276        Some{left_cell(row, column)} : Maybe<&2, F32>},
277    second_target: {TensorBackward.sole_target(second_of(pass)) == Some{b} : Maybe<&2, Nat>},
278    second_dimensions: {TensorBackward.sole_dimensions(second_of(pass)) == Some{Shape.dimensions(right)} : Maybe<&2, List<&2, U32>>},
279    second_value: @+row: Nat -> @row_below: {Nat.is_lt(row, right_rows) == True{} : Bool} ->
280      @+column: Nat -> @column_below: {Nat.is_lt(column, right_columns) == True{} : Bool} ->
281      {TensorBackward.sole_point(second_of(pass), U32.from_nat(Nat.add(Nat.mul(row, right_columns), column))) ==
282        Some{right_cell(row, column)} : Maybe<&2, F32>}}