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 ../tensor/elementwise_totality.bend as Totality
7	import ./reduce.bend as Reduce
8	import ./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.
17	def 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	
25	def 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`.
33	def 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.
43	def 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.
53	law sum_dims_every_axis:
54	  for +axis: U32
55	  {SumDims.reduces([], axis) == True{} : Bool}
56	
57	law 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.
64	law 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	
69	law 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.
78	law 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	
83	law 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	
91	law 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.
101	law sum_dims_rank_zero_limit:
102	  {SumDims.limit([]) == 1 : U32}
103	
104	law 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.
111	law 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	
117	law 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	
124	law 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.
132	law 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.
140	law 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.
146	law 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.
155	law 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.
164	law 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>}