src/tensor/reduction.bend

bendtorch-core@0.1.0.0 · 187 lines · 22 definitions

1import Base
2import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
3import bendtorch-core@0.1.0.0/src/storage/main.bend as Storage
4import bendtorch-core@0.1.0.0/src/storage/model.bend as Model
5import bendtorch-core@0.1.0.0/src/storage/refinement.bend as Refinement
6import bendtorch-core@0.1.0.0/src/tensor/core.bend as Core
7
8def scalar_shape() -> Shape.Shape:
9  Shape.SShape{Nil{}, 1n, {==}, Unit{}, Unit{}}
10
11def sequential_sum(values: List<&2, F32>, accumulator: F32) -> F32:
12  match values:
13    case Nil{}:
14      accumulator
15    case value <> tail:
16      sequential_sum(tail, F32.add(accumulator, value))
17
18def read_step(-shape: Shape.Shape,
19  rec: Core.Tensor<shape> -> F32 -> Result<&2, &1, Core.Error, Core.Tensor<shape> & F32>,
20  accumulator: F32, pair: Core.Tensor<shape> & Result<&2, &2, Core.Error, F32>) ->
21  Result<&2, &1, Core.Error, Core.Tensor<shape> & F32>:
22  (tensor, result) = pair
23  match result:
24    case Fail{error}:
25      Fail{error}
26    case Done{value}:
27      rec(tensor, F32.add(accumulator, value))
28
29def sum_decided(-shape: Shape.Shape,
30  rec: Core.Tensor<shape> -> F32 -> Result<&2, &1, Core.Error, Core.Tensor<shape> & F32>, tensor: Core.Tensor<shape>, +cursor: Nat,
31  +index: U32, accumulator: F32, decision: Shape.BoolCase<Nat.is_eq(U32.to_nat(index), cursor)>) ->
32  Result<&2, &1, Core.Error, Core.Tensor<shape> & F32>:
33  match decision:
34    case Shape.BFalse{equal}:
35      Fail{Core.StorageError{Storage.InvalidIndex{}}}
36    case Shape.BTrue{equal}:
37      read_step(shape, rec, accumulator, Core.read_flat(shape, tensor, index))
38
39def natural_sum_loop(fuel: Nat, -shape: Shape.Shape, tensor: Core.Tensor<shape>, +cursor: Nat,
40  accumulator: F32) -> Result<&2, &1, Core.Error, Core.Tensor<shape> & F32>:
41  match fuel:
42    case 0n:
43      Done{(tensor, accumulator)}
44    case 1n+tail:
45      +index = U32.from_nat(cursor)
46      sum_decided(shape, tensor => value => natural_sum_loop(tail, shape, tensor, 1n+cursor, value), tensor, cursor, index, accumulator,
47        Shape.bool_case(Nat.is_eq(U32.to_nat(index), cursor)))
48
49def sum_loop(fuel: Nat, -shape: Shape.Shape, tensor: Core.Tensor<shape>, +index: U32,
50  accumulator: F32) -> Result<&2, &1, Core.Error, Core.Tensor<shape> & F32>:
51  natural_sum_loop(fuel, shape, tensor, U32.to_nat(index), accumulator)
52
53def scalar_result(-input_shape: Shape.Shape, tensor: Core.Tensor<input_shape>,
54  result: Result<&2, &1, Core.Error, Core.Tensor<scalar_shape()>>) ->
55  Result<&2, &1, Core.Error, Core.Tensor<input_shape> & Core.Tensor<scalar_shape()>>:
56  match result:
57    case Fail{error}:
58      Fail{error}
59    case Done{scalar}:
60      Done{(tensor, scalar)}
61
62def sum_result(-shape: Shape.Shape, result: Result<&2, &1, Core.Error, Core.Tensor<shape> & F32>) ->
63  Result<&2, &1, Core.Error, Core.Tensor<shape> & Core.Tensor<scalar_shape()>>:
64  match result:
65    case Fail{error}:
66      Fail{error}
67    case Done{pair}:
68      (tensor, value) = pair
69      scalar_result(shape, tensor, Core.new(scalar_shape(), value))
70
71# The fuel + 1 words of the array from slot `cursor`, whose read `pair` is
72# already made, added to `accumulator` left to right. Each step reads the
73# next word and hands it on as an argument: one tail call with no checked
74# read and no closure per word.
75def array_run(fuel: Nat, +cursor: Nat, accumulator: F32, pair: Array<F32> & F32) -> F32 & Array<F32>:
76  match fuel pair:
77    case 0n Tuple{array, word}:
78      (F32.add(accumulator, word), array)
79    case 1n+rest Tuple{array, word}:
80      array_run(rest, 1n+cursor, F32.add(accumulator, word), Array.get(F32, array, U32.from_nat(1n+cursor)))
81
82# The first `fuel` words of the array from slot `cursor`, added to
83# `accumulator` left to right: the F32 adds of `natural_sum_loop`, in the
84# same order.
85def array_sum(fuel: Nat, +cursor: Nat, accumulator: F32, array: Array<F32>) -> F32 & Array<F32>:
86  match fuel:
87    case 0n:
88      (accumulator, array)
89    case 1n+rest:
90      array_run(rest, cursor, accumulator, Array.get(F32, array, U32.from_nat(cursor)))
91
92def array_from(pair: Array<F32> & F32) -> Array<F32>:
93  Refinement.first(F32, pair)
94
95type Fact<-P: Type> is Data:
96  Fact{-proof: P}
97
98def kept_balanced(-depth: Nat, -original: Array<F32>, -balanced: Model.Balanced(F32, depth, original), -kept: Array<F32>,
99  -same: {kept == original : Array<F32>}) -> Fact<Model.Balanced(F32, depth, kept)>:
100  Fact{
101    %Equal.sym(Array<F32>, kept, original, same) : Model.Balanced(F32, depth, _)
102    balanced}
103
104# The array a read hands back is the array it read.
105def read_keeps(-depth: Nat, -original: Array<F32>, -balanced: Model.Balanced(F32, depth, original), -array: Array<F32>,
106  -index: U32, -pair: Array<F32> & F32, -exact: {pair == Array.get(F32, array, index) : Array<F32> & F32},
107  -same: {array == original : Array<F32>}) -> Fact<{array_from(pair) == original : Array<F32>}>:
108  Fact{Equal.trans(Array<F32>, array_from(pair), array_from(Array.get(F32, array, index)), original,
109    Equal.cong(Array<F32> & F32, Array<F32>, p => array_from(p), pair, Array.get(F32, array, index), exact),
110    Equal.trans(Array<F32>, array_from(Array.get(F32, array, index)), array, original,
111      Refinement.law_get_retains(F32, depth, array, index,
112        {\{Fact: p => p} : Fact<(Model.Balanced(F32, depth, array))> -> Model.Balanced(F32, depth, array)}(
113          kept_balanced(depth, original, balanced, array, same))),
114      same))}
115
116# A run hands back the array it read.
117def array_run_keeps(-depth: Nat, -original: Array<F32>, -balanced: Model.Balanced(F32, depth, original), fuel: Nat, +cursor: Nat,
118  -accumulator: F32, -array: Array<F32>, pair: Array<F32> & F32,
119  -exact: {pair == Array.get(F32, array, U32.from_nat(cursor)) : Array<F32> & F32}, -same: {array == original : Array<F32>}) ->
120  Fact<{Pair.snd(F32, Array<F32>, array_run(fuel, cursor, accumulator, pair)) == original : Array<F32>}>:
121  match fuel pair:
122    case 0n Tuple{kept, word}:
123      read_keeps(depth, original, balanced, array, U32.from_nat(cursor), (kept, word), exact, same)
124    case 1n+rest Tuple{kept, word}:
125      array_run_keeps(depth, original, balanced, rest, 1n+cursor, F32.add(accumulator, word), kept,
126        Array.get(F32, kept, U32.from_nat(1n+cursor)), {==},
127        {\{Fact: p => p} : Fact<({array_from((kept, word)) == original : Array<F32>})> -> {kept == original : Array<F32>}}(
128          read_keeps(depth, original, balanced, array, U32.from_nat(cursor), (kept, word), exact, same)))
129
130def array_sum_keeps(-depth: Nat, -original: Array<F32>, -balanced: Model.Balanced(F32, depth, original), fuel: Nat, +cursor: Nat,
131  -accumulator: F32, array: Array<F32>, -same: {array == original : Array<F32>}) ->
132  Fact<{Pair.snd(F32, Array<F32>, array_sum(fuel, cursor, accumulator, array)) == original : Array<F32>}>:
133  match fuel:
134    case 0n:
135      Fact{same}
136    case 1n+rest:
137      array_run_keeps(depth, original, balanced, rest, cursor, accumulator, array, Array.get(F32, array, U32.from_nat(cursor)), {==}, same)
138
139def owned_sum(-shape: Shape.Shape, +depth: Nat, +length: Nat, -array: Array<F32>,
140  -balanced: Model.Balanced(F32, depth, array), -valid: {Storage.bounds(depth, length) == True{} : Bool},
141  -exact: {length == Shape.logical_numel(Shape.dimensions(shape)) : Nat},
142  pair: F32 & Array<F32>, -kept: Fact<{Pair.snd(F32, Array<F32>, pair) == array : Array<F32>}>) ->
143  Result<&2, &1, Core.Error, Core.Tensor<shape> & F32>:
144  (value, retained) = pair
145  Done{(Core.TTensor{Storage.Owned{depth, length, retained,
146    {\{Fact: p => p} : Fact<(Model.Balanced(F32, depth, retained))> -> Model.Balanced(F32, depth, retained)}(
147      kept_balanced(depth, array, balanced, retained,
148        {\{Fact: p => p} : Fact<({Pair.snd(F32, Array<F32>, (value, retained)) == array : Array<F32>})> -> {retained == array : Array<F32>}}(kept))),
149    valid}, exact}, value)}
150
151# The first `count` words of the tensor added left to right from +0.0, next
152# to the tensor, which keeps every word.
153def tensor_sum(-shape: Shape.Shape, +count: Nat, tensor: Core.Tensor<shape>) -> Result<&2, &1, Core.Error, Core.Tensor<shape> & F32>:
154  Core.TTensor{storage, exact} = tensor
155  match storage:
156    case Storage.Empty{}:
157      Done{(Core.TTensor{Storage.Empty{}, exact}, 0.0)}
158    case Storage.Owned{+depth, +length, array, balanced, valid}:
159      owned_sum(shape, depth, length, array, balanced, valid, exact, array_sum(count, 0n, 0.0, array),
160        array_sum_keeps(depth, array, balanced, count, 0n, 0.0, array, {==}))
161
162def sum(+shape: Shape.Shape, tensor: Core.Tensor<shape>) ->
163  Result<&2, &1, Core.Error, Core.Tensor<shape> & Core.Tensor<scalar_shape()>>:
164  sum_result(shape, tensor_sum(shape, Shape.elements(shape), tensor))
165
166def pack_result(+shape: Shape.Shape,
167  result: Result<&2, &1, Core.Error, Core.Tensor<shape> & Core.Tensor<scalar_shape()>>) ->
168  Result<&2, &1, Core.Error, Core.Dynamic & Core.Dynamic>:
169  match result:
170    case Fail{error}:
171      Fail{error}
172    case Done{pair}:
173      (tensor, scalar) = pair
174      Done{(Core.pack(shape, tensor), Core.pack(scalar_shape(), scalar))}
175
176def sum_dynamic(dynamic: Core.Dynamic) -> Result<&2, &1, Core.Error, Core.Dynamic & Core.Dynamic>:
177  Core.Packed{+shape, tensor} = dynamic
178  pack_result(shape, sum(shape, tensor))
179
180def sum_observation(-shape: Shape.Shape,
181  result: Result<&2, &1, Core.Error, Core.Tensor<shape> & Core.Tensor<scalar_shape()>>) -> Maybe<&2, F32>:
182  match result:
183    case Fail{error}:
184      None{}
185    case Done{pair}:
186      (tensor, scalar) = pair
187      Core.observe(scalar_shape(), scalar, 0)