src/tensor/projection_LAWS.bend
bendtorch-core@0.1.0.0 · 73 lines · 9 definitions
1import Base
2import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
3import bendtorch-core@0.1.0.0/src/shape/LAWS.bend as ShapeLaws
4import bendtorch-core@0.1.0.0/src/tensor/core.bend as Core
5import bendtorch-core@0.1.0.0/src/tensor/elementwise.bend as Elementwise
6
7def pair_offsets(left: Maybe<&2, U32>, right: Maybe<&2, U32>) -> Maybe<&1, U32 & U32>:
8 match left right:
9 case None{} _:
10 None{}
11 case Some{left} None{}:
12 None{}
13 case Some{left} Some{right}:
14 Some{(left, right)}
15
16def offset_spec(+plan: Shape.BroadcastPlan, +position: Nat) -> Maybe<&1, U32 & U32>:
17 +left = Shape.extents(Shape.left_shape(plan))
18 +right = Shape.extents(Shape.right_shape(plan))
19 +output = Shape.extents(Shape.output_shape(plan))
20 +index = ShapeLaws.fin_of(Shape.numel(output), position)
21 pair_offsets(
22 ShapeLaws.flatten_spec(left, ShapeLaws.project_left_coordinate(left, right, output, Shape.alignment(plan), index)),
23 ShapeLaws.flatten_spec(right, ShapeLaws.project_right_coordinate(left, right, output, Shape.alignment(plan), index)))
24
25def indices_result(result: Result<&2, &1, Core.Error, U32 & U32>) -> Maybe<&1, U32 & U32>:
26 match result:
27 case Fail{error}:
28 None{}
29 case Done{indices}:
30 Some{indices}
31
32def indices_accepted(result: Result<&2, &1, Core.Error, U32 & U32>) -> Bool:
33 match result:
34 case Fail{error}:
35 False{}
36 case Done{indices}:
37 True{}
38
39def specified_expected(operation: Elementwise.Operation, left: Core.Dynamic, right: Core.Dynamic,
40 offsets: Maybe<&1, U32 & U32>) -> Maybe<&2, F32>:
41 match offsets:
42 case None{}:
43 None{}
44 case Some{indices}:
45 (left_index, right_index) = indices
46 Elementwise.expected_left(operation, right, right_index, Core.checked_observe_dynamic(left, left_index))
47
48law actual_flat_projection:
49 for +plan: Shape.BroadcastPlan
50 for +index: Shape.FlatIndex<Shape.output_shape(plan)>
51 {indices_result(Elementwise.flat_input_indices(plan, Done{index})) ==
52 offset_spec(plan, ShapeLaws.position(Shape.output_shape(plan), index)) : Maybe<&1, U32 & U32>}
53
54law input_indices_refines:
55 for +plan: Shape.BroadcastPlan
56 for +index: U32
57 {indices_result(Elementwise.input_indices(plan, index)) ==
58 offset_spec(plan, U32.to_nat(index)) : Maybe<&1, U32 & U32>}
59
60law input_indices_acceptance:
61 for +plan: Shape.BroadcastPlan
62 for +index: U32
63 {indices_accepted(Elementwise.input_indices(plan, index)) ==
64 Nat.is_lt(U32.to_nat(index), Shape.numel(Shape.extents(Shape.output_shape(plan)))) : Bool}
65
66law expected_projection_refines:
67 for +operation: Elementwise.Operation
68 for +plan: Shape.BroadcastPlan
69 for left: Core.Dynamic
70 for right: Core.Dynamic
71 for +index: U32
72 {Elementwise.expected(operation, plan, left, right, index) ==
73 specified_expected(operation, left, right, offset_spec(plan, U32.to_nat(index))) : Maybe<&2, F32>}