Move protobuf deps to be optional.

PiperOrigin-RevId: 466123822

setup.py

@@ -70,7 +70,6 @@ def generate_proto(source):
         'scipy>=1.5',
         'typing_extensions',
         'etils[epath]',
-        'protobuf>=3.13,<4',
     ],
     extras_require={
         # Minimum jaxlib version; used in testing.
@@ -90,6 +89,9 @@ def generate_proto(source):
                 # Required by cloud_tpu_init.py
                 'requests'],
 
+        # $ pip install jax[australis]
+        'australis': ['protobuf>=3.13,<4'],
+
         # CUDA installations require adding jax releases URL; e.g.
         # Cuda installation defaulting to a CUDA and Cudnn version defined above.
         # $ pip install jax[cuda] -f https://storage.googleapis.com/jax-releases/jax_releases.html