src/shape/flat.bend
bendtorch-core@0.1.0.0 · 56 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/shape/nat_proofs.bend as Nats
4import bendtorch-core@0.1.0.0/src/shape/coordinate_proofs.bend as Coordinates
5import bendtorch-core@0.1.0.0/src/shape/projection_proofs.bend as Projection
6
7def hollow(+output: Shape.Shape, +value: Nat, zero: {0n == Shape.numel(Shape.extents(output)) : Nat}, bound: {True{} == Nat.is_lt(value, Shape.elements(output)) : Bool}) -> Empty:
8 +size = Equal.trans(Nat, Shape.elements(output), Shape.numel(Shape.extents(output)), 0n, Shape.elements_exact(output), Equal.sym(Nat, 0n, Shape.numel(Shape.extents(output)), zero))
9 %Equal.trans(Bool, True{}, Nat.is_lt(value, 0n), False{}, Equal.trans(Bool, True{}, Nat.is_lt(value, Shape.elements(output)), Nat.is_lt(value, 0n), bound, Equal.cong(Nat, Bool, count => Nat.is_lt(value, count), Shape.elements(output), 0n, size)), Projection.lt_zero_false(value)) : Shape.truth(_)
10 Unit{}
11
12def left_bound(+plan: Shape.BroadcastPlan, -rows: List<&2, Shape.Row>, exact: {rows == Shape.index_table(Shape.extents(Shape.left_shape(plan)), Shape.extents(Shape.right_shape(plan)), Shape.extents(Shape.output_shape(plan))) : List<&2, Shape.Row>}, +value: Nat, bound: {True{} == Nat.is_lt(value, Shape.elements(Shape.output_shape(plan))) : Bool}) -> {True{} == Nat.is_lt(Shape.walk_left(rows, value), Shape.elements(Shape.left_shape(plan))) : Bool}:
13 +left = Shape.extents(Shape.left_shape(plan))
14 +right = Shape.extents(Shape.right_shape(plan))
15 +output = Shape.extents(Shape.output_shape(plan))
16 +input = Shape.left_shape(plan)
17 +fits = Coordinates.size_fits(Shape.output_shape(plan), value, bound)
18 +index = Coordinates.fin_from(Shape.numel(output), value, fits)
19 +projected = Shape.project_left(left, right, output, Shape.alignment(plan), Shape.unflatten(output, index))
20 +table = Shape.index_table(left, right, output)
21 %Equal.sym(List<&2, Shape.Row>, rows, table, exact) : {True{} == Nat.is_lt(Shape.walk_left(_, value), Shape.elements(input)) : Bool}
22 %Coordinates.fin_from_value(Shape.numel(output), value, fits) : {True{} == Nat.is_lt(Shape.walk_left(table, _), Shape.elements(input)) : Bool}
23 %Equal.sym(Nat, Shape.walk_left(table, Shape.Fin.value(Shape.numel(output), index)), Shape.row_major(Shape.extents(input), projected), Projection.walk_left_value(left, right, output, Shape.alignment(plan), index)) : {True{} == Nat.is_lt(_, Shape.elements(input)) : Bool}
24 %Equal.sym(Nat, Shape.elements(input), Shape.numel(Shape.extents(input)), Shape.elements_exact(input)) : {True{} == Nat.is_lt(Shape.row_major(Shape.extents(input), projected), _) : Bool}
25 %Nats.lt_true(Shape.row_major(Shape.extents(input), projected), Shape.numel(Shape.extents(input)), Coordinates.row_major_lt(Shape.extents(input), projected)) : {True{} == _ : Bool}
26 {==}
27
28def project_left_index(+plan: Shape.BroadcastPlan, map: Shape.Index<Shape.extents(Shape.left_shape(plan)), Shape.extents(Shape.right_shape(plan)), Shape.extents(Shape.output_shape(plan))>, index: Shape.FlatIndex<Shape.output_shape(plan)>) -> Shape.FlatIndex<Shape.left_shape(plan)>:
29 match map index:
30 case Shape.Indexed{rows, exact} Shape.FIndex{+value, bound}:
31 Shape.FIndex{Shape.walk_left(rows, value), left_bound(plan, rows, exact, value, bound)}
32 case Shape.Unindexed{zero} Shape.FIndex{+value, bound}:
33 Empty.absurd(Shape.FlatIndex<Shape.left_shape(plan)>, hollow(Shape.output_shape(plan), value, zero, bound))
34
35def right_bound(+plan: Shape.BroadcastPlan, -rows: List<&2, Shape.Row>, exact: {rows == Shape.index_table(Shape.extents(Shape.left_shape(plan)), Shape.extents(Shape.right_shape(plan)), Shape.extents(Shape.output_shape(plan))) : List<&2, Shape.Row>}, +value: Nat, bound: {True{} == Nat.is_lt(value, Shape.elements(Shape.output_shape(plan))) : Bool}) -> {True{} == Nat.is_lt(Shape.walk_right(rows, value), Shape.elements(Shape.right_shape(plan))) : Bool}:
36 +left = Shape.extents(Shape.left_shape(plan))
37 +right = Shape.extents(Shape.right_shape(plan))
38 +output = Shape.extents(Shape.output_shape(plan))
39 +input = Shape.right_shape(plan)
40 +fits = Coordinates.size_fits(Shape.output_shape(plan), value, bound)
41 +index = Coordinates.fin_from(Shape.numel(output), value, fits)
42 +projected = Shape.project_right(left, right, output, Shape.alignment(plan), Shape.unflatten(output, index))
43 +table = Shape.index_table(left, right, output)
44 %Equal.sym(List<&2, Shape.Row>, rows, table, exact) : {True{} == Nat.is_lt(Shape.walk_right(_, value), Shape.elements(input)) : Bool}
45 %Coordinates.fin_from_value(Shape.numel(output), value, fits) : {True{} == Nat.is_lt(Shape.walk_right(table, _), Shape.elements(input)) : Bool}
46 %Equal.sym(Nat, Shape.walk_right(table, Shape.Fin.value(Shape.numel(output), index)), Shape.row_major(Shape.extents(input), projected), Projection.walk_right_value(left, right, output, Shape.alignment(plan), index)) : {True{} == Nat.is_lt(_, Shape.elements(input)) : Bool}
47 %Equal.sym(Nat, Shape.elements(input), Shape.numel(Shape.extents(input)), Shape.elements_exact(input)) : {True{} == Nat.is_lt(Shape.row_major(Shape.extents(input), projected), _) : Bool}
48 %Nats.lt_true(Shape.row_major(Shape.extents(input), projected), Shape.numel(Shape.extents(input)), Coordinates.row_major_lt(Shape.extents(input), projected)) : {True{} == _ : Bool}
49 {==}
50
51def project_right_index(+plan: Shape.BroadcastPlan, map: Shape.Index<Shape.extents(Shape.left_shape(plan)), Shape.extents(Shape.right_shape(plan)), Shape.extents(Shape.output_shape(plan))>, index: Shape.FlatIndex<Shape.output_shape(plan)>) -> Shape.FlatIndex<Shape.right_shape(plan)>:
52 match map index:
53 case Shape.Indexed{rows, exact} Shape.FIndex{+value, bound}:
54 Shape.FIndex{Shape.walk_right(rows, value), right_bound(plan, rows, exact, value, bound)}
55 case Shape.Unindexed{zero} Shape.FIndex{+value, bound}:
56 Empty.absurd(Shape.FlatIndex<Shape.right_shape(plan)>, hollow(Shape.output_shape(plan), value, zero, bound))