src/autograd/sum_dims_LAWS.bend

bendtorch-core@0.1.0.0 · 171 lines · 21 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/layout.bend as Layout
5import bendtorch-core@0.1.0.0/src/tensor/transforms.bend as Transforms
6import bendtorch-core@0.1.0.0/src/tensor/elementwise_totality.bend as Totality
7import bendtorch-core@0.1.0.0/src/autograd/reduce.bend as Reduce
8import bendtorch-core@0.1.0.0/src/autograd/sum_dims.bend as SumDims
9
10# PyTorch's sum(dim, keepdim) is reduce_to onto the kept shape, then a
11# reshape to the dropped dimensions. These laws state the axis bookkeeping
12# and the two steps; the reduce_to laws (reduce_LAWS.bend) give the words of
13# the first and the reshape laws (transform_LAWS.bend) say the second moves
14# none.
15
16# reduce_to's result as a sum's: its error as a ReduceError.
17def as_sum(pair: Core.Dynamic & Result<&2, &1, Reduce.Error, Core.Dynamic>) -> Core.Dynamic & Result<&2, &1, SumDims.Error, Core.Dynamic>:
18  (input, result) = pair
19  match result:
20    case Fail{error}:
21      (input, Fail{SumDims.ReduceError{error}})
22    case Done{output}:
23      (input, Done{output})
24
25def reshape_as_sum(result: Result<&2, &1, Transforms.Error, Core.Dynamic>) -> Result<&2, &1, SumDims.Error, Core.Dynamic>:
26  match result:
27    case Fail{error}:
28      Fail{SumDims.TransformError{error}}
29    case Done{output}:
30      Done{output}
31
32# A sum's result with its output reshaped to `dimensions`.
33def reshaped(+dimensions: List<&2, U32>, pair: Core.Dynamic & Result<&2, &1, SumDims.Error, Core.Dynamic>) ->
34  Core.Dynamic & Result<&2, &1, SumDims.Error, Core.Dynamic>:
35  (input, result) = pair
36  match result:
37    case Fail{error}:
38      (input, Fail{error})
39    case Done{output}:
40      (input, reshape_as_sum(Transforms.reshape_dynamic(output, dimensions)))
41
42# What a plan holds: the kept shape's dimensions and the dropped ones.
43def PlanKeeps(+axes: List<&2, U32>, +dimensions: List<&2, U32>, result: Result<&2, &2, SumDims.Error, SumDims.Plan>) -> Type:
44  match result:
45    case Fail{error}:
46      Unit
47    case Done{SumDims.Plan{kept, dropped}}:
48      {Shape.dimensions(kept) == SumDims.kept_dimensions(axes, dimensions) : List<&2, U32>} &
49        {dropped == SumDims.dropped_dimensions(axes, dimensions) : List<&2, U32>}
50
51# An empty axis list reduces every axis, as PyTorch's dim=[]; a listed axis
52# is reduced exactly when it is in the list.
53law sum_dims_every_axis:
54  for +axis: U32
55  {SumDims.reduces([], axis) == True{} : Bool}
56
57law sum_dims_listed_axis:
58  for +head: U32
59  for +tail: List<&2, U32>
60  for +axis: U32
61  {SumDims.reduces(head <> tail, axis) == Layout.contains(axis, head <> tail) : Bool}
62
63# The kept dimensions count axes from 0, and set each reduced axis to 1.
64law sum_dims_kept_start:
65  for +axes: List<&2, U32>
66  for +dimensions: List<&2, U32>
67  {SumDims.kept_dimensions(axes, dimensions) == SumDims.kept_from(axes, dimensions, 0) : List<&2, U32>}
68
69law sum_dims_kept_step:
70  for +axes: List<&2, U32>
71  for +extent: U32
72  for +rest: List<&2, U32>
73  for +axis: U32
74  {SumDims.kept_from(axes, extent <> rest, axis) ==
75    Bool.pick(U32, SumDims.reduces(axes, axis), 1, extent) <> SumDims.kept_from(axes, rest, U32.add(axis, 1)) : List<&2, U32>}
76
77# The dropped dimensions count axes from 0, and leave out each reduced axis.
78law sum_dims_dropped_start:
79  for +axes: List<&2, U32>
80  for +dimensions: List<&2, U32>
81  {SumDims.dropped_dimensions(axes, dimensions) == SumDims.dropped_from(axes, dimensions, 0) : List<&2, U32>}
82
83law sum_dims_dropped_reduced:
84  for +axes: List<&2, U32>
85  for +extent: U32
86  for +rest: List<&2, U32>
87  for +axis: U32
88  for reduced: {SumDims.reduces(axes, axis) == True{} : Bool}
89  {SumDims.dropped_from(axes, extent <> rest, axis) == SumDims.dropped_from(axes, rest, U32.add(axis, 1)) : List<&2, U32>}
90
91law sum_dims_dropped_kept:
92  for +axes: List<&2, U32>
93  for +extent: U32
94  for +rest: List<&2, U32>
95  for +axis: U32
96  for kept: {SumDims.reduces(axes, axis) == False{} : Bool}
97  {SumDims.dropped_from(axes, extent <> rest, axis) == extent <> SumDims.dropped_from(axes, rest, U32.add(axis, 1)) : List<&2, U32>}
98
99# An input has its rank's axes, and a rank-0 input has axis 0, as in
100# PyTorch.
101law sum_dims_rank_zero_limit:
102  {SumDims.limit([]) == 1 : U32}
103
104law sum_dims_rank_limit:
105  for +extent: U32
106  for +rest: List<&2, U32>
107  {SumDims.limit(extent <> rest) == U32.from_nat(Shape.rank_dimensions(extent <> rest)) : U32}
108
109# The axes are refused as permute refuses them, and in PyTorch's order: an
110# axis at or past the limit first, then an axis listed twice.
111law sum_dims_axis_out_of_bounds:
112  for +axes: List<&2, U32>
113  for +limit: U32
114  for outside: {Layout.bounded(axes, limit) == False{} : Bool}
115  {SumDims.refusal(axes, limit) == Some{Layout.AxisOutOfBounds{}} : Maybe<&2, Layout.Error>}
116
117law sum_dims_duplicate_axis:
118  for +axes: List<&2, U32>
119  for +limit: U32
120  for inside: {Layout.bounded(axes, limit) == True{} : Bool}
121  for repeated: {Layout.unique(axes) == False{} : Bool}
122  {SumDims.refusal(axes, limit) == Some{Layout.DuplicateAxis{}} : Maybe<&2, Layout.Error>}
123
124law sum_dims_axes_accepted:
125  for +axes: List<&2, U32>
126  for +limit: U32
127  for inside: {Layout.bounded(axes, limit) == True{} : Bool}
128  for unique: {Layout.unique(axes) == True{} : Bool}
129  {SumDims.refusal(axes, limit) == None{} : Maybe<&2, Layout.Error>}
130
131# A refused list of axes fails the plan with the refusal.
132law sum_dims_plan_refused:
133  for +shape: Shape.Shape
134  for +axes: List<&2, U32>
135  for error: Layout.Error
136  for refused: {SumDims.refusal(axes, SumDims.limit(Shape.dimensions(shape))) == Some{error} : Maybe<&2, Layout.Error>}
137  {SumDims.plan(shape, axes) == Fail{SumDims.TransformError{Transforms.LayoutError{error}}} : Result<&2, &2, SumDims.Error, SumDims.Plan>}
138
139# A plan holds the kept dimensions as its shape and the dropped dimensions.
140law sum_dims_plan_dimensions:
141  for +shape: Shape.Shape
142  for +axes: List<&2, U32>
143  PlanKeeps(axes, Shape.dimensions(shape), SumDims.plan(shape, axes))
144
145# A refused plan fails the sum with its error and hands the input back.
146law sum_dims_refused:
147  for input: Core.Dynamic
148  for +axes: List<&2, U32>
149  for keepdim: Bool
150  for error: SumDims.Error
151  for refused: {SumDims.plan(Totality.dynamic_shape(input), axes) == Fail{error} : Result<&2, &2, SumDims.Error, SumDims.Plan>}
152  {SumDims.sum_dims(input, axes, keepdim) == (input, Fail{error}) : Core.Dynamic & Result<&2, &1, SumDims.Error, Core.Dynamic>}
153
154# With keepdim, the sum is exactly reduce_to onto the kept shape.
155law sum_dims_keepdim:
156  for input: Core.Dynamic
157  for +axes: List<&2, U32>
158  for +kept: Shape.Shape
159  for +dropped: List<&2, U32>
160  for planned: {SumDims.plan(Totality.dynamic_shape(input), axes) == Done{SumDims.Plan{kept, dropped}} : Result<&2, &2, SumDims.Error, SumDims.Plan>}
161  {SumDims.sum_dims(input, axes, True{}) == as_sum(Reduce.reduce_to(input, kept)) : Core.Dynamic & Result<&2, &1, SumDims.Error, Core.Dynamic>}
162
163# Without keepdim, it is the keepdim sum reshaped to the dropped dimensions.
164law sum_dims_drops:
165  for input: Core.Dynamic
166  for +axes: List<&2, U32>
167  for +kept: Shape.Shape
168  for +dropped: List<&2, U32>
169  for planned: {SumDims.plan(Totality.dynamic_shape(input), axes) == Done{SumDims.Plan{kept, dropped}} : Result<&2, &2, SumDims.Error, SumDims.Plan>}
170  {SumDims.sum_dims(input, axes, False{}) == reshaped(dropped, SumDims.sum_dims(input, axes, True{})) :
171    Core.Dynamic & Result<&2, &1, SumDims.Error, Core.Dynamic>}