grain should not have jax dependency

[SamuelMarks]

@ 77c56dc with Python (CPython) 3.10.16, setuptools 75.6.0, wheel 0.45.1 on x86_64 Linux:

$ python -m pip install .
$ python -c 'import grain.python as grain'
Traceback (most recent call last):
  File "<string>", line 1, in <module>
  File "grain/grain/python.py", line 21, in <module>
    from . import python_experimental as experimental
  File "grain/grain/python_experimental.py", line 48, in <module>
    from ._src.python.dataset.transformations.packing import FirstFitPackIterDataset
  File "grain/grain/_src/python/dataset/transformations/packing.py", line 21, in <module>
    from grain._src.python.dataset.transformations import packing_packed_batch
  File "grain/grain/_src/python/dataset/transformations/packing_packed_batch.py", line 23, in <module>
    import jax
ModuleNotFoundError: No module named 'jax'

But at that same commit your README reads:
https://github.com/google/grain/blob/77c56dc/README.md?#L30

Grain is designed to work with JAX models but it does not require JAX to run and can be used with other frameworks as well.

[iindyk]

thanks for flagging this, let me prepare a pr with a fix