1	import Base
2	import ../shape/core.bend as Shape
3	import ../tensor/core.bend as Core
4	import ../tensor/layout.bend as Layout
5	import ../tensor/transforms.bend as Transforms
6	import ./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	
21	type 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.
26	type 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.
30	def 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.
38	def 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	
45	def kept_dimensions(+axes: List<&2, U32>, dimensions: List<&2, U32>) -> List<&2, U32>:
46	  kept_from(axes, dimensions, 0)
47	
48	def 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.
56	def 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	
63	def 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.
68	def 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.
77	def 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	
86	def refusal(+axes: List<&2, U32>, +limit: U32) -> Maybe<&2, Layout.Error>:
87	  refused(Layout.bounded(axes, limit), Layout.unique(axes))
88	
89	def 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	
96	def 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.
105	def 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	
108	def 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	
115	def 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	
122	def 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	
130	def 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.
137	def 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	
141	def 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.
151	def 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	
155	def 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)