<jupyter_start><jupyter_text>Using PRNG key<jupyter_code>key = random.PRNGKey(0)
X = random.normal(key, (100, 2))
print(X.shape)<jupyter_output>(100, 2)<jupyter_text>VMAP for auto-batching!<jupyter_code>%timeit vmap(pred_yhat, in_axes=(0, None))(X, theta).block_until_ready()
%timeit jnp.stack([pred_yhat(x, theta) for x in X]).block_until_ready()
%timeit X@theta[1:] + theta[0]<jupyter_output>893 µs ± 2.61 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)<jupyter_text>JIT for speedup<jupyter_code>%timeit vmap(jit(pred_yhat), in_axes=(0, None))(X, theta).block_until_ready()
pred_y_hat_vector = lambda X, theta: vmap(jit(pred_yhat), in_axes=(0, None))(X, theta)
def cost(X, y, theta):
    y_hat_vec = pred_y_hat_vector(X, theta)
    error_square = jnp.linalg.norm(y - y_hat_vec, 2)
    return error_square
theta_gt = jnp.array([1., 4., 5.])
y_gt = vmap(jit(pred_yhat), in_axes=(0, None))(X, theta_gt) + 0.2*random.normal(key, (100, 1))
cost(X, y_gt, theta_gt)<jupyter_output><empty_output><jupyter_text>Our initial estimates (theta) are not good<jupyter_code>cost(X, y_gt, theta)
grad_theta = grad(cost, argnums=[2])
lr = 0.001
for i in range(50):
    cost_val = cost(X, y_gt, theta)
    print(i, cost_val)
    grad_theta_val = grad_theta(X, y_gt, theta)[0]
    theta = theta - lr*grad_theta_val
theta
theta_gt<jupyter_output><empty_output><jupyter_text>Using Optax instead of manually writing SGD<jupyter_code>#optimizer = optax.adam(learning_rate=0.01)
optimizer = optax.sgd(learning_rate=0.001)

theta = jnp.array([-1., 2., 3.])
opt_state = optimizer.init(theta)
opt_state
for i in range(50):
    cost_val = cost(X, y_gt, theta)
    print(i, cost_val)
    grad_theta_val = grad_theta(X, y_gt, theta)[0]
    updates, opt_state = optimizer.update(grad_theta_val, opt_state)
    theta = optax.apply_updates(theta, updates)<jupyter_output>0 377.96906
1 365.99265
2 354.02716
3 342.07303
4 330.13083
5 318.20102
6 306.28424
7 294.38098
8 282.49225
9 270.61868
10 258.76102
11 246.92035
12 235.09787
13 223.29465
14 211.51213
15 199.75198
16 188.01602
17 176.30638
18 164.62575
19 152.97719
20 141.36456
21 129.7927
22 118.26772
23 106.79769
24 95.3937
25 84.071465
26 72.85466
27 61.781315
28 50.91854
29 40.401173
30 30.558014
31 22.45481
32 18.97314
33 18.814388
34 18.81421
35 18.814215
36 18.814215
37 18.814215
38 18.814215
39 18.814215
40 18.814215
41 18.814215
42 18.814215
43 18.814215
44 18.814215
45 18.814215
46 18.814215
47 18.814215
48 18.814215
49 18.814215<jupyter_text>Is JAX quicker (even on CPU?!)Gaussian Processes need Cholesky decompositions. Can we get a speedup using JAX instead of Numpy?<jupyter_code>a = np.random.randn(1000, 1000)
b = a.T@a
%timeit np.linalg.cholesky(b)
b = jnp.array(b)
%timeit jnp.linalg.cholesky(b)<jupyter_output>2.19 ms ± 67.9 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
{'fraction_non_alphanumeric': 0.13089373680506686, 'fraction_numerical': 0.1959887403237157, 'mean_word_length': 5.048936170212766, 'pattern_counts': {'":': 0, '<': 21, '<?xml version=': 0, '>': 21, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '6', 'dir': 'jupyter-structured-clean-dedup', 'id': '105390', 'license': 'BSD-3-Clause', 'n_tokens_mistral': 1637, 'n_tokens_neox': 1246, 'n_words': 297, 'path': 'notebooks/jax/.ipynb_checkpoints/introduction-jax-checkpoint.ipynb', 'repo_name': 'prog-ml/prog-ml.github.io', 'size': '2842'}