search

neel04 / ReAct

Exploration with adding recurrent priors to attention based models
3 stars 0 forks source link

Porting to equinox (Jax) #8

Open neel04 opened 1 year ago

neel04 commented 1 year ago

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 commented 1 year ago

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