src/tensor/axis_spec.bend
bendtorch-core@0.1.0.0 · 108 lines · 18 definitions
1import Base
2import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
3
4def first(-A: Data, values: List<&2, A>) -> Maybe<&2, A>:
5 match values:
6 case Nil{}:
7 None{}
8 case head <> tail:
9 Some{head}
10
11def lookup(-A: Data, values: List<&2, A>, index: Nat) -> Maybe<&2, A>:
12 first(A, List.drop(&2, A, values, index))
13
14def prepend(-A: Data, head: Maybe<&2, A>, tail: Maybe<&2, List<&2, A>>) -> Maybe<&2, List<&2, A>>:
15 match head tail:
16 case Some{value} Some{values}:
17 Some{value <> values}
18 case _ _:
19 None{}
20
21def gather(-A: Data, +values: List<&2, A>, axes: List<&2, Nat>) -> Maybe<&2, List<&2, A>>:
22 match axes:
23 case Nil{}:
24 Some{Nil{}}
25 case axis <> tail:
26 prepend(A, lookup(A, values, axis), gather(A, values, tail))
27
28def position_next(value: Maybe<&2, Nat>) -> Maybe<&2, Nat>:
29 match value:
30 case None{}:
31 None{}
32 case Some{index}:
33 Some{1n+index}
34
35def position_decided(decision: Bool, rec: Maybe<&2, Nat>) -> Maybe<&2, Nat>:
36 match decision:
37 case True{}:
38 Some{0n}
39 case False{}:
40 position_next(rec)
41
42def position(axes: List<&2, Nat>, +axis: Nat) -> Maybe<&2, Nat>:
43 match axes:
44 case Nil{}:
45 None{}
46 case head <> tail:
47 position_decided(Nat.is_eq(head, axis), position(tail, axis))
48
49def select(-A: Data, +values: List<&2, A>, index: Maybe<&2, Nat>) -> Maybe<&2, A>:
50 match index:
51 case None{}:
52 None{}
53 case Some{value}:
54 lookup(A, values, value)
55
56def scatter(-A: Data, rank: Nat, +axes: List<&2, Nat>, +coordinates: List<&2, A>, +cursor: Nat) -> Maybe<&2, List<&2, A>>:
57 match rank:
58 case 0n:
59 Some{Nil{}}
60 case 1n+p:
61 prepend(A, select(A, coordinates, position(axes, cursor)), scatter(A, p, axes, coordinates, 1n+cursor))
62
63def indices(rank: Nat, +cursor: Nat) -> List<&2, Nat>:
64 match rank:
65 case 0n:
66 Nil{}
67 case 1n+p:
68 cursor <> indices(p, 1n+cursor)
69
70def inverse(+rank: Nat, +axes: List<&2, Nat>) -> Maybe<&2, List<&2, Nat>>:
71 scatter(Nat, rank, axes, indices(rank, 0n), 0n)
72
73def values(-A: Data, result: Maybe<&2, List<&2, A>>) -> List<&2, A>:
74 match result:
75 case None{}:
76 Nil{}
77 case Some{items}:
78 items
79
80def order(-A: Data, +items: List<&2, A>, axes: List<&2, Nat>) -> List<&2, A>:
81 values(A, gather(A, items, axes))
82
83def inverse_axes(+rank: Nat, axes: List<&2, Nat>) -> List<&2, Nat>:
84 values(Nat, inverse(rank, axes))
85
86def contains(+axis: Nat, axes: List<&2, Nat>) -> Bool:
87 match axes:
88 case Nil{}:
89 False{}
90 case head <> tail:
91 Bool.or(Nat.is_eq(axis, head), contains(axis, tail))
92
93def unique(axes: List<&2, Nat>) -> Bool:
94 match axes:
95 case Nil{}:
96 True{}
97 case +head <> +tail:
98 Bool.and(Bool.not(contains(head, tail)), unique(tail))
99
100def bounded(axes: List<&2, Nat>, +rank: Nat) -> Bool:
101 match axes:
102 case Nil{}:
103 True{}
104 case head <> tail:
105 Bool.and(Nat.is_lt(head, rank), bounded(tail, rank))
106
107def valid(+rank: Nat, +axes: List<&2, Nat>) -> Bool:
108 Bool.and(Nat.is_eq(List.length(&2, Nat, axes), rank), Bool.and(bounded(axes, rank), unique(axes)))