Python and Viv · Gradients
Gradients by hand and by grad
What to notice
- The Python gradient is derived by hand: for the mean squared error,
2 X^T (X w - y) / n. The Viv gradient isgrad(loss, w, X, y): the compiler writes the reverse pass oflossitself, and every engine runs that generated code. Nothing on the Viv side states the formula. - Both check the gradient against central differences with
H = 1e-6; the Python side does it in one comprehension overnp.eye(2), the Viv side in a loop that builds each unit vector (Viv has noeyeyet). grad(f, w, ...)differentiates with respect to the first argument and passes the rest through unchanged, and the result has the shape ofw.- For one variable, Python writes the derivative
dfout (product rule). Viv'sderiv(fn(x) -> sin(x) * x * x, 1.3)computes it in forward mode. w0.copy()iscopy(w0);X.T @ rwould betranspose(X) @ r, but the Viv side never needs it.- The descent prints the same loss to 8 decimals in both, although the gradients are computed by different formulas: they agree to rounding.
Python gradient.py · 48 lines
Viv gradient.viv · 54 lines
py2viv translated from gradient.py
What py2viv adds before the program
(nothing)
(nothing)
273 lines: a header and the guard functions its translation calls
1###
2Translated from gradient.py by py2viv 0.2.0 (numeric Python -> Viv).
3Python source sha256: 19c4669bc3cffd85c02e378e67bb02ed6207d554be6532e07810cd516f33f30b
4Each print(...) is one print(...) line; fmt(x) writes a float as Python's repr does.
5A `raise`, a failed assert and every guard `stop` with Python's exception line.
6###
7
8#| 0, 1, ..., n - 1.
9def py_iota(n) {
10 a = zeros(n)
11 for i to n {
12 a[i] = i
13 }
14 return a
15}
16
17#| Every entry of a real array, row-major, as a vector (a copy).
18def py_flat(a: Arr) {
19 return take(a, py_iota(numel(a)))
20}
21
22#| A copy of a real array (np.copy, a.copy(), np.array(a)).
23def py_copy(a: Arr) {
24 return a * 1.0
25}
26
27#| An int result, as Python's int: a -0.0 becomes 0 (an int has no sign on zero), and the run
28#| stops at 2^53, past which a double no longer holds every integer (Python's int does).
29def py_int(x) {
30 if abs(x) >= 9007199254740992 {
31 stop("py2viv: an integer reached 2^53; Python's int is exact there and a double is not")
32 }
33 return x + 0.0
34}
35
36#| 0.0 with the sign of w (C's copysign(0.0, w)).
37def py_zero_f(w) {
38 if w < 0 or (w == 0 and 1.0 / w < 0) {
39 return -0.0
40 }
41 return 0.0
42}
43
44#| Floored division (k = 0) or remainder (k = 1) of doubles, as CPython's float_divmod and
45#| NumPy's npy_divmod compute them (fmod first, then the sign fix and the rounding fix).
46def py_fdm(a, b, k) {
47 m = mod(a, b)
48 d = (a - m) / b
49 if m != 0 {
50 if (b < 0) != (m < 0) {
51 m = m + b
52 d = d - 1.0
53 }
54 }
55 else {
56 m = py_zero_f(b)
57 }
58 if k == 1 {
59 return m
60 }
61 if d != 0 {
62 f = floor(d)
63 if d - f > 0.5 {
64 f = f + 1.0
65 }
66 return f
67 }
68 return py_zero_f(a / b)
69}
70
71#| a / b with a float: ZeroDivisionError on a zero divisor.
72def py_div_f(a, b) {
73 if b == 0 {
74 stop("ZeroDivisionError: float division by zero")
75 }
76 return a / b
77}
78
79#| a % b of two ints.
80def py_mod_i(a, b) {
81 if b == 0 {
82 stop("ZeroDivisionError: integer modulo by zero")
83 }
84 return py_int(py_fdm(a, b, 1))
85}
86
87#| NumPy's pairwise sum of f[lo], ..., f[lo + n - 1] (numpy/core/src/umath/loops_utils.h):
88#| below 8 terms left to right from 0; to 128, eight running sums combined as a tree, then
89#| the rest; above, the two halves (the first a multiple of 8 long) summed apart.
90def py_pw(f: Arr, lo, n) {
91 if n < 8 {
92 s = 0.0
93 for i to n {
94 s += f[lo + i]
95 }
96 return s
97 }
98 if n <= 128 {
99 r = zeros(8)
100 for j to 8 {
101 r[j] = f[lo + j]
102 }
103 i = 8
104 while i < n - n % 8 {
105 for j to 8 {
106 r[j] += f[lo + i + j]
107 }
108 i += 8
109 }
110 s = ((r[0] + r[1]) + (r[2] + r[3])) + ((r[4] + r[5]) + (r[6] + r[7]))
111 while i < n {
112 s += f[lo + i]
113 i += 1
114 }
115 return s
116 }
117 n2 = floor(n / 2)
118 n2 = n2 - n2 % 8
119 return py_pw(f, lo, n2) + py_pw(f, lo + n2, n - n2)
120}
121
122#| NumPy's sum of f[lo], ..., f[lo + n - 1] along one axis: from 0.0, a pairwise sum of each
123#| block of NumPy's buffer size (8192 entries), added in order.
124def py_blk(f: Arr, lo, n) {
125 bufsize = 8192
126 s = 0.0
127 b = 0
128 while b < n {
129 s = s + py_pw(f, lo + b, min(bufsize, n - b))
130 b += bufsize
131 }
132 return s
133}
134
135#| np.sum(a), a.sum() of a float array: NumPy's sum over the entries in row-major order (py_blk).
136def py_npsum(a: Arr) {
137 f = py_flat(a)
138 return py_blk(f, 0, len(f))
139}
140
141#| np.isclose(a, b, rtol, atol) of two numbers: |a - b| <= atol + rtol |b| when both are
142#| finite, a == b otherwise (a NaN is never close).
143def py_npisclose(a, b, rtol, atol) {
144 if abs(a) < 1.0 / 0.0 and abs(b) < 1.0 / 0.0 {
145 return abs(a - b) <= atol + rtol * abs(b)
146 }
147 return a == b
148}
149
150#| np.allclose(a, b, rtol, atol) of two arrays of one shape: np.isclose entry by entry.
151def py_allclose(a: Arr, b: Arr, rtol, atol) {
152 fa = py_flat(a)
153 fb = py_flat(b)
154 if len(fa) != len(fb) {
155 stop("py2viv: np.allclose of two arrays of different sizes (NumPy broadcasts them)")
156 }
157 for i to len(fa) {
158 if not py_npisclose(fa[i], fb[i], rtol, atol) {
159 return false
160 }
161 }
162 return true
163}
164
165#| The shape NumPy broadcasts two shapes da and db to (aligned at the last axis, a length 1
166#| stretched); the run stops where NumPy raises ValueError.
167def py_bdims(da: Arr, db: Arr) {
168 n = len(da)
169 if len(db) > n {
170 n = len(db)
171 }
172 e = zeros(n)
173 for i to n {
174 x = 1
175 y = 1
176 if i >= n - len(da) {
177 x = da[i - n + len(da)]
178 }
179 if i >= n - len(db) {
180 y = db[i - n + len(db)]
181 }
182 if x == y or y == 1 {
183 e[i] = x
184 }
185 else {
186 if x == 1 {
187 e[i] = y
188 }
189 else {
190 stop("ValueError: operands could not be broadcast together with shapes")
191 }
192 }
193 }
194 return e
195}
196
197#| The flat indices, in row-major order of the shape e, of the entries of an array of shape d
198#| that broadcasting to e reads (a length 1 of d is read again along that axis of e).
199def py_bidx(d: Arr, e: Arr) {
200 n = len(e)
201 m = len(d)
202 st = zeros(n)
203 s = 1
204 for k to m {
205 j = m - 1 - k
206 if d[j] != 1 {
207 st[j + n - m] = s
208 }
209 s = s * d[j]
210 }
211 idx = zeros(1)
212 for i to n {
213 w = e[i]
214 nxt = zeros(len(idx) * w)
215 for p to len(idx) {
216 for c to w {
217 nxt[p * w + c] = idx[p] + c * st[i]
218 }
219 }
220 idx = nxt
221 }
222 return idx
223}
224
225#| a broadcast to the shape e (a shape py_bdims gave for a and another array): a itself when
226#| it has that shape, else a new array whose entries are copies of a's.
227def py_bto(a: Arr, e: Arr) {
228 d = shape(a)
229 if len(d) == len(e) {
230 if all(d == e) {
231 return a
232 }
233 }
234 return py_shaped(take(a, py_bidx(d, e)), e)
235}
236
237#| a + b of two real arrays, broadcast as NumPy broadcasts them.
238def py_badd(a: Arr, b: Arr) {
239 e = py_bdims(shape(a), shape(b))
240 return py_bto(a, e) + py_bto(b, e)
241}
242
243#| a - b of two real arrays, broadcast.
244def py_bsub(a: Arr, b: Arr) {
245 e = py_bdims(shape(a), shape(b))
246 return py_bto(a, e) - py_bto(b, e)
247}
248
249#| np.allclose(a, b, rtol, atol) of two arrays, broadcast.
250def py_ballclose(a: Arr, b: Arr, rtol, atol) {
251 e = py_bdims(shape(a), shape(b))
252 return py_allclose(py_bto(a, e), py_bto(b, e), rtol, atol)
253}
254
255#| The entries of f, row-major, as an array of shape e (a vector of 1 to 4 lengths).
256def py_shaped(f: Arr, e: Arr) {
257 if len(e) == 1 {
258 return f
259 }
260 if len(e) == 2 {
261 z = zeros(e[0], e[1])
262 }
263 else {
264 if len(e) == 3 {
265 z = zeros(e[0], e[1], e[2])
266 }
267 else {
268 z = zeros(e[0], e[1], e[2], e[3])
269 }
270 }
271 put(z, py_iota(numel(z)), f)
272 return z
273}
1 imports
4import numpy as np
2# nothing to import: grad and deriv are builtins
(nothing)
2 data: points near the line y = 1 + 2 x
7X = np.array([[1.0, 0.0], [1.0, 0.5], [1.0, 1.0], [1.0, 1.5], [1.0, 2.0]])
8y = np.array([1.1, 1.9, 3.2, 3.9, 5.1])
5X = [[1.0, 0.0], [1.0, 0.5], [1.0, 1.0], [1.0, 1.5], [1.0, 2.0]]
6y = [1.1, 1.9, 3.2, 3.9, 5.1]
279X = [[1.0, 0.0], [1.0, 0.5], [1.0, 1.0], [1.0, 1.5], [1.0, 2.0]]
280y = [1.1, 1.9, 3.2, 3.9, 5.1]
3 the loss
11def loss(w, X, y):
12 r = X @ w - y
13 return np.sum(r * r) / len(y)
9def loss(w, X, y) {
10 r = X @ w - y
11 return sum(r * r) / len(y)
12}
282def loss(w: Arr, X: Arr, y: Arr) {
283 r = py_bsub(X @ w, y)
284 return py_div_f(py_npsum(r * r), len(y))
285}
4 its gradient
16def loss_grad(w, X, y):
17 r = X @ w - y
18 return 2.0 * (X.T @ r) / len(y)
15def loss_grad(w, X, y) {
16 return grad(loss, w, X, y)
17}
287def loss_grad(w: Arr, X: Arr, y: Arr) {
288 r = py_bsub(X @ w, y)
289 return 2.0 * (transpose(X) @ r) / len(y)
290}
5 check the gradient at one point
21w0 = np.array([0.5, 0.5])
22H = 1e-6
23fd = np.array([(loss(w0 + H * e, X, y) - loss(w0 - H * e, X, y)) / (2 * H) for e in np.eye(2)])
24print("gradient ", " ".join(f"{g:.6f}" for g in loss_grad(w0, X, y)))
25print("differences", " ".join(f"{g:.6f}" for g in fd))
20w0 = [0.5, 0.5]
21H = 1e-6
22fd = zeros(len(w0))
23for i to len(w0) {
24 e = zeros(len(w0))
25 e[i] = 1.0
26 fd[i] = (loss(w0 + H * e, X, y) - loss(w0 - H * e, X, y)) / (2 * H)
27}
28g0 = loss_grad(w0, X, y)
29print("gradient ", f"{g0[0]:.6f} {g0[1]:.6f}")
30print("differences", f"{fd[0]:.6f} {fd[1]:.6f}")
292w0 = [0.5, 0.5]
293H = 1e-06
294py_s2 = eye(2, 2, 0)
295py_l1 = zeros(len(py_s2))
296for py_i3 to len(py_s2) {
297 py_e4 = slab(py_s2, py_i3)
298 py_l1[py_i3] = py_div_f(loss(py_badd(w0, H * py_e4), X, y) - loss(py_bsub(w0, H * py_e4), X, y), 2 * H)
299}
300fd = py_l1
301py_l5 = []
302py_s6 = loss_grad(w0, X, y)
303for py_i7 to len(py_s6) {
304 py_g8 = py_s6[py_i7]
305 py_l5 = append(py_l5, fmt(py_g8, ".6f"))
306}
307print("gradient ", join(" ", py_l5))
308print("differences", join(" ", [fmt(py_g11, ".6f") for py_g11 in fd]))
6 gradient descent
28RATE = 0.1
29STEPS = 500
30REPORT = 100
31w = w0.copy()
32for step in range(STEPS):
33 w = w - RATE * loss_grad(w, X, y)
34 if (step + 1) % REPORT == 0:
35 print(f"step {step + 1}: loss {loss(w, X, y):.8f}")
36print(f"w = {w[0]:.6f}, {w[1]:.6f}")
33RATE = 0.1
34STEPS = 500
35REPORT = 100
36w = copy(w0)
37for step to STEPS {
38 w = w - RATE * loss_grad(w, X, y)
39 if (step + 1) % REPORT == 0 {
40 print(f"step {step + 1:d}: loss {loss(w, X, y):.8f}")
41 }
42}
43print(f"w = {w[0]:.6f}, {w[1]:.6f}")
310RATE = 0.1
311STEPS = 500
312REPORT = 100
313w = py_copy(w0)
314for step to STEPS {
315 w = py_bsub(w, RATE * loss_grad(w, X, y))
316 if py_mod_i(py_int(step + 1), REPORT) == 0 {
317 print("step ", fmt((py_int(step + 1)) + 0.0, "d"), ": loss ", fmt(loss(w, X, y), ".8f"), sep="")
318 }
319}
320print("w = ", fmt(w[0], ".6f"), ", ", fmt(w[1], ".6f"), sep="")
7 a derivative of one variable
39def f(x):
40 return np.sin(x) * x * x
41
42def df(x):
43 return np.cos(x) * x * x + 2.0 * x * np.sin(x)
44
45print(f"f'(1.3) = {df(1.3):.12f}")
46print(f"f'(1.3) = {deriv(fn(x) -> sin(x) * x * x, 1.3):.12f}")
322def f(x) {
323 return sin(x) * x * x
324}
325def df(x) {
326 return cos(x) * x * x + 2.0 * x * sin(x)
327}
328print("f'(1.3) =", fmt(df(1.3), ".12f"))
8 checks
48assert np.allclose(loss_grad(w0, X, y), fd, atol=1e-6)
49test "gradient" {
50 assert_allclose(g0, fd, 1e-6)
51}
330if not (py_ballclose(loss_grad(w0, X, y), fd, 1e-05, 1e-06)) {
331 stop("AssertionError")
332}
What they print
identical — the two programs print the same text.
Dimmed lines are left out of the comparison. Each Viv test block prints one PASS line; the Python checks the same facts with assert, which prints nothing.
$ python3 gradient.py
gradient -4.080000 -5.580000differences -4.080000 -5.580000step 100: loss 0.01440692step 200: loss 0.01440000step 300: loss 0.01440000step 400: loss 0.01440000step 500: loss 0.01440000w = 1.040000, 2.000000f'(1.3) = 2.957324302460
$ viv run gradient.viv
gradient -4.080000 -5.580000differences -4.080000 -5.580000step 100: loss 0.01440692step 200: loss 0.01440000step 300: loss 0.01440000step 400: loss 0.01440000step 500: loss 0.01440000w = 1.040000, 2.000000f'(1.3) = 2.957324302460PASS gradient
$ viv run gradient.py (py2viv’s translation)
The translation prints exactly what the Python prints.
gradient -4.080000 -5.580000differences -4.080000 -5.580000step 100: loss 0.01440692step 200: loss 0.01440000step 300: loss 0.01440000step 400: loss 0.01440000step 500: loss 0.01440000w = 1.040000, 2.000000f'(1.3) = 2.957324302460