Porting to equinox (Jax)

[neel04]

I was working on an equinox implementation here: https://colab.research.google.com/drive/1pIki9a6CG5Ch8dH94vbcpCC6-tE8Fbfj?usp=sharing

However, I was getting nan loss after a couple epochs and decided to leave this for another time.

If anyone wants to contribute to triangulate the issue, you're more than welcome! 🤗

(Eventually, I'll probably end up making a seperate Jax version of this repo)

[neel04]

A lot of the components I mentioned in the writeup aren't implemented here, but they shouldn't be too hard. For simpler tasks, like reverse_string one could even forgo adversarial_perturbation.

Check the WandB commit for more information