src/autograd/sum_dims_PROOF.bend
bendtorch-core@0.1.0.0 · 139 lines · 5 definitions
1import Base
2import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
3import bendtorch-core@0.1.0.0/src/storage/arith.bend as Arith
4import bendtorch-core@0.1.0.0/src/tensor/core.bend as Core
5import bendtorch-core@0.1.0.0/src/tensor/layout.bend as Layout
6import bendtorch-core@0.1.0.0/src/tensor/transforms.bend as Transforms
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
9import bendtorch-core@0.1.0.0/src/autograd/sum_dims_LAWS.bend as Laws
10
11def Laws.sum_dims_every_axis(axis):
12 {==}
13
14def Laws.sum_dims_listed_axis(head, tail, axis):
15 {==}
16
17def Laws.sum_dims_kept_start(axes, dimensions):
18 {==}
19
20def Laws.sum_dims_kept_step(axes, extent, rest, axis):
21 {==}
22
23def Laws.sum_dims_dropped_start(axes, dimensions):
24 {==}
25
26def Laws.sum_dims_dropped_reduced(axes, extent, rest, axis, reduced):
27 %Equal.sym(Bool, SumDims.reduces(axes, axis), True{}, reduced) :
28 {SumDims.dropped_next(_, extent, SumDims.dropped_from(axes, rest, U32.add(axis, 1))) ==
29 SumDims.dropped_from(axes, rest, U32.add(axis, 1)) : List<&2, U32>}
30 {==}
31
32def Laws.sum_dims_dropped_kept(axes, extent, rest, axis, kept):
33 %Equal.sym(Bool, SumDims.reduces(axes, axis), False{}, kept) :
34 {SumDims.dropped_next(_, extent, SumDims.dropped_from(axes, rest, U32.add(axis, 1))) ==
35 extent <> SumDims.dropped_from(axes, rest, U32.add(axis, 1)) : List<&2, U32>}
36 {==}
37
38def Laws.sum_dims_rank_zero_limit():
39 {==}
40
41def Laws.sum_dims_rank_limit(extent, rest):
42 {==}
43
44def Laws.sum_dims_axis_out_of_bounds(axes, limit, outside):
45 %Equal.sym(Bool, Layout.bounded(axes, limit), False{}, outside) :
46 {SumDims.refused(_, Layout.unique(axes)) == Some{Layout.AxisOutOfBounds{}} : Maybe<&2, Layout.Error>}
47 {==}
48
49def Laws.sum_dims_duplicate_axis(axes, limit, inside, repeated):
50 %Equal.sym(Bool, Layout.bounded(axes, limit), True{}, inside) :
51 {SumDims.refused(_, Layout.unique(axes)) == Some{Layout.DuplicateAxis{}} : Maybe<&2, Layout.Error>}
52 %Equal.sym(Bool, Layout.unique(axes), False{}, repeated) :
53 {SumDims.refused(True{}, _) == Some{Layout.DuplicateAxis{}} : Maybe<&2, Layout.Error>}
54 {==}
55
56def Laws.sum_dims_axes_accepted(axes, limit, inside, unique):
57 %Equal.sym(Bool, Layout.bounded(axes, limit), True{}, inside) :
58 {SumDims.refused(_, Layout.unique(axes)) == None{} : Maybe<&2, Layout.Error>}
59 %Equal.sym(Bool, Layout.unique(axes), True{}, unique) :
60 {SumDims.refused(True{}, _) == None{} : Maybe<&2, Layout.Error>}
61 {==}
62
63def Laws.sum_dims_plan_refused(shape, axes, error, refused):
64 %Equal.sym(Maybe<&2, Layout.Error>, SumDims.refusal(axes, SumDims.limit(Shape.dimensions(shape))), Some{error}, refused) :
65 {SumDims.plan_checked(axes, Shape.dimensions(shape), _) == Fail{SumDims.TransformError{Transforms.LayoutError{error}}} :
66 Result<&2, &2, SumDims.Error, SumDims.Plan>}
67 {==}
68
69def built_keeps(+axes: List<&2, U32>, +dimensions: List<&2, U32>,
70 built: Shape.ConstructResult<SumDims.kept_dimensions(axes, dimensions)>) ->
71 Laws.PlanKeeps(axes, dimensions,
72 SumDims.plan_built(SumDims.kept_dimensions(axes, dimensions), SumDims.dropped_dimensions(axes, dimensions), built)):
73 match built:
74 case Shape.SError{error}:
75 Unit{}
76 case Shape.SBuilt{+kept, preserved}:
77 (preserved, {==})
78
79def checked_keeps(+axes: List<&2, U32>, +dimensions: List<&2, U32>, checked: Maybe<&2, Layout.Error>) ->
80 Laws.PlanKeeps(axes, dimensions, SumDims.plan_checked(axes, dimensions, checked)):
81 match checked:
82 case Some{error}:
83 Unit{}
84 case None{}:
85 built_keeps(axes, dimensions, Shape.construct(SumDims.kept_dimensions(axes, dimensions)))
86
87def Laws.sum_dims_plan_dimensions(shape, axes):
88 checked_keeps(axes, Shape.dimensions(shape), SumDims.refusal(axes, SumDims.limit(Shape.dimensions(shape))))
89
90def Laws.sum_dims_refused(input, axes, keepdim, error, refused):
91 Core.Packed{+shape, tensor} = input
92 %Equal.sym(Result<&2, &2, SumDims.Error, SumDims.Plan>, SumDims.plan(shape, axes), Fail{error}, refused) :
93 {SumDims.sum_checked(Core.Packed{shape, tensor}, keepdim, _) == (Core.Packed{shape, tensor}, Fail{error}) :
94 Core.Dynamic & Result<&2, &1, SumDims.Error, Core.Dynamic>}
95 {==}
96
97
98def kept_as_sum(+dropped: List<&2, U32>, pair: Core.Dynamic & Result<&2, &1, Reduce.Error, Core.Dynamic>) ->
99 {SumDims.summed(dropped, True{}, pair) == Laws.as_sum(pair) : Core.Dynamic & Result<&2, &1, SumDims.Error, Core.Dynamic>}:
100 match pair:
101 case Tuple{input, Fail{error}}:
102 {==}
103 case Tuple{input, Done{output}}:
104 {==}
105
106def Laws.sum_dims_keepdim(input, axes, kept, dropped, planned):
107 Core.Packed{+shape, tensor} = input
108 %Equal.sym(Result<&2, &2, SumDims.Error, SumDims.Plan>, SumDims.plan(shape, axes), Done{SumDims.Plan{kept, dropped}}, planned) :
109 {SumDims.sum_checked(Core.Packed{shape, tensor}, True{}, _) == Laws.as_sum(Reduce.reduce_to(Core.Packed{shape, tensor}, kept)) :
110 Core.Dynamic & Result<&2, &1, SumDims.Error, Core.Dynamic>}
111 kept_as_sum(dropped, Reduce.reduce_to(Core.Packed{shape, tensor}, kept))
112
113def reshape_renamed(result: Result<&2, &1, Transforms.Error, Core.Dynamic>) ->
114 {SumDims.transformed(result) == Laws.reshape_as_sum(result) : Result<&2, &1, SumDims.Error, Core.Dynamic>}:
115 match result:
116 case Fail{error}:
117 {==}
118 case Done{output}:
119 {==}
120
121
122def dropped_reshaped(+dropped: List<&2, U32>, pair: Core.Dynamic & Result<&2, &1, Reduce.Error, Core.Dynamic>) ->
123 {SumDims.summed(dropped, False{}, pair) == Laws.reshaped(dropped, SumDims.summed(dropped, True{}, pair)) :
124 Core.Dynamic & Result<&2, &1, SumDims.Error, Core.Dynamic>}:
125 match pair:
126 case Tuple{input, Fail{error}}:
127 {==}
128 case Tuple{input, Done{output}}:
129 Equal.cong(Result<&2, &1, SumDims.Error, Core.Dynamic>, Core.Dynamic & Result<&2, &1, SumDims.Error, Core.Dynamic>,
130 result => (input, result), SumDims.transformed(Transforms.reshape_dynamic(output, dropped)),
131 Laws.reshape_as_sum(Transforms.reshape_dynamic(output, dropped)), reshape_renamed(Transforms.reshape_dynamic(output, dropped)))
132
133def Laws.sum_dims_drops(input, axes, kept, dropped, planned):
134 Core.Packed{+shape, tensor} = input
135 %Equal.sym(Result<&2, &2, SumDims.Error, SumDims.Plan>, SumDims.plan(shape, axes), Done{SumDims.Plan{kept, dropped}}, planned) :
136 {SumDims.sum_checked(Core.Packed{shape, tensor}, False{}, _) ==
137 Laws.reshaped(dropped, SumDims.sum_checked(Core.Packed{shape, tensor}, True{}, _)) :
138 Core.Dynamic & Result<&2, &1, SumDims.Error, Core.Dynamic>}
139 dropped_reshaped(dropped, Reduce.reduce_to(Core.Packed{shape, tensor}, kept))