Autograd in JAX and PyTorch
Contents
Autograd in JAX and PyTorch¶
Basic Imports¶
import torch
from jax import grad
import jax.numpy as jnp
Creating scalar variables in PyTorch¶
x_torch = torch.autograd.Variable(torch.tensor(1.), requires_grad=True)
y_torch = torch.autograd.Variable(torch.tensor(1.), requires_grad=True)
Creating scalar variables in JAX¶
x_jax = jnp.array(1.)
y_jax = jnp.array(1.)
WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
Defining a loss on scalar inputs¶
def loss(x, y):
return x*x + y*y
Computing the loss on PyTorch input¶
l_torch = loss(x_torch, y_torch)
l_torch
tensor(2., grad_fn=<AddBackward0>)
Computing the loss on JAX input¶
l_jax = loss(x_jax, y_jax)
Computing the gradient on PyTorch input¶
l_torch.backward()
x_torch.grad, y_torch.grad
(tensor(2.), tensor(2.))
Computing the gradient on JAX input¶
grad_loss = grad(loss, argnums=[0, 1])
grad_loss(x_jax, y_jax)
(DeviceArray(2., dtype=float32, weak_type=True),
DeviceArray(2., dtype=float32, weak_type=True))
Repeating the same procedure as above for both libraries but instead using vector function¶
def loss(theta):
return theta.T@theta
theta_torch = torch.autograd.Variable(torch.tensor([1., 1.]), requires_grad=True)
theta_torch
tensor([1., 1.], requires_grad=True)
l = loss(theta_torch)
l
tensor(2., grad_fn=<DotBackward0>)
l.backward()
theta_torch.grad
tensor([2., 2.])
theta_jax = jnp.array([1., 1.])
loss(theta_jax)
DeviceArray(2., dtype=float32)
grad_loss = grad(loss, argnums=[0])
grad_loss(theta_jax)
(DeviceArray([2., 2.], dtype=float32),)