Skip to content

Porting to equinox (Jax) #8

@neel04

Description

@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)

Metadata

Metadata

Assignees

No one assigned

    Labels

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions