src/tensor/contracts/fold.bend
bendtorch-core@0.1.0.0 · 128 lines · 15 definitions
1import Base
2import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
3import bendtorch-core@0.1.0.0/src/storage/LAWS.bend as StorageLaws
4import bendtorch-core@0.1.0.0/src/tensor/core.bend as Core
5import bendtorch-core@0.1.0.0/src/tensor/contracts/replayed.bend as Replayed
6
7
8
9
10
11
12
13def chosen(~O: Data, ~g: O -> F32 -> F32 -> F32, +operation: O, kept: Bool, accumulator: F32, word: F32) -> F32:
14 match kept:
15 case False{}:
16 accumulator
17 case True{}:
18 g(operation, accumulator, word)
19
20def step(~O: Data, ~g: O -> F32 -> F32 -> F32, +operation: O, +kept: Bool,
21 rec: Core.Dynamic -> F32 -> Core.Dynamic & Result<&2, &1, Core.Error, F32>,
22 accumulator: F32, pair: Core.Dynamic & Result<&2, &2, Core.Error, F32>) -> Core.Dynamic & Result<&2, &1, Core.Error, F32>:
23 (input, result) = pair
24 match result:
25 case Fail{error}:
26 (input, Fail{error})
27 case Done{x}:
28 rec(input, chosen(~O, ~g, operation, kept, accumulator, x))
29
30def loop(~O: Data, ~g: O -> F32 -> F32 -> F32, ~P: Data, ~ia: P -> Nat -> Nat, ~keep: P -> Nat -> Bool,
31 fuel: Nat, +operation: O, +plan: P, +cursor: Nat, input: Core.Dynamic, accumulator: F32) ->
32 Core.Dynamic & Result<&2, &1, Core.Error, F32>:
33 match fuel:
34 case 0n:
35 (input, Done{accumulator})
36 case 1n+rest:
37 step(~O, ~g, operation, keep(plan, cursor), input => value => loop(~O, ~g, ~P, ~ia, ~keep, rest, operation, plan, 1n+cursor, input, value),
38 accumulator, Core.read_dynamic(input, U32.from_nat(ia(plan, cursor))))
39
40def run(~O: Data, ~g: O -> F32 -> F32 -> F32, ~P: Data, ~ia: P -> Nat -> Nat, ~keep: P -> Nat -> Bool,
41 +operation: O, +plan: P, +count: Nat, input: Core.Dynamic) -> Core.Dynamic & Result<&2, &1, Core.Error, F32>:
42 loop(~O, ~g, ~P, ~ia, ~keep, count, operation, plan, 0n, input, 0.0)
43
44def pair_right(~O: Data, ~g: O -> F32 -> F32 -> F32 -> F32, +operation: O,
45 rec: Core.Dynamic -> Core.Dynamic -> F32 -> (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, F32>,
46 left: Core.Dynamic, accumulator: F32, x: F32, pair: Core.Dynamic & Result<&2, &2, Core.Error, F32>) ->
47 (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, F32>:
48 (right, result) = pair
49 match result:
50 case Fail{error}:
51 ((left, right), Fail{error})
52 case Done{y}:
53 rec(left, right, g(operation, accumulator, x, y))
54
55def pair_left(~O: Data, ~g: O -> F32 -> F32 -> F32 -> F32, +operation: O, +right_index: U32,
56 rec: Core.Dynamic -> Core.Dynamic -> F32 -> (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, F32>,
57 right: Core.Dynamic, accumulator: F32, pair: Core.Dynamic & Result<&2, &2, Core.Error, F32>) ->
58 (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, F32>:
59 (left, result) = pair
60 match result:
61 case Fail{error}:
62 ((left, right), Fail{error})
63 case Done{x}:
64 pair_right(~O, ~g, operation, rec, left, accumulator, x, Core.read_dynamic(right, right_index))
65
66def pair_loop(~O: Data, ~g: O -> F32 -> F32 -> F32 -> F32, ~P: Data, ~ia: P -> Nat -> Nat, ~ib: P -> Nat -> Nat,
67 fuel: Nat, +operation: O, +plan: P, +cursor: Nat, left: Core.Dynamic, right: Core.Dynamic, accumulator: F32) ->
68 (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, F32>:
69 match fuel:
70 case 0n:
71 ((left, right), Done{accumulator})
72 case 1n+rest:
73 pair_left(~O, ~g, operation, U32.from_nat(ib(plan, cursor)),
74 left => right => value => pair_loop(~O, ~g, ~P, ~ia, ~ib, rest, operation, plan, 1n+cursor, left, right, value),
75 right, accumulator, Core.read_dynamic(left, U32.from_nat(ia(plan, cursor))))
76
77def pair_run(~O: Data, ~g: O -> F32 -> F32 -> F32 -> F32, ~P: Data, ~ia: P -> Nat -> Nat, ~ib: P -> Nat -> Nat,
78 +operation: O, +plan: P, +count: Nat, left: Core.Dynamic, right: Core.Dynamic) ->
79 (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, F32>:
80 pair_loop(~O, ~g, ~P, ~ia, ~ib, count, operation, plan, 0n, left, right, 0.0)
81
82def total(~O: Data, ~g: O -> F32 -> F32 -> F32, ~P: Data, ~ia: P -> Nat -> Nat, ~keep: P -> Nat -> Bool,
83 +operation: O, +plan: P, +shape: Shape.Shape, +fill: F32, +ops: List<&2, StorageLaws.Op>, +count: Nat) -> F32:
84 match count:
85 case 0n:
86 0.0
87 case 1n+previous:
88 chosen(~O, ~g, operation, keep(plan, previous), total(~O, ~g, ~P, ~ia, ~keep, operation, plan, shape, fill, ops, previous),
89 Replayed.value(shape, fill, ops, U32.from_nat(ia(plan, previous))))
90
91def pair_total(~O: Data, ~g: O -> F32 -> F32 -> F32 -> F32, ~P: Data, ~ia: P -> Nat -> Nat, ~ib: P -> Nat -> Nat,
92 +operation: O, +plan: P, +left: Shape.Shape, +left_fill: F32, +left_ops: List<&2, StorageLaws.Op>,
93 +right: Shape.Shape, +right_fill: F32, +right_ops: List<&2, StorageLaws.Op>, +count: Nat) -> F32:
94 match count:
95 case 0n:
96 0.0
97 case 1n+previous:
98 g(operation, pair_total(~O, ~g, ~P, ~ia, ~ib, operation, plan, left, left_fill, left_ops, right, right_fill, right_ops, previous),
99 Replayed.value(left, left_fill, left_ops, U32.from_nat(ia(plan, previous))),
100 Replayed.value(right, right_fill, right_ops, U32.from_nat(ib(plan, previous))))
101
102def value(result: Core.Dynamic & Result<&2, &1, Core.Error, F32>) -> Maybe<&2, F32>:
103 (input, outcome) = result
104 match outcome:
105 case Fail{error}:
106 None{}
107 case Done{word}:
108 Some{word}
109
110def input_point(result: Core.Dynamic & Result<&2, &1, Core.Error, F32>, +index: U32) -> Maybe<&2, F32>:
111 (input, outcome) = result
112 Core.dynamic_read_value(Core.read_dynamic(input, index))
113
114def pair_value(result: (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, F32>) -> Maybe<&2, F32>:
115 (inputs, outcome) = result
116 match outcome:
117 case Fail{error}:
118 None{}
119 case Done{word}:
120 Some{word}
121
122def left_point(result: (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, F32>, +index: U32) -> Maybe<&2, F32>:
123 ((left, right), outcome) = result
124 Core.dynamic_read_value(Core.read_dynamic(left, index))
125
126def right_point(result: (Core.Dynamic & Core.Dynamic) & Result<&2, &1, Core.Error, F32>, +index: U32) -> Maybe<&2, F32>:
127 ((left, right), outcome) = result
128 Core.dynamic_read_value(Core.read_dynamic(right, index))