def ProductWith

bendtorch-core@0.1.0.0 · src/tensor/contracts/LAWS.bend:282

import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
import bendtorch-core@0.1.0.0/src/storage/LAWS.bend as StorageLaws
import bendtorch-core@0.1.0.0/src/tensor/core.bend as Core
import bendtorch-core@0.1.0.0/src/tensor/lifecycle.bend as Lifecycle
import bendtorch-core@0.1.0.0/src/tensor/matmul.bend as Matrix
import bendtorch-core@0.1.0.0/src/tensor/dot_fold.bend as DotFold
import bendtorch-core@0.1.0.0/src/tensor/contracts/replayed.bend as Replayed

282# What the transport form of product_replayed asks of P, as MatmulWith asks
283# it of a matmul: P holds for both inputs, replayed with any histories that
284# keep their words, next to the output, of the plan's shape, replayed over
285# +0.0 with any history whose cell (row, column) holds the total
286# `product_ordered` states.
287def ProductWith(+lo: Matrix.Order, +ro: Matrix.Order, +left: Shape.Shape, +right: Shape.Shape, +plan: Matrix.Product<lo, ro, left, right>,
288  +left_fill: F32, +left_ops: List<&2, StorageLaws.Op>, +right_fill: F32, +right_ops: List<&2, StorageLaws.Op>,
289  -P: Result<&2, &1, Matrix.Error, Core.Dynamic & Core.Dynamic & Core.Dynamic> -> Type) -> Type:
290  @+left_kept: List<&2, StorageLaws.Op> ->
291  @left_words: {words(left, left_fill, left_kept, Shape.logical_numel(Shape.dimensions(left))) ==
292    words(left, left_fill, left_ops, Shape.logical_numel(Shape.dimensions(left))) : List<&2, F32>} ->
293  @-left_fits: {Lifecycle.logical_length(Replayed.storage(left, left_fill, left_kept)) == Shape.logical_numel(Shape.dimensions(left)) : Nat} ->
294  @+right_kept: List<&2, StorageLaws.Op> ->
295  @right_words: {words(right, right_fill, right_kept, Shape.logical_numel(Shape.dimensions(right))) ==
296    words(right, right_fill, right_ops, Shape.logical_numel(Shape.dimensions(right))) : List<&2, F32>} ->
297  @-right_fits: {Lifecycle.logical_length(Replayed.storage(right, right_fill, right_kept)) == Shape.logical_numel(Shape.dimensions(right)) : Nat} ->
298  @+output_ops: List<&2, StorageLaws.Op> ->
299  @cells: (@+row: Nat -> @row_below: {Nat.is_lt(row, product_rows(lo, ro, left, right, plan)) == True{} : Bool} ->
300    @+column: Nat -> @column_below: {Nat.is_lt(column, product_columns(lo, ro, left, right, plan)) == True{} : Bool} ->
301    {Replayed.value(Matrix.product_shape(lo, ro, left, right, plan), 0.0, output_ops,
302      U32.from_nat(Nat.add(Nat.mul(row, product_columns(lo, ro, left, right, plan)), column))) ==
303      DotFold.dot_fold(product_factors(row, column, lo, ro, product_inner(lo, ro, left, right, plan),
304        product_rows(lo, ro, left, right, plan), product_columns(lo, ro, left, right, plan),
305        left, left_fill, left_ops, right, right_fill, right_ops, product_inner(lo, ro, left, right, plan)), 0.0) : F32}) ->
306  @-output_fits: {Lifecycle.logical_length(Replayed.storage(Matrix.product_shape(lo, ro, left, right, plan), 0.0, output_ops)) ==
307    Shape.logical_numel(Shape.dimensions(Matrix.product_shape(lo, ro, left, right, plan))) : Nat} ->
308    P(Done{(Core.Packed{left, Core.TTensor{Replayed.storage(left, left_fill, left_kept), left_fits}},
309      Core.Packed{right, Core.TTensor{Replayed.storage(right, right_fill, right_kept), right_fits}},
310      Core.Packed{Matrix.product_shape(lo, ro, left, right, plan),
311        Core.TTensor{Replayed.storage(Matrix.product_shape(lo, ro, left, right, plan), 0.0, output_ops), output_fits}})})