adding orbax.checkpoint version to requirements.txt

[lihenglin]

so I just found out using latest orbax.checkpoint package (0.4.8) would result in the following error
AttributeError: module 'jax.experimental.array_serialization.serialization' has no attribute '_DEFAULT_DRIVER'

But this can be fixed by downgrading to 0.3.5.
Probably the authors would like to add a version number for this package to requirements.txt?