[alevskaya]

It's still early days for JAX, but there's neural tangents https://arxiv.org/abs/1912.02803 and reformer https://arxiv.org/abs/2001.04451 from iclr.

[nl]

I agree about it being early days.

Reformer is a good example that I'd missed.

Neural Tangents is another paper demoing a framework.