src/tensor/contracts/fold_records.bend

bendtorch-core@0.1.0.0 · 79 lines · 8 definitions

1import Base
2import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
3import bendtorch-core@0.1.0.0/src/tensor/core.bend as Core
4import bendtorch-core@0.1.0.0/src/tensor/matmul.bend as Matrix
5import bendtorch-core@0.1.0.0/src/tensor/contracts/fold.bend as Fold
6
7# What a fold result promises, stated only through the result: it succeeds
8# with `total`, and the returned input still holds every word. `total` and
9# `word` are the specification; an instance fills them in.
10type FoldContract<-input: Shape.Shape, -total: F32, -word: U32 -> F32,
11  -result: Core.Dynamic & Result<&2, &1, Core.Error, F32>> is Type:
12  FoldContract{
13    value: {Fold.value(result) == Some{total} : Maybe<&2, F32>},
14    kept: @+other: U32 -> @inside: {Nat.is_lt(U32.to_nat(other), Shape.logical_numel(Shape.dimensions(input))) == True{} : Bool} ->
15      {Fold.input_point(result, other) == Some{word(other)} : Maybe<&2, F32>}}
16
17# The same for a pair fold: the total, and every word of both returned inputs.
18type PairFoldContract<-left: Shape.Shape, -right: Shape.Shape, -total: F32, -left_word: U32 -> F32, -right_word: U32 -> F32,
19  -result: (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, F32>> is Type:
20  PairFoldContract{
21    value: {Fold.pair_value(result) == Some{total} : Maybe<&2, F32>},
22    left_kept: @+other: U32 -> @inside: {Nat.is_lt(U32.to_nat(other), Shape.logical_numel(Shape.dimensions(left))) == True{} : Bool} ->
23      {Fold.left_point(result, other) == Some{left_word(other)} : Maybe<&2, F32>},
24    right_kept: @+other: U32 -> @inside: {Nat.is_lt(U32.to_nat(other), Shape.logical_numel(Shape.dimensions(right))) == True{} : Bool} ->
25      {Fold.right_point(result, other) == Some{right_word(other)} : Maybe<&2, F32>}}
26
27def product_succeeded(result: Result<&2, &1, Matrix.Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>) -> Bool:
28  Result.is_done(&2, &1, Matrix.Error, Core.Dynamic & Core.Dynamic & Core.Dynamic, result)
29
30def product_extents(result: Result<&2, &1, Matrix.Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>) -> Maybe<&2, List<&2, Nat>>:
31  match result:
32    case Fail{error}:
33      None{}
34    case Done{values}:
35      (left, right, product) = values
36      Core.Packed{shape, tensor} = product
37      Some{Shape.extents(shape)}
38
39def product_point(result: Result<&2, &1, Matrix.Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>, +index: U32) -> Maybe<&2, F32>:
40  match result:
41    case Fail{error}:
42      None{}
43    case Done{values}:
44      (left, right, product) = values
45      Core.dynamic_read_value(Core.read_dynamic(product, index))
46
47def product_left_point(result: Result<&2, &1, Matrix.Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>, +index: U32) -> Maybe<&2, F32>:
48  match result:
49    case Fail{error}:
50      None{}
51    case Done{values}:
52      (left, right, product) = values
53      Core.dynamic_read_value(Core.read_dynamic(left, index))
54
55def product_right_point(result: Result<&2, &1, Matrix.Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>, +index: U32) -> Maybe<&2, F32>:
56  match result:
57    case Fail{error}:
58      None{}
59    case Done{values}:
60      (left, right, product) = values
61      Core.dynamic_read_value(Core.read_dynamic(right, index))
62
63# What a matrix product promises, stated only through the result: it
64# succeeds with a rows x columns output whose cell (row, column), at index
65# row * columns + column, holds `cell(row, column)`, and both returned inputs
66# still hold every word.
67type MatmulContract<-left: Shape.Shape, -right: Shape.Shape, -rows: Nat, -columns: Nat, -cell: Nat -> Nat -> F32,
68  -left_word: U32 -> F32, -right_word: U32 -> F32,
69  -result: Result<&2, &1, Matrix.Error, Core.Dynamic & Core.Dynamic & Core.Dynamic>> is Type:
70  MatmulContract{
71    successful: {product_succeeded(result) == True{} : Bool},
72    extents: {product_extents(result) == Some{[rows, columns]} : Maybe<&2, List<&2, Nat>>},
73    value: @+row: Nat -> @row_below: {Nat.is_lt(row, rows) == True{} : Bool} ->
74      @+column: Nat -> @column_below: {Nat.is_lt(column, columns) == True{} : Bool} ->
75      {product_point(result, U32.from_nat(Nat.add(Nat.mul(row, columns), column))) == Some{cell(row, column)} : Maybe<&2, F32>},
76    left_kept: @+other: U32 -> @inside: {Nat.is_lt(U32.to_nat(other), Shape.logical_numel(Shape.dimensions(left))) == True{} : Bool} ->
77      {product_left_point(result, other) == Some{left_word(other)} : Maybe<&2, F32>},
78    right_kept: @+other: U32 -> @inside: {Nat.is_lt(U32.to_nat(other), Shape.logical_numel(Shape.dimensions(right))) == True{} : Bool} ->
79      {product_right_point(result, other) == Some{right_word(other)} : Maybe<&2, F32>}}