Dependency issue of 0.0.9

[wdhongtw]

v0.0.9 has a direct dependency on flax,

https://github.com/google/torchax/blob/main/torchax/checkpoint.py#L18

but does not declare flax as a dependency in Python package level. (not installed indirectly if install torchax)

Should consider adding this dependency or update README doc for how to install and use torchax properly.

[wdhongtw]

Mitigated in #71.

But it's probably good to include an optional dependency for flax.