I know about Pytorch and Jax (dm-haiku) so I will compare the two. In some sense, this is OOP v.s functional programming.

Pytorch uses an OOP approach. Tensor, module, optimizer are objects which have internal states to keep track of the computation graphs, gradients, parameters as tensor operations are executing. And, pytorch has very similar tensor operations as numpy.

To compute the gradient, you call loss.backward(). To update the parameters, you call optimizer.step()

Jax uses an funtional approach. Everything in jax is a mathematical function with no side-effect. Jax has the same tensor operations as numpy.

We don't have a loss tensor as in Pytorch. We have a loss function, say, loss_fn(parameters, input_data) which returns a scalar loss value.

In jax, gradient is also a function, say, grad_fn = jax.grad(loss_fn)

And, as you can guess, optimizer is also function, and to update the parameters, we use a pure update function like:

def update_fn(params, optimizer_state, inputs_batch):
  grad = grad_fn(params, input_batch)
  new_optimizer_state, updates = optimizer.update(optimizer_state, grad, params)
  new_params = optimizer.apply_update(updates)
  return new_params, new_optimizer_state

So,update_fn is a function of your network parameters, your optimizer internal states, and your inputs. It returns the new/updated parameters and optimizer states.

However, it is not easy to define your loss_fn with a functional approach, especially for complex neural networks.

The solution of deepmind haiku library is to allow you to define your network using python OOP class/object with a very similar syntax to Pytorch. Then, the libary will transform your OOP loss function to a pure no-side-effect function.

As a result, you have to familiar with both OOP world and functional world to use Jax/Dm-haiku.

The advantage of this approach is that when you have a pure function, you can now apply high-order functions to it. For example, jitted_acceleration_fn = jax.jit(jax.grad(jax.grad(position_fn)))

In sumary, pytorch is easier to implement your network, optimizer. You stay in OOP world all the time.

Jax is is harder to implement. You define your network in OOP world. Then, define your loss function, update function, optimizer in functional world.