Comment by krapht

Should mention Jax in title, otherwise people won't know why they should use this instead of torch.