Request for new PyPI release

[yhtang]

JEstimator 0.3.3 on PyPI is incompatible with the latest JAX due to the relocation of the PartitionSpec endpoint:

  File "/opt/t5x/t5x/optimizers.py", line 38, in <module>
    from jestimator import amos
  File "/usr/local/lib/python3.10/dist-packages/jestimator/__init__.py", line 20, in <module>
    from . import amos_helper
  File "/usr/local/lib/python3.10/dist-packages/jestimator/amos_helper.py", line 24, in <module>
    from jax.experimental.pjit import PartitionSpec
ImportError: cannot import name 'PartitionSpec' from 'jax.experimental.pjit' (/opt/jax-source/jax/experimental/pjit.py)

This is causing issues in T5X.

While a workaround can be made in T5X to install jestimator from the development head, it may be better to publish a 'stable' version on PyPI so that depending projects can more easily refer to.