traversal_grouped.bend
stelliferous@0.0.3.0 · 101 lines · 0 definitions
1# Shape-only groups, ordered folds from +0, then adjacent-pair trees. The
2# reader is affine; the same traversal serves weight and bias gradients.
3import Base
4import stelliferous@0.0.3.0/traversal_loop.bend as Loop
5
6type Range is Data:
7 Range{first: Nat,count: Nat}
8
9def rows(remaining: Nat,+first: Nat,+height: Nat,+width: Nat,+span: Nat) -> List<&2,Range>:
10 match remaining:
11 case 0n: []
12 case 1n+rest:
13 Range{Nat.mul(first,width),Nat.mul(Nat.min(span,Nat.sub(height,first)),width)} <> rows(rest,Nat.add(first,span),height,width,span)
14
15def columns(remaining: Nat,+row: Nat,+first: Nat,+width: Nat,+group: Nat) -> List<&2,Range>:
16 match remaining:
17 case 0n: []
18 case 1n+rest:
19 Range{Nat.add(Nat.mul(row,width),first),Nat.min(group,Nat.sub(width,first))} <> columns(rest,row,Nat.add(first,group),width,group)
20
21def cat(left: List<&2,Range>,+right: List<&2,Range>) -> List<&2,Range>:
22 match left:
23 case []: right
24 case head <> tail: head <> cat(tail,right)
25
26def wide_rows(remaining: Nat,+row: Nat,+width: Nat,+group: Nat) -> List<&2,Range>:
27 match remaining:
28 case 0n: []
29 case 1n+rest:
30 cat(columns(Nat.div(Nat.add(width,Nat.sub(group,1n)),group),row,0n,width,group),wide_rows(rest,1n+row,width,group))
31
32def ranges_selected(wide: Bool,+height: Nat,+width: Nat,+group: Nat) -> List<&2,Range>:
33 match wide:
34 case True{}: wide_rows(height,0n,width,group)
35 case False{}:
36 +span = Nat.div(group,width)
37 rows(Nat.div(Nat.add(height,Nat.sub(span,1n)),span),0n,height,width,span)
38
39def ranges(+height: Nat,+width: Nat,+group: Nat) -> List<&2,Range>:
40 ranges_selected(Nat.is_lt(group,width),height,width,group)
41
42def pairs(~ownership: Quant,~Element: Kind(ownership),~Context: Data,
43 ~combine: Context -> Element -> Element -> Element,+context: Context,values: List<ownership,Element>) -> List<ownership,Element>:
44 match values:
45 case []: []
46 case a <> []: [a]
47 case a <> b <> tail: combine(context,a,b) <> pairs(~ownership,~Element,~Context,~combine,context,tail)
48
49# The original length bounds the number of tree levels. A singleton is
50# returned directly, including its sign bit; there is no final addition.
51def levels(~ownership: Quant,~Element: Kind(ownership),~Context: Data,
52 ~zero: Context -> Element,~combine: Context -> Element -> Element -> Element,
53 fuel: Nat,+context: Context,values: List<ownership,Element>) -> Element:
54 match fuel values:
55 case 0n Nil{}: zero(context)
56 case 0n head <> tail: head
57 case 1n+rest Nil{}: zero(context)
58 case 1n+rest head <> []: head
59 case 1n+rest a <> b <> tail: levels(~ownership,~Element,~Context,~zero,~combine,rest,context,pairs(~ownership,~Element,~Context,~combine,context,a <> b <> tail))
60
61def tree(+values: List<&2,F32>) -> F32:
62 levels(~&2,~F32,~Unit,~(context => 0.0),~(context => F32.add),List.length(&2,F32,values),Unit{},values)
63
64type State<-Reader: Type,-Context: Data> is Type:
65 State{reader: Reader,context: Context,position: Nat,total: F32}
66
67def advanced(~Reader: Type,~Context: Data,+context: Context,position: Nat,total: F32,pair: Reader & F32) -> State<Reader,Context>:
68 (reader,value) = pair
69 State{reader,context,1n+position,(total + value : F32)}
70
71def step(~Reader: Type,~Context: Data,~read: Reader -> Context -> Nat -> Reader & F32,state: State<Reader,Context>) -> State<Reader,Context>:
72 State{reader,+context,+position,total} = state
73 advanced(~Reader,~Context,context,position,total,read(reader,context,position))
74
75def finished(~Reader: Type,~Context: Data,state: State<Reader,Context>) -> Reader & F32:
76 State{reader,context,position,total} = state
77 (reader,total)
78
79def fold(~Reader: Type,~Context: Data,~read: Reader -> Context -> Nat -> Reader & F32,count: Nat,first: Nat,reader: Reader,context: Context) -> Reader & F32:
80 finished(~Reader,~Context,Loop.run(~State<Reader,Context>,~step(~Reader,~Context,~read),count,State{reader,context,first,0.0}))
81
82def prepend(~Reader: Type,value: F32,pair: Reader & List<&2,F32>) -> Reader & List<&2,F32>:
83 (reader,values) = pair
84 (reader,value <> values)
85
86def grouped_tail(~Reader: Type,pair: Reader & F32,tail: Reader -> Reader & List<&2,F32>) -> Reader & List<&2,F32>:
87 (reader,value) = pair
88 prepend(~Reader,value,tail(reader))
89
90def fold_ranges(~Reader: Type,~Context: Data,~read: Reader -> Context -> Nat -> Reader & F32,ranges: List<&2,Range>,reader: Reader,+context: Context) -> Reader & List<&2,F32>:
91 match ranges:
92 case []: (reader,[])
93 case Range{first,count} <> rest:
94 grouped_tail(~Reader,fold(~Reader,~Context,~read,count,first,reader,context),reader => fold_ranges(~Reader,~Context,~read,rest,reader,context))
95
96def reduced(~Reader: Type,pair: Reader & List<&2,F32>) -> Reader & F32:
97 (reader,values) = pair
98 (reader,tree(values))
99
100def run(~Reader: Type,~Context: Data,~read: Reader -> Context -> Nat -> Reader & F32,ranges: List<&2,Range>,reader: Reader,context: Context) -> Reader & F32:
101 reduced(~Reader,fold_ranges(~Reader,~Context,~read,ranges,reader,context))