Replace flax usage with equinox
Open
enhancement
- Dominant language
- Python
- Stars
- 43
- Forks
- 6
- Avg merge
- 3d 22h
- Merged PRs (30d)
- 10
Description
### What's the new feature?
In `coreax.networks.py` we have the class `ScoreNetwork` which uses `flax.linen.Module`, as-well as the use of `flax.training.train_state.TrainState` in the same file. We could instead use Equinox functionality for the same purpose, and remove the need to use flax at all. Note that Equinox is used elsewhere in the codebase anyway.
### What value does this add?
Reliance on fewer third party packages.
### Is there an alternative you've considered?
_No response_
### Additional context
_No response_
Contributor guide
Assessment
This issue has not been assessed yet.