src/autograd/sum_dims.bend

bendtorch-core@0.1.0.0 · 160 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/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/autograd/reduce.bend as Reduce
7
8# PyTorch's sum(input, dim, keepdim): the sum over the listed axes. It is
9# reduce_to onto the kept shape, the input's dimensions with every reduced
10# axis set to 1, so each output cell is the ordered fold of its fiber from
11# +0.0 on the parallel build. With keepdim False the result is then reshaped
12# to the dropped dimensions, the input's without the reduced axes, which
13# moves no word. An empty axis list reduces every axis, as PyTorch's dim=[]
14# does; a whole-tensor sum with no dim stays Reduction.sum.
15#
16# The axes are checked as PyTorch checks them: any axis at or past the rank
17# (past 0 for a rank-0 input, whose only axis is 0) is AxisOutOfBounds, then
18# an axis listed twice is DuplicateAxis, each with the code permute gives it.
19# Negative axes are not taken: an axis is a U32.
20
21type Error is Data:
22  ReduceError{error: Reduce.Error}
23  TransformError{error: Transforms.Error}
24
25# The kept shape, built with Shape.construct, and the dropped dimensions.
26type Plan is Data:
27  Plan{kept: Shape.Shape, dropped: List<&2, U32>}
28
29# Whether the sum reduces axis `axis`: every axis for an empty list.
30def reduces(+axes: List<&2, U32>, +axis: U32) -> Bool:
31  match axes:
32    case Nil{}:
33      True{}
34    case +head <> +tail:
35      Layout.contains(axis, head <> tail)
36
37# The dimensions from axis `axis` on, each reduced one set to 1.
38def kept_from(+axes: List<&2, U32>, dimensions: List<&2, U32>, +axis: U32) -> List<&2, U32>:
39  match dimensions:
40    case Nil{}:
41      Nil{}
42    case extent <> rest:
43      Bool.pick(U32, reduces(axes, axis), 1, extent) <> kept_from(axes, rest, U32.add(axis, 1))
44
45def kept_dimensions(+axes: List<&2, U32>, dimensions: List<&2, U32>) -> List<&2, U32>:
46  kept_from(axes, dimensions, 0)
47
48def dropped_next(reduced: Bool, extent: U32, rest: List<&2, U32>) -> List<&2, U32>:
49  match reduced:
50    case True{}:
51      rest
52    case False{}:
53      extent <> rest
54
55# The dimensions from axis `axis` on, without the reduced ones.
56def dropped_from(+axes: List<&2, U32>, dimensions: List<&2, U32>, +axis: U32) -> List<&2, U32>:
57  match dimensions:
58    case Nil{}:
59      Nil{}
60    case extent <> rest:
61      dropped_next(reduces(axes, axis), extent, dropped_from(axes, rest, U32.add(axis, 1)))
62
63def dropped_dimensions(+axes: List<&2, U32>, dimensions: List<&2, U32>) -> List<&2, U32>:
64  dropped_from(axes, dimensions, 0)
65
66# How many axes an input of these dimensions has for PyTorch's check: its
67# rank, and 1 for a rank-0 input.
68def limit(dimensions: List<&2, U32>) -> U32:
69  match dimensions:
70    case Nil{}:
71      1
72    case +extent <> +rest:
73      U32.from_nat(Shape.rank_dimensions(extent <> rest))
74
75# PyTorch's refusal of the axes, permute's checks in permute's order: an axis
76# at or past the limit, then an axis listed twice.
77def refused(bounded: Bool, unique: Bool) -> Maybe<&2, Layout.Error>:
78  match bounded unique:
79    case False{} unique:
80      Some{Layout.AxisOutOfBounds{}}
81    case True{} False{}:
82      Some{Layout.DuplicateAxis{}}
83    case True{} True{}:
84      None{}
85
86def refusal(+axes: List<&2, U32>, +limit: U32) -> Maybe<&2, Layout.Error>:
87  refused(Layout.bounded(axes, limit), Layout.unique(axes))
88
89def plan_built(-kept: List<&2, U32>, +dropped: List<&2, U32>, built: Shape.ConstructResult<kept>) -> Result<&2, &2, Error, Plan>:
90  match built:
91    case Shape.SError{error}:
92      Fail{ReduceError{Reduce.TensorError{Core.ShapeError{error}}}}
93    case Shape.SBuilt{+shape, preserved}:
94      Done{Plan{shape, dropped}}
95
96def plan_checked(+axes: List<&2, U32>, +dimensions: List<&2, U32>, checked: Maybe<&2, Layout.Error>) -> Result<&2, &2, Error, Plan>:
97  match checked:
98    case Some{error}:
99      Fail{TransformError{Transforms.LayoutError{error}}}
100    case None{}:
101      plan_built(kept_dimensions(axes, dimensions), dropped_dimensions(axes, dimensions),
102        Shape.construct(kept_dimensions(axes, dimensions)))
103
104# The sum's plan for an input of this shape, or the refusal of its axes.
105def plan(+shape: Shape.Shape, +axes: List<&2, U32>) -> Result<&2, &2, Error, Plan>:
106  plan_checked(axes, Shape.dimensions(shape), refusal(axes, limit(Shape.dimensions(shape))))
107
108def transformed(result: Result<&2, &1, Transforms.Error, Core.Dynamic>) -> Result<&2, &1, Error, Core.Dynamic>:
109  match result:
110    case Fail{error}:
111      Fail{TransformError{error}}
112    case Done{output}:
113      Done{output}
114
115def shaped(+dropped: List<&2, U32>, keepdim: Bool, output: Core.Dynamic) -> Result<&2, &1, Error, Core.Dynamic>:
116  match keepdim:
117    case True{}:
118      Done{output}
119    case False{}:
120      transformed(Transforms.reshape_dynamic(output, dropped))
121
122def squeezed(+dropped: List<&2, U32>, keepdim: Bool, result: Result<&2, &1, Reduce.Error, Core.Dynamic>) ->
123  Result<&2, &1, Error, Core.Dynamic>:
124  match result:
125    case Fail{error}:
126      Fail{ReduceError{error}}
127    case Done{output}:
128      shaped(dropped, keepdim, output)
129
130def summed(+dropped: List<&2, U32>, keepdim: Bool, pair: Core.Dynamic & Result<&2, &1, Reduce.Error, Core.Dynamic>) ->
131  Core.Dynamic & Result<&2, &1, Error, Core.Dynamic>:
132  (input, result) = pair
133  (input, squeezed(dropped, keepdim, result))
134
135# The sum of a planned input: reduce_to onto the kept shape, then, without
136# keepdim, the reshape to the dropped dimensions. The input owner comes back.
137def sum_planned(input: Core.Dynamic, plan: Plan, keepdim: Bool) -> Core.Dynamic & Result<&2, &1, Error, Core.Dynamic>:
138  Plan{+kept, +dropped} = plan
139  summed(dropped, keepdim, Reduce.reduce_to(input, kept))
140
141def sum_checked(input: Core.Dynamic, keepdim: Bool, planned: Result<&2, &2, Error, Plan>) ->
142  Core.Dynamic & Result<&2, &1, Error, Core.Dynamic>:
143  match planned:
144    case Fail{error}:
145      (input, Fail{error})
146    case Done{plan}:
147      sum_planned(input, plan, keepdim)
148
149# PyTorch's input.sum(dim=axes, keepdim=keepdim). The input owner always comes
150# back unchanged next to the result.
151def sum_dims(input: Core.Dynamic, +axes: List<&2, U32>, keepdim: Bool) -> Core.Dynamic & Result<&2, &1, Error, Core.Dynamic>:
152  Core.Packed{+shape, tensor} = input
153  sum_checked(Core.Packed{shape, tensor}, keepdim, plan(shape, axes))
154
155def error_code(error: Error) -> U32:
156  match error:
157    case ReduceError{error}:
158      Reduce.error_code(error)
159    case TransformError{error}:
160      Transforms.error_code(error)