1 import Base 2 import ../shape/core.bend as Shape 3 import ../storage/arith.bend as Arith 4 import ../tensor/core.bend as Core 5 import ../tensor/layout.bend as Layout 6 import ../tensor/transforms.bend as Transforms 7 import ./reduce.bend as Reduce 8 import ./sum_dims.bend as SumDims 9 import ./sum_dims_LAWS.bend as Laws 10 11 def Laws.sum_dims_every_axis(axis): 12 {==} 13 14 def Laws.sum_dims_listed_axis(head, tail, axis): 15 {==} 16 17 def Laws.sum_dims_kept_start(axes, dimensions): 18 {==} 19 20 def Laws.sum_dims_kept_step(axes, extent, rest, axis): 21 {==} 22 23 def Laws.sum_dims_dropped_start(axes, dimensions): 24 {==} 25 26 def 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 32 def 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 38 def Laws.sum_dims_rank_zero_limit(): 39 {==} 40 41 def Laws.sum_dims_rank_limit(extent, rest): 42 {==} 43 44 def 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 49 def 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 56 def 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 63 def 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 69 def built_keeps(+axes: List<&2, U32>, +dimensions: List<&2, U32>, 70 built: Shape.ConstructResult) -> 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 79 def 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 87 def Laws.sum_dims_plan_dimensions(shape, axes): 88 checked_keeps(axes, Shape.dimensions(shape), SumDims.refusal(axes, SumDims.limit(Shape.dimensions(shape)))) 89 90 def 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 # With keepdim, summed only renames reduce_to's error. 98 def 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 106 def 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 113 def 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 # Without keepdim, summed reshapes what it gives with keepdim. 122 def 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 133 def 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))