<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 Automatic differentiation with JAX
Main features
Numpy wrapper
Auto-vectorization
Auto-parallelization (SPMD paradigm)
Auto-differentiation
XLA backend and JIT support
How to compute gradient of your objective?
Define it as a standard Python function
Call jax.grad and voila!
Do not forget to wrap these functions with jax.jit to speed up
Step1:  By default, JAX exploits single-precision numbers float32
You can enable double precision (float64) by hands.
Step2:  Random numbers in JAX
JAX focuses on the reproducibility of the runs
Analogue of random seed is the necessary argument of all functions that generate something random
More details and references on the design of random submodule are here
Step3:  Forward mode vs. backward mode
Step4:  Forward mode vs. backward mode
Step5:  Hessian-by-vector product
<ASSISTANT_TASK:>
Python Code:
import jax
import jax.numpy as jnp
Explanation: Automatic differentiation with JAX
Main features
Numpy wrapper
Auto-vectorization
Auto-parallelization (SPMD paradigm)
Auto-differentiation
XLA backend and JIT support
How to compute gradient of your objective?
Define it as a standard Python function
Call jax.grad and voila!
Do not forget to wrap these functions with jax.jit to speed up
End of explanation
from jax.config import config
config.update("jax_enable_x64", True)
@jax.jit
def f(x, A, b):
    res = A @ x - b
    return res @ res
gradf = jax.grad(f, argnums=0, has_aux=False)
Explanation: By default, JAX exploits single-precision numbers float32
You can enable double precision (float64) by hands.
End of explanation
n = 1000
x = jax.random.normal(jax.random.PRNGKey(0), (n, ))
A = jax.random.normal(jax.random.PRNGKey(0), (n, n))
b = jax.random.normal(jax.random.PRNGKey(0), (n, ))
print("Check correctness", jnp.linalg.norm(gradf(x, A, b) - 2 * A.T @ (A @ x - b)))
print("Compare speed")
print("Analytical gradient")
%timeit 2 * A.T @ (A @ x - b)
print("Grad function")
%timeit gradf(x, A, b).block_until_ready()
jit_gradf = jax.jit(gradf)
print("Jitted grad function")
%timeit jit_gradf(x, A, b).block_until_ready()
hess_func = jax.jit(jax.hessian(f))
print("Check correctness", jnp.linalg.norm(2 * A.T @ A - hess_func(x, A, b)))
print("Time for hessian")
%timeit hess_func(x, A, b).block_until_ready()
print("Emulate hessian and check correctness", 
      jnp.linalg.norm(jax.jit(hess_func)(x, A, b) - jax.jacfwd(jax.jacrev(f))(x, A, b)))
print("Time of emulating hessian")
hess_umul_func = jax.jit(jax.jacfwd(jax.jacrev(f)))
%timeit hess_umul_func(x, A, b).block_until_ready()
Explanation: Random numbers in JAX
JAX focuses on the reproducibility of the runs
Analogue of random seed is the necessary argument of all functions that generate something random
More details and references on the design of random submodule are here
End of explanation
fmode_f = jax.jit(jax.jacfwd(f))
bmode_f = jax.jit(jax.jacrev(f))
print("Check correctness", jnp.linalg.norm(fmode_f(x, A, b) - bmode_f(x, A, b)))
print("Forward mode")
%timeit fmode_f(x, A, b).block_until_ready()
print("Backward mode")
%timeit bmode_f(x, A, b).block_until_ready()
Explanation: Forward mode vs. backward mode: $m \ll n$
End of explanation
def fvec(x, A, b):
    y = A @ x + b
    return jnp.exp(y - jnp.max(y)) / jnp.sum(jnp.exp(y - jnp.max(y)))
grad_fvec = jax.jit(jax.grad(fvec))
jac_fvec = jax.jacobian(fvec)
fmode_fvec = jax.jit(jax.jacfwd(fvec))
bmode_fvec = jax.jit(jax.jacrev(fvec))
n = 1000
m = 1000
x = jax.random.normal(jax.random.PRNGKey(0), (n, ))
A = jax.random.normal(jax.random.PRNGKey(0), (m, n))
b = jax.random.normal(jax.random.PRNGKey(0), (m, ))
J = jac_fvec(x, A, b)
print(J.shape)
grad_fvec(x, A, b)
print("Check correctness", jnp.linalg.norm(fmode_fvec(x, A, b) - bmode_fvec(x, A, b)))
print("Check shape", fmode_fvec(x, A, b).shape, bmode_fvec(x, A, b).shape)
print("Time forward mode")
%timeit fmode_fvec(x, A, b).block_until_ready()
print("Time backward mode")
%timeit bmode_fvec(x, A, b).block_until_ready()
n = 10
m = 1000
x = jax.random.normal(jax.random.PRNGKey(0), (n, ))
A = jax.random.normal(jax.random.PRNGKey(0), (m, n))
b = jax.random.normal(jax.random.PRNGKey(0), (m, ))
print("Check correctness", jnp.linalg.norm(fmode_fvec(x, A, b) - bmode_fvec(x, A, b)))
print("Check shape", fmode_fvec(x, A, b).shape, bmode_fvec(x, A, b).shape)
print("Time forward mode")
%timeit fmode_fvec(x, A, b).block_until_ready()
print("Time backward mode")
%timeit bmode_fvec(x, A, b).block_until_ready()
Explanation: Forward mode vs. backward mode: $m \geq n$
End of explanation
def hvp(f, x, z, *args):
    def g(x):
        return f(x, *args)
    return jax.jvp(jax.grad(g), (x,), (z,))[1]
n = 3000
x = jax.random.normal(jax.random.PRNGKey(0), (n, ))
A = jax.random.normal(jax.random.PRNGKey(0), (n, n))
b = jax.random.normal(jax.random.PRNGKey(0), (n, ))
z = jax.random.normal(jax.random.PRNGKey(0), (n, ))
print("Check correctness", jnp.linalg.norm(2 * A.T @ (A @ z) - hvp(f, x, z, A, b)))
print("Time for hvp by hands")
%timeit (2 * A.T @ (A @ z)).block_until_ready()
print("Time for hvp via jvp, NO jit")
%timeit hvp(f, x, z, A, b).block_until_ready()
print("Time for hvp via jvp, WITH jit")
%timeit jax.jit(hvp, static_argnums=0)(f, x, z, A, b).block_until_ready()
Explanation: Hessian-by-vector product
End of explanation
<END_TASK>