~/jendhub

src/autograd/gradient.bend

bendtorch-core@0.1.0.0 · 150 lines · 22 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/main.bend as Storage
4import bendtorch-core@0.1.0.0/src/tensor/core.bend as Core
5import bendtorch-core@0.1.0.0/src/tensor/layout.bend as Layout
6
7# Gradient slots for reverse differentiation. A slot starts Unset. The first
8# contribution whose dimensions agree with the slot's shape moves in unchanged;
9# every later one is added in place, position by position in ascending flat
10# order: total[i] = total[i] + contribution[i], one F32 add per position, never
11# an overwrite. A contribution of any other shape is refused, and the slot and
12# the contribution come back unchanged. Like PyTorch's .grad, grad reads an
13# Unset slot as None; gradient instead materializes zeros of the slot's shape.
14# reset returns a slot to Unset, like PyTorch's zero_grad(set_to_none=True).
15
16type Error is Data:
17  TensorError{error: Core.Error}
18  ShapeMismatch{}
19
20type Slot is Type:
21  Unset{}
22  Held{gradient: Core.Dynamic}
23
24type Outcome is Type:
25  Accumulated{slot: Slot}
26  Refused{slot: Slot, contribution: Core.Dynamic, error: Error}
27
28type State<-shape: Shape.Shape> is Type:
29  State{contribution: Core.Dynamic, total: Core.Tensor<shape>}
30
31type Step<-shape: Shape.Shape> is Type:
32  Continue{state: State<shape>}
33  Stop{contribution: Core.Dynamic, total: Core.Tensor<shape>, error: Core.Error}
34
35def step_write(-shape: Shape.Shape, contribution: Core.Dynamic,
36  pair: Core.Tensor<shape> & Result<&2, &2, Core.Error, Unit>) -> Step<shape>:
37  (total, result) = pair
38  match result:
39    case Fail{error}:
40      Stop{contribution, total, error}
41    case Done{unit}:
42      Continue{State{contribution, total}}
43
44def step_total(-shape: Shape.Shape, +index: U32, contribution: Core.Dynamic, value: F32,
45  pair: Core.Tensor<shape> & Result<&2, &2, Core.Error, F32>) -> Step<shape>:
46  (total, result) = pair
47  match result:
48    case Fail{error}:
49      Stop{contribution, total, error}
50    case Done{+accumulated}:
51      step_write(shape, contribution, Core.write_flat(shape, total, index, F32.add(accumulated, value)))
52
53def step_contribution(-shape: Shape.Shape, +index: U32, total: Core.Tensor<shape>,
54  pair: Core.Dynamic & Result<&2, &2, Core.Error, F32>) -> Step<shape>:
55  (contribution, result) = pair
56  match result:
57    case Fail{error}:
58      Stop{contribution, total, error}
59    case Done{value}:
60      step_total(shape, index, contribution, value, Core.read_flat(shape, total, index))
61
62def step(-shape: Shape.Shape, +index: U32, state: State<shape>) -> Step<shape>:
63  State{contribution, total} = state
64  step_contribution(shape, index, total, Core.read_dynamic(contribution, index))
65
66def continue_step(+shape: Shape.Shape, rec: State<shape> -> Outcome, result: Step<shape>) -> Outcome:
67  match result:
68    case Stop{contribution, total, error}:
69      Refused{Held{Core.Packed{shape, total}}, contribution, TensorError{error}}
70    case Continue{state}:
71      rec(state)
72
73def conversion_error(+shape: Shape.Shape, state: State<shape>) -> Outcome:
74  State{contribution, total} = state
75  Refused{Held{Core.Packed{shape, total}}, contribution, TensorError{Core.StorageError{Storage.InvalidIndex{}}}}
76
77def decided(+shape: Shape.Shape, +cursor: Nat, +index: U32, state: State<shape>, rec: State<shape> -> Outcome,
78  decision: Shape.BoolCase<Nat.is_eq(U32.to_nat(index), cursor)>) -> Outcome:
79  match decision:
80    case Shape.BFalse{equal}:
81      conversion_error(shape, state)
82    case Shape.BTrue{equal}:
83      continue_step(shape, rec, step(shape, index, state))
84
85def loop(fuel: Nat, +shape: Shape.Shape, +cursor: Nat, state: State<shape>) -> Outcome:
86  match fuel:
87    case 0n:
88      State{contribution, total} = state
89      Accumulated{Held{Core.Packed{shape, total}}}
90    case 1n+remaining:
91      +index = U32.from_nat(cursor)
92      decided(shape, cursor, index, state, state => loop(remaining, shape, 1n+cursor, state),
93        Shape.bool_case(Nat.is_eq(U32.to_nat(index), cursor)))
94
95def held_checked(contribution: Core.Dynamic, +total_shape: Shape.Shape, total: Core.Tensor<total_shape>, same: Bool) -> Outcome:
96  match same:
97    case False{}:
98      Refused{Held{Core.Packed{total_shape, total}}, contribution, ShapeMismatch{}}
99    case True{}:
100      loop(Shape.elements(total_shape), total_shape, 0n, State{contribution, total})
101
102def held(+shape: Shape.Shape, contribution: Core.Dynamic, gradient: Core.Dynamic) -> Outcome:
103  Core.Packed{+total_shape, total} = gradient
104  held_checked(contribution, total_shape, total, Layout.words_same(Shape.dimensions(total_shape), Shape.dimensions(shape)))
105
106def agreed(+shape: Shape.Shape, slot: Slot, contribution: Core.Dynamic) -> Outcome:
107  match slot:
108    case Unset{}:
109      Accumulated{Held{contribution}}
110    case Held{gradient}:
111      held(shape, contribution, gradient)
112
113def checked(+shape: Shape.Shape, slot: Slot, +contribution_shape: Shape.Shape, tensor: Core.Tensor<contribution_shape>, same: Bool) -> Outcome:
114  match same:
115    case False{}:
116      Refused{slot, Core.Packed{contribution_shape, tensor}, ShapeMismatch{}}
117    case True{}:
118      agreed(shape, slot, Core.Packed{contribution_shape, tensor})
119
120def accumulate(+shape: Shape.Shape, slot: Slot, contribution: Core.Dynamic) -> Outcome:
121  Core.Packed{+contribution_shape, tensor} = contribution
122  checked(shape, slot, contribution_shape, tensor,
123    Layout.words_same(Shape.dimensions(contribution_shape), Shape.dimensions(shape)))
124
125# The gradient a slot holds, or None when it never received one or was reset.
126def grad(slot: Slot) -> Maybe<&1, Core.Dynamic>:
127  match slot:
128    case Unset{}:
129      None{}
130    case Held{gradient}:
131      Some{gradient}
132
133# Drops a held gradient: the next contribution moves in unchanged.
134def reset(slot: Slot) -> Slot:
135  Unset{}
136
137# The gradient a slot holds; zeros of shape when it never received one.
138def gradient(+shape: Shape.Shape, slot: Slot) -> Result<&2, &1, Core.Error, Core.Dynamic>:
139  match slot:
140    case Unset{}:
141      Core.construct_values(shape, Core.new(shape, 0.0))
142    case Held{gradient}:
143      Done{gradient}
144
145def error_code(error: Error) -> U32:
146  match error:
147    case TensorError{error}:
148      Core.error_code(error)
149    case ShapeMismatch{}:
150      102