Remove flax dependency

[smorad]

Since October, flax has failed to support the latest python. I use flashbax in my code and it is currently breaking my builds due this flax issue.

This PR replaces the only use of flax (dataclasses) with a simple alternative with almost no dependencies.

[sash-a]

Thanks for this, that is reasonable since we only use it in one place, however seeing as in other places we use chex.dataclass. I'd rather use that instead of adding an extra dependency. Is there a way to do static types/non-pytree nodes with chex?

I'm also a bit weary of jax_dataclasses as it hasn't had much activity recently, so flax is better since I'm sure it will be maintained for a while.

[smorad]

@sash-a I also wanted to use chex but it doesn't look like they support the pytree_node argument. The flax struct is fairly modular and can be extracted from the flax library. I've added this as a new file.

[smorad]

I am pretty sure the one test failure is unrelated to this commit.

flashbax/vault/__init__.py:14: in <module>
    from flashbax.vault.vault import Vault
flashbax/vault/vault.py:27: in <module>
    from etils import epath  # type: ignore
E   ModuleNotFoundError: No module named 'etils'

All tests passed on my machine.

[sash-a]

Thanks for this, I'm a bit worried about potential extra maintenance from hving this in flashbax. Do you know how long until they upgrade to the latest python?

[smorad]

@sash-a I have no idea, the issue has been open for two months and no update for one month.

[sash-a]

Can you add the following to your pyproject.toml, it should fix the dependency issue 😄

    'orbax-checkpoint==0.11.20',
    'etils==1.13.0',
    'importlib-resources==6.5.2',

[smorad]

Seems the orbax-checkpoint dep no longer supports python 3.9 used by the GitHub action. I tested locally and can install everything using python 3.10. I will update the runner to use 3.10 and see if that works.