src/autograd/reduce_LAWS.bend

bendtorch-core@0.1.0.0 · 172 lines · 24 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/elementwise_totality.bend as Totality
5import bendtorch-core@0.1.0.0/src/tensor/projection_LAWS.bend as Projection
6import bendtorch-core@0.1.0.0/src/autograd/reduce.bend as Reduce
7
8def owner(pair: Core.Dynamic & Result<&2, &1, Reduce.Error, Core.Dynamic>) -> Core.Dynamic:
9  (input, result) = pair
10  input
11
12def succeeded(pair: Core.Dynamic & Result<&2, &1, Reduce.Error, Core.Dynamic>) -> Bool:
13  (input, result) = pair
14  Result.is_done(&2, &1, Reduce.Error, Core.Dynamic, result)
15
16def word_at(values: List<&2, F32>, +position: Nat) -> Maybe<&2, F32>:
17  match values position:
18    case Nil{} _:
19      None{}
20    case value <> rest 0n:
21      Some{value}
22    case value <> rest 1n+next:
23      word_at(rest, next)
24
25def add_word(total: Maybe<&2, F32>, value: Maybe<&2, F32>) -> Maybe<&2, F32>:
26  match total value:
27    case None{} _:
28      None{}
29    case Some{sum} None{}:
30      None{}
31    case Some{sum} Some{word}:
32      Some{F32.add(sum, word)}
33
34def matched(total: Maybe<&2, F32>, value: Maybe<&2, F32>, same: Bool) -> Maybe<&2, F32>:
35  match same:
36    case False{}:
37      total
38    case True{}:
39      add_word(total, value)
40
41def contribute(+values: List<&2, F32>, +target_index: U32, total: Maybe<&2, F32>, offsets: Maybe<&1, U32 & U32>) -> Maybe<&2, F32>:
42  match offsets:
43    case None{}:
44      None{}
45    case Some{pair}:
46      (target_offset, source_offset) = pair
47      matched(total, word_at(values, U32.to_nat(source_offset)), Nat.is_eq(U32.to_nat(target_offset), U32.to_nat(target_index)))
48
49# The specified accumulation: start at +0.0 and visit contribution positions
50# 0, 1, ..., count - 1 in ascending order; each position whose independent
51# broadcast projection, Projection.offset_spec, has target_offset equal to
52# target_index adds the word at its source_offset with one F32 add.
53def fiber(+plan: Shape.BroadcastPlan, +values: List<&2, F32>, +target_index: U32, +count: Nat) -> Maybe<&2, F32>:
54  match count:
55    case 0n:
56      Some{0.0}
57    case 1n+position:
58      contribute(values, target_index, fiber(plan, values, target_index, position), Projection.offset_spec(plan, position))
59
60def planned_fiber(+values: List<&2, F32>, +target_index: U32, +count: Nat,
61  plan: Result<&2, &2, Shape.Error, Shape.BroadcastPlan>) -> Maybe<&2, F32>:
62  match plan:
63    case Fail{error}:
64      None{}
65    case Done{+plan}:
66      fiber(plan, values, target_index, count)
67
68# values lists the contribution's words: the checked observation at every
69# U32 position is the word at that position.
70def Listed(-input: Core.Dynamic, +values: List<&2, F32>) -> Type:
71  @+source: U32 -> {Core.checked_observe_dynamic(input, source) == word_at(values, U32.to_nat(source)) : Maybe<&2, F32>}
72
73def MappingWith(-input: Core.Dynamic, +target: Shape.Shape, plan: Result<&2, &2, Shape.Error, Shape.BroadcastPlan>,
74  pair: Core.Dynamic & Result<&2, &1, Reduce.Error, Core.Dynamic>) -> Type:
75  (retained, result) = pair
76  match result:
77    case Fail{error}:
78      Unit
79    case Done{output}:
80      {Totality.dimensions(output) == Shape.dimensions(target) : List<&2, U32>} &
81        (@+values: List<&2, F32> -> @-listed: Listed(input, values) ->
82          @+index: U32 -> @inside: {Nat.is_lt(U32.to_nat(index), Shape.logical_numel(Shape.dimensions(target))) == True{} : Bool} ->
83            {Core.checked_observe_dynamic(output, index) ==
84              planned_fiber(values, index, Shape.logical_numel(Totality.dimensions(input)), plan) : Maybe<&2, F32>})
85
86def Mapping(input: Core.Dynamic, +target: Shape.Shape, pair: Core.Dynamic & Result<&2, &1, Reduce.Error, Core.Dynamic>) -> Type:
87  MappingWith(input, target, Shape.broadcast_checked(target, Totality.dynamic_shape(input)), pair)
88
89def ZeroValues(+target: Shape.Shape, pair: Core.Dynamic & Result<&2, &1, Reduce.Error, Core.Dynamic>) -> Type:
90  (retained, result) = pair
91  match result:
92    case Fail{error}:
93      Unit
94    case Done{output}:
95      @+index: U32 -> @inside: {Nat.is_lt(U32.to_nat(index), Shape.logical_numel(Shape.dimensions(target))) == True{} : Bool} ->
96        {Core.checked_observe_dynamic(output, index) == Some{0.0} : Maybe<&2, F32>}
97
98type Retained<-input: Core.Dynamic, -pair: Core.Dynamic & Result<&2, &1, Reduce.Error, Core.Dynamic>> is Data:
99  Retained{-kept: Totality.SameSource(input, owner(pair))}
100
101type MappingModel<-input: Core.Dynamic, -target: Shape.Shape,
102  -pair: Core.Dynamic & Result<&2, &1, Reduce.Error, Core.Dynamic>> is Data:
103  MappingModel{-mapping: Mapping(input, target, pair)}
104
105type Success<-pair: Core.Dynamic & Result<&2, &1, Reduce.Error, Core.Dynamic>> is Data:
106  Success{-successful: {succeeded(pair) == True{} : Bool}}
107
108type Length<-input: Core.Dynamic, -values: List<&2, F32>> is Data:
109  Length{-exact: {List.length(&2, F32, values) == Shape.logical_numel(Totality.dimensions(input)) : Nat}}
110
111type Zeros<-target: Shape.Shape, -pair: Core.Dynamic & Result<&2, &1, Reduce.Error, Core.Dynamic>> is Data:
112  Zeros{-zeros: ZeroValues(target, pair)}
113
114# The contribution owner comes back with the same dimensions and observations.
115law reduce_retains_input:
116  for input: Core.Dynamic
117  for +target: Shape.Shape
118  Retained<input, Reduce.reduce_to(input, target)>
119
120# On success the output has the target's dimensions, and for every list of
121# words that lists the contribution, each in-bounds target position holds the
122# ordered fiber sum of those words.
123law reduce_mapping:
124  for input: Core.Dynamic
125  for +target: Shape.Shape
126  MappingModel<input, target, Reduce.reduce_to(input, target)>
127
128# A list of words that lists the contribution has exactly the contribution's
129# logical length.
130law listed_length:
131  for input: Core.Dynamic
132  for +values: List<&2, F32>
133  for listed: Listed(input, values)
134  Length<input, values>
135
136# A target whose right-aligned broadcast with the contribution shape is the
137# contribution shape itself always succeeds.
138law reduce_valid_success:
139  for input: Core.Dynamic
140  for +target: Shape.Shape
141  for +alignment: Shape.Broadcast<Shape.extents(target), Shape.extents(Totality.dynamic_shape(input)), Shape.extents(Totality.dynamic_shape(input))>
142  Success<Reduce.reduce_to(input, target)>
143
144# Reducing an empty contribution gives zeros at every target position, not an
145# empty tensor: for example [0, 3] reduced to [1, 3] is three zeros.
146law reduce_empty_zeros:
147  for input: Core.Dynamic
148  for +target: Shape.Shape
149  for empty: {Shape.logical_numel(Totality.dimensions(input)) == 0n : Nat}
150  Zeros<target, Reduce.reduce_to(input, target)>
151
152# When the broadcast planner rejects the target and contribution shapes,
153# reduce_to fails with exactly that planner error and returns the contribution
154# owner. For example [2] to [0] fails with IncompatibleBroadcast instead of
155# returning an empty tensor.
156law reduce_shape_error:
157  for input: Core.Dynamic
158  for +target: Shape.Shape
159  for error: Shape.Error
160  for rejected: {Shape.broadcast_checked(target, Totality.dynamic_shape(input)) == Fail{error} : Result<&2, &2, Shape.Error, Shape.BroadcastPlan>}
161  {Reduce.reduce_to(input, target) == (input, Fail{Reduce.TensorError{Core.ShapeError{error}}}) : Core.Dynamic & Result<&2, &1, Reduce.Error, Core.Dynamic>}
162
163# When the planned broadcast of the target is not the contribution's shape
164# (for example a [2, 3] target for a [3] contribution), reduce_to fails with
165# NotReducible and returns the contribution owner.
166law reduce_not_reducible:
167  for input: Core.Dynamic
168  for +target: Shape.Shape
169  for +plan: Shape.BroadcastPlan
170  for planned: {Shape.broadcast_checked(target, Totality.dynamic_shape(input)) == Done{plan} : Result<&2, &2, Shape.Error, Shape.BroadcastPlan>}
171  for differ: {Shape.dimensions(Shape.output_shape(plan)) != Totality.dimensions(input) : List<&2, U32>}
172  {Reduce.reduce_to(input, target) == (input, Fail{Reduce.NotReducible{}}) : Core.Dynamic & Result<&2, &1, Reduce.Error, Core.Dynamic>}