We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Seg Fault after saving checkpoints
Hi, 

I am getting a seg fault sometimes after the model has saved the checkpoint. It is not every checkpoint and seems to be random which checkpoints it crashes after. I am not sure if it is related to issue #340

For example, I am running `prompt_tuning/scripts/sst2-demo-xxl.sh`, and the output is below. 

```
317 18:14:56.525280 140415323761728 utils.py:138] Saved Numpy Arrays for step 1104000 to gs://nicl/checkpoint_models/sst/full_dataset/prompt-tuning/t5-11b/numpy_checkpoints/checkpoint_1104000
I0317 18:14:56.604028 140415323761728 checkpoints.py:600] Saving checkpoint for step 1104000 to gs://nicl/checkpoint_models/sst/full_dataset/prompt-tuning/t5-11b/checkpoint_1104000.tmp-1647540896
I0317 18:14:56.614308 140622481194048 checkpoints.py:600] Saving checkpoint for step 1104000 to gs://nicl/checkpoint_models/sst/full_dataset/prompt-tuning/t5-11b/checkpoint_1104000.tmp-1647540896
I0317 18:14:56.624289 140590966570048 checkpoints.py:600] Saving checkpoint for step 1104000 to gs://nicl/checkpoint_models/sst/full_dataset/prompt-tuning/t5-11b/checkpoint_1104000.tmp-1647540896
I0317 18:14:56.653718 140272509271104 checkpoints.py:600] Saving checkpoint for step 1104000 to gs://nicl/checkpoint_models/sst/full_dataset/prompt-tuning/t5-11b/checkpoint_1104000.tmp-1647540896
Fatal Python error: Segmentation fault


Thread 0x00007fdb1dc01700 (most recent call first):
  File "/home/dptam/.local/lib/python3.8/site-packages/jax/interpreters/pxla.py", line 664 in _sda_value
  File "/home/dptam/.local/lib/python3.8/site-packages/jax/_src/device_array.py", line 266 in __array__
  File "/home/dptam/.local/lib/python3.8/site-packages/t5x/checkpoints.py", line 447 in <lambda>
  File "/home/dptam/.local/lib/python3.8/site-packages/t5x/checkpoint_importer.py", line 84 in get
  File "/usr/lib/python3.8/concurrent/futures/thread.py", line 57 in run
  File "/usr/lib/python3.8/concurrent/futures/thread.py", line 80 in _worker
  File "/usr/lib/python3.8/threading.py", line 870 in run
  File "/usr/lib/python3.8/threading.py", line 932 in _bootstrap_inner
  File "/usr/lib/python3.8/threading.py", line 890 in _bootstrap

Thread 0x00007f56809df700 (most recent call first):
  File "/usr/lib/python3.8/concurrent/futures/thread.py", line 78 in _worker
  File "/usr/lib/python3.8/threading.py", line 870 in run
  File "/usr/lib/python3.8/threading.py", line 932 in _bootstrap_inner
  File "/usr/lib/python3.8/threading.py", line 890 in _bootstrap

  Thread 0x00007f56c7aad700 (most recent call first):
  File "/usr/lib/python3.8/concurrent/futures/thread.py", line 78 in _worker
  File "/usr/lib/python3.8/threading.py", line 870 in run
  File "/usr/lib/python3.8/threading.py", line 932 in _bootstrap_inner
  File "/usr/lib/python3.8/threading.py", line 890 in _bootstrap
Thread 0x00007fdde29efc40 (most recent call first):
  File "/home/dptam/.local/lib/python3.8/site-packages/t5x/checkpoints.py", line 693 in _write_array
https://symbolize.stripped_domain/r/?trace=7fdde2e4203b,7fdde2e420bf,e,5ef27540f,e,26f7c5aff,f,b15f59df&map= 
E0317 18:14:57.770066  341059 process_state.cc:1062] RAW: Signal 11 raised at PC: 0x7fdde2e4203b while already in FailureSignalHandler!
E0317 18:14:57.770096  341059 process_state.cc:1065] RAW: tid: 341059 raised new signal
    @                0xf       1440  (unknown)
    @        0x25ed159b0  (unknown)  (unknown)
    @               0x10   76231216  (unknown)
    @        0x261cdc840  (unknown)  (unknown)
    @        0x2dfdd4780  (unknown)  (unknown)
    @        0x5f1f8a120  (unknown)  (unknown)
https://symbolize.stripped_domain/r/?trace=7fdde301ffd3,7fddd98d57f9,7fdde2e420bf,7,e,25ed159af,f,261cdc83f,2dfdd477f,5f1f8a11f&map=7a511a57244151c993b16b37978e7ed7:7fddcaefd000-7fddd9c3fd50 
E0317 18:14:57.818885  341068 coredump_hook.cc:365] RAW: Remote crash data gathering hook invoked.
E0317 18:14:57.818900  341068 coredump_hook.cc:411] RAW: Skipping coredump since rlimit was 0 at process start.
E0317 18:14:57.818919  341068 client.cc:221] RAW: Coroner client retries enabled (b/136286901), will retry for up to 30 sec.
E0317 18:14:57.818922  341068 coredump_hook.cc:473] RAW: Sending fingerprint to remote end.
E0317 18:14:57.818928  341068 coredump_socket.cc:124] RAW: Stat failed errno=2 on socket /var/google/services/logmanagerd/remote_coredump.socket
E0317 18:14:57.818933  341068 coredump_hook.cc:477] RAW: Cannot send fingerprint to Coroner: [NOT_FOUND] Missing crash reporting socket. Is the listener running?
E0317 18:14:57.818938  341068 coredump_hook.cc:550] RAW: Discarding core.
prompt_tuning/scripts/sst2-demo-xxl.sh: line 37: 337643 Segmentation fault      (core dumped) python3 -m t5x.train --gin_search_paths="${T5X_DIR},${FLAXFORMER_DIR},${PROMPT_DIR}" --gin_file="prompt_tuning/configs/models/t5_1_1_xxl_prompt.gin" --gin_file="prompt_tuning/configs/prompts/from_class_labels.gin" --gin_file="prompt_tuning/configs/runs/prompt_finetune.gin" --gin.CLASS_LABELS="['positive', 'negative']" --gin.MODEL_DIR="'${MODEL_DIR}'" --gin.MIXTURE_OR_TASK_NAME="'taskless_glue_sst2_v200_examples'" --gin.MIXTURE_OR_TASK_MODULE="'prompt_tuning.data.glue'" --gin.TASK_FEATURE_LENGTHS="{'inputs': 512, 'targets': 8}" --gin.INITIAL_CHECKPOINT_PATH="'${PRETRAINED_MODEL}'" --gin.TRAIN_STEPS="1_212_000" --gin.USE_CACHED_TASKS="False" --gin.BATCH_SIZE="16" --gin.partitioning.PjitPartitioner.model_parallel_submesh="(4, 4, 1, 2)" --tfds_data_dir=${TFDS_DATA_DIR}
##### Command execution on worker 3 failed with return code 139. Continuing.
prompt_tuning/scripts/sst2-demo-xxl.sh: line 37: 334750 Aborted                 (core dumped) python3 -m t5x.train --gin_search_paths="${T5X_DIR},${FLAXFORMER_DIR},${PROMPT_DIR}" --gin_file="prompt_tuning/configs/models/t5_1_1_xxl_prompt.gin" --gin_file="prompt_tuning/configs/prompts/from_class_labels.gin" --gin_file="prompt_tuning/configs/runs/prompt_finetune.gin" --gin.CLASS_LABELS="['positive', 'negative']" --gin.MODEL_DIR="'${MODEL_DIR}'" --gin.MIXTURE_OR_TASK_NAME="'taskless_glue_sst2_v200_examples'" --gin.MIXTURE_OR_TASK_MODULE="'prompt_tuning.data.glue'" --gin.TASK_FEATURE_LENGTHS="{'inputs': 512, 'targets': 8}" --gin.INITIAL_CHECKPOINT_PATH="'${PRETRAINED_MODEL}'" --gin.TRAIN_STEPS="1_212_000" --gin.USE_CACHED_TASKS="False" --gin.BATCH_SIZE="16" --gin.partitioning.PjitPartitioner.model_parallel_submesh="(4, 4, 1, 2)" --tfds_data_dir=${TFDS_DATA_DIR}
##### Command execution on worker 1 failed with return code 134. Continuing.
prompt_tuning/scripts/sst2-demo-xxl.sh: line 37: 335504 Aborted                 (core dumped) python3 -m t5x.train --gin_search_paths="${T5X_DIR},${FLAXFORMER_DIR},${PROMPT_DIR}" --gin_file="prompt_tuning/configs/models/t5_1_1_xxl_prompt.gin" --gin_file="prompt_tuning/configs/prompts/from_class_labels.gin" --gin_file="prompt_tuning/configs/runs/prompt_finetune.gin" --gin.CLASS_LABELS="['positive', 'negative']" --gin.MODEL_DIR="'${MODEL_DIR}'" --gin.MIXTURE_OR_TASK_NAME="'taskless_glue_sst2_v200_examples'" --gin.MIXTURE_OR_TASK_MODULE="'prompt_tuning.data.glue'" --gin.TASK_FEATURE_LENGTHS="{'inputs': 512, 'targets': 8}" --gin.INITIAL_CHECKPOINT_PATH="'${PRETRAINED_MODEL}'" --gin.TRAIN_STEPS="1_212_000" --gin.USE_CACHED_TASKS="False" --gin.BATCH_SIZE="16" --gin.partitioning.PjitPartitioner.model_parallel_submesh="(4, 4, 1, 2)" --tfds_data_dir=${TFDS_DATA_DIR}
##### Command execution on worker 0 failed with return code 134. Continuing.
```

Thanks

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `setup.py`
Content:
```
1 # Copyright 2022 The T5X Authors.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #     http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 
15 """Install T5X."""
16 
17 import os
18 import sys
19 import setuptools
20 
21 # To enable importing version.py directly, we add its path to sys.path.
22 version_path = os.path.join(os.path.dirname(__file__), 't5x')
23 sys.path.append(version_path)
24 from version import __version__  # pylint: disable=g-import-not-at-top
25 
26 # Get the long description from the README file.
27 with open('README.md') as fp:
28   _LONG_DESCRIPTION = fp.read()
29 
30 _jax_version = '0.2.27'
31 _jaxlib_version = '0.1.76'
32 
33 setuptools.setup(
34     name='t5x',
35     version=__version__,
36     description='T5-eXtended in JAX',
37     long_description=_LONG_DESCRIPTION,
38     long_description_content_type='text/markdown',
39     author='Google Inc.',
40     author_email='no-reply@google.com',
41     url='http://github.com/google-research/t5x',
42     license='Apache 2.0',
43     packages=setuptools.find_packages(),
44     package_data={
45         '': ['**/*.gin'],  # not all subdirectories may have __init__.py.
46     },
47     scripts=[],
48     install_requires=[
49         'absl-py',
50         'cached_property',
51         # TODO(adarob): Replace with 'clu' once >0.0.6 is released.
52         'clu @ git+https://github.com/google/CommonLoopUtils#egg=clu',
53         'flax @ git+https://github.com/google/flax#egg=flax',
54         'gin-config',
55         f'jax >= {_jax_version}',
56         f'jaxlib >= {_jaxlib_version}',
57         'numpy',
58         'seqio-nightly',
59         't5',
60         'tensorflow',
61         'tensorstore',
62     ],
63     extras_require={
64         'gcp': [
65             'gevent', 'google-api-python-client', 'google-compute-engine',
66             'google-cloud-storage', 'oauth2client'
67         ],
68         'test': ['pytest'],
69 
70         # Cloud TPU requirements.
71         'tpu': [f'jax[tpu] >= {_jax_version}'],
72     },
73     classifiers=[
74         'Development Status :: 4 - Beta',
75         'Intended Audience :: Developers',
76         'Intended Audience :: Science/Research',
77         'License :: OSI Approved :: Apache Software License',
78         'Topic :: Scientific/Engineering :: Artificial Intelligence',
79     ],
80     keywords='text nlp machinelearning',
81 )
82
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -58,7 +58,7 @@
         'seqio-nightly',
         't5',
         'tensorflow',
-        'tensorstore',
+        'tensorstore >= 0.1.20',
     ],
     extras_require={
         'gcp': [

{"golden_diff": "diff --git a/setup.py b/setup.py\n--- a/setup.py\n+++ b/setup.py\n@@ -58,7 +58,7 @@\n         'seqio-nightly',\n         't5',\n         'tensorflow',\n-        'tensorstore',\n+        'tensorstore >= 0.1.20',\n     ],\n     extras_require={\n         'gcp': [\n", "issue": "Seg Fault after saving checkpoints\nHi, \r\n\r\nI am getting a seg fault sometimes after the model has saved the checkpoint. It is not every checkpoint and seems to be random which checkpoints it crashes after. I am not sure if it is related to issue #340\r\n\r\nFor example, I am running `prompt_tuning/scripts/sst2-demo-xxl.sh`, and the output is below. \r\n\r\n```\r\n317 18:14:56.525280 140415323761728 utils.py:138] Saved Numpy Arrays for step 1104000 to gs://nicl/checkpoint_models/sst/full_dataset/prompt-tuning/t5-11b/numpy_checkpoints/checkpoint_1104000\r\nI0317 18:14:56.604028 140415323761728 checkpoints.py:600] Saving checkpoint for step 1104000 to gs://nicl/checkpoint_models/sst/full_dataset/prompt-tuning/t5-11b/checkpoint_1104000.tmp-1647540896\r\nI0317 18:14:56.614308 140622481194048 checkpoints.py:600] Saving checkpoint for step 1104000 to gs://nicl/checkpoint_models/sst/full_dataset/prompt-tuning/t5-11b/checkpoint_1104000.tmp-1647540896\r\nI0317 18:14:56.624289 140590966570048 checkpoints.py:600] Saving checkpoint for step 1104000 to gs://nicl/checkpoint_models/sst/full_dataset/prompt-tuning/t5-11b/checkpoint_1104000.tmp-1647540896\r\nI0317 18:14:56.653718 140272509271104 checkpoints.py:600] Saving checkpoint for step 1104000 to gs://nicl/checkpoint_models/sst/full_dataset/prompt-tuning/t5-11b/checkpoint_1104000.tmp-1647540896\r\nFatal Python error: Segmentation fault\r\n\r\n\r\nThread 0x00007fdb1dc01700 (most recent call first):\r\n  File \"/home/dptam/.local/lib/python3.8/site-packages/jax/interpreters/pxla.py\", line 664 in _sda_value\r\n  File \"/home/dptam/.local/lib/python3.8/site-packages/jax/_src/device_array.py\", line 266 in __array__\r\n  File \"/home/dptam/.local/lib/python3.8/site-packages/t5x/checkpoints.py\", line 447 in <lambda>\r\n  File \"/home/dptam/.local/lib/python3.8/site-packages/t5x/checkpoint_importer.py\", line 84 in get\r\n  File \"/usr/lib/python3.8/concurrent/futures/thread.py\", line 57 in run\r\n  File \"/usr/lib/python3.8/concurrent/futures/thread.py\", line 80 in _worker\r\n  File \"/usr/lib/python3.8/threading.py\", line 870 in run\r\n  File \"/usr/lib/python3.8/threading.py\", line 932 in _bootstrap_inner\r\n  File \"/usr/lib/python3.8/threading.py\", line 890 in _bootstrap\r\n\r\nThread 0x00007f56809df700 (most recent call first):\r\n  File \"/usr/lib/python3.8/concurrent/futures/thread.py\", line 78 in _worker\r\n  File \"/usr/lib/python3.8/threading.py\", line 870 in run\r\n  File \"/usr/lib/python3.8/threading.py\", line 932 in _bootstrap_inner\r\n  File \"/usr/lib/python3.8/threading.py\", line 890 in _bootstrap\r\n\r\n  Thread 0x00007f56c7aad700 (most recent call first):\r\n  File \"/usr/lib/python3.8/concurrent/futures/thread.py\", line 78 in _worker\r\n  File \"/usr/lib/python3.8/threading.py\", line 870 in run\r\n  File \"/usr/lib/python3.8/threading.py\", line 932 in _bootstrap_inner\r\n  File \"/usr/lib/python3.8/threading.py\", line 890 in _bootstrap\r\nThread 0x00007fdde29efc40 (most recent call first):\r\n  File \"/home/dptam/.local/lib/python3.8/site-packages/t5x/checkpoints.py\", line 693 in _write_array\r\nhttps://symbolize.stripped_domain/r/?trace=7fdde2e4203b,7fdde2e420bf,e,5ef27540f,e,26f7c5aff,f,b15f59df&map= \r\nE0317 18:14:57.770066  341059 process_state.cc:1062] RAW: Signal 11 raised at PC: 0x7fdde2e4203b while already in FailureSignalHandler!\r\nE0317 18:14:57.770096  341059 process_state.cc:1065] RAW: tid: 341059 raised new signal\r\n    @                0xf       1440  (unknown)\r\n    @        0x25ed159b0  (unknown)  (unknown)\r\n    @               0x10   76231216  (unknown)\r\n    @        0x261cdc840  (unknown)  (unknown)\r\n    @        0x2dfdd4780  (unknown)  (unknown)\r\n    @        0x5f1f8a120  (unknown)  (unknown)\r\nhttps://symbolize.stripped_domain/r/?trace=7fdde301ffd3,7fddd98d57f9,7fdde2e420bf,7,e,25ed159af,f,261cdc83f,2dfdd477f,5f1f8a11f&map=7a511a57244151c993b16b37978e7ed7:7fddcaefd000-7fddd9c3fd50 \r\nE0317 18:14:57.818885  341068 coredump_hook.cc:365] RAW: Remote crash data gathering hook invoked.\r\nE0317 18:14:57.818900  341068 coredump_hook.cc:411] RAW: Skipping coredump since rlimit was 0 at process start.\r\nE0317 18:14:57.818919  341068 client.cc:221] RAW: Coroner client retries enabled (b/136286901), will retry for up to 30 sec.\r\nE0317 18:14:57.818922  341068 coredump_hook.cc:473] RAW: Sending fingerprint to remote end.\r\nE0317 18:14:57.818928  341068 coredump_socket.cc:124] RAW: Stat failed errno=2 on socket /var/google/services/logmanagerd/remote_coredump.socket\r\nE0317 18:14:57.818933  341068 coredump_hook.cc:477] RAW: Cannot send fingerprint to Coroner: [NOT_FOUND] Missing crash reporting socket. Is the listener running?\r\nE0317 18:14:57.818938  341068 coredump_hook.cc:550] RAW: Discarding core.\r\nprompt_tuning/scripts/sst2-demo-xxl.sh: line 37: 337643 Segmentation fault      (core dumped) python3 -m t5x.train --gin_search_paths=\"${T5X_DIR},${FLAXFORMER_DIR},${PROMPT_DIR}\" --gin_file=\"prompt_tuning/configs/models/t5_1_1_xxl_prompt.gin\" --gin_file=\"prompt_tuning/configs/prompts/from_class_labels.gin\" --gin_file=\"prompt_tuning/configs/runs/prompt_finetune.gin\" --gin.CLASS_LABELS=\"['positive', 'negative']\" --gin.MODEL_DIR=\"'${MODEL_DIR}'\" --gin.MIXTURE_OR_TASK_NAME=\"'taskless_glue_sst2_v200_examples'\" --gin.MIXTURE_OR_TASK_MODULE=\"'prompt_tuning.data.glue'\" --gin.TASK_FEATURE_LENGTHS=\"{'inputs': 512, 'targets': 8}\" --gin.INITIAL_CHECKPOINT_PATH=\"'${PRETRAINED_MODEL}'\" --gin.TRAIN_STEPS=\"1_212_000\" --gin.USE_CACHED_TASKS=\"False\" --gin.BATCH_SIZE=\"16\" --gin.partitioning.PjitPartitioner.model_parallel_submesh=\"(4, 4, 1, 2)\" --tfds_data_dir=${TFDS_DATA_DIR}\r\n##### Command execution on worker 3 failed with return code 139. Continuing.\r\nprompt_tuning/scripts/sst2-demo-xxl.sh: line 37: 334750 Aborted                 (core dumped) python3 -m t5x.train --gin_search_paths=\"${T5X_DIR},${FLAXFORMER_DIR},${PROMPT_DIR}\" --gin_file=\"prompt_tuning/configs/models/t5_1_1_xxl_prompt.gin\" --gin_file=\"prompt_tuning/configs/prompts/from_class_labels.gin\" --gin_file=\"prompt_tuning/configs/runs/prompt_finetune.gin\" --gin.CLASS_LABELS=\"['positive', 'negative']\" --gin.MODEL_DIR=\"'${MODEL_DIR}'\" --gin.MIXTURE_OR_TASK_NAME=\"'taskless_glue_sst2_v200_examples'\" --gin.MIXTURE_OR_TASK_MODULE=\"'prompt_tuning.data.glue'\" --gin.TASK_FEATURE_LENGTHS=\"{'inputs': 512, 'targets': 8}\" --gin.INITIAL_CHECKPOINT_PATH=\"'${PRETRAINED_MODEL}'\" --gin.TRAIN_STEPS=\"1_212_000\" --gin.USE_CACHED_TASKS=\"False\" --gin.BATCH_SIZE=\"16\" --gin.partitioning.PjitPartitioner.model_parallel_submesh=\"(4, 4, 1, 2)\" --tfds_data_dir=${TFDS_DATA_DIR}\r\n##### Command execution on worker 1 failed with return code 134. Continuing.\r\nprompt_tuning/scripts/sst2-demo-xxl.sh: line 37: 335504 Aborted                 (core dumped) python3 -m t5x.train --gin_search_paths=\"${T5X_DIR},${FLAXFORMER_DIR},${PROMPT_DIR}\" --gin_file=\"prompt_tuning/configs/models/t5_1_1_xxl_prompt.gin\" --gin_file=\"prompt_tuning/configs/prompts/from_class_labels.gin\" --gin_file=\"prompt_tuning/configs/runs/prompt_finetune.gin\" --gin.CLASS_LABELS=\"['positive', 'negative']\" --gin.MODEL_DIR=\"'${MODEL_DIR}'\" --gin.MIXTURE_OR_TASK_NAME=\"'taskless_glue_sst2_v200_examples'\" --gin.MIXTURE_OR_TASK_MODULE=\"'prompt_tuning.data.glue'\" --gin.TASK_FEATURE_LENGTHS=\"{'inputs': 512, 'targets': 8}\" --gin.INITIAL_CHECKPOINT_PATH=\"'${PRETRAINED_MODEL}'\" --gin.TRAIN_STEPS=\"1_212_000\" --gin.USE_CACHED_TASKS=\"False\" --gin.BATCH_SIZE=\"16\" --gin.partitioning.PjitPartitioner.model_parallel_submesh=\"(4, 4, 1, 2)\" --tfds_data_dir=${TFDS_DATA_DIR}\r\n##### Command execution on worker 0 failed with return code 134. Continuing.\r\n```\r\n\r\nThanks\n", "before_files": [{"content": "# Copyright 2022 The T5X Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Install T5X.\"\"\"\n\nimport os\nimport sys\nimport setuptools\n\n# To enable importing version.py directly, we add its path to sys.path.\nversion_path = os.path.join(os.path.dirname(__file__), 't5x')\nsys.path.append(version_path)\nfrom version import __version__  # pylint: disable=g-import-not-at-top\n\n# Get the long description from the README file.\nwith open('README.md') as fp:\n  _LONG_DESCRIPTION = fp.read()\n\n_jax_version = '0.2.27'\n_jaxlib_version = '0.1.76'\n\nsetuptools.setup(\n    name='t5x',\n    version=__version__,\n    description='T5-eXtended in JAX',\n    long_description=_LONG_DESCRIPTION,\n    long_description_content_type='text/markdown',\n    author='Google Inc.',\n    author_email='no-reply@google.com',\n    url='http://github.com/google-research/t5x',\n    license='Apache 2.0',\n    packages=setuptools.find_packages(),\n    package_data={\n        '': ['**/*.gin'],  # not all subdirectories may have __init__.py.\n    },\n    scripts=[],\n    install_requires=[\n        'absl-py',\n        'cached_property',\n        # TODO(adarob): Replace with 'clu' once >0.0.6 is released.\n        'clu @ git+https://github.com/google/CommonLoopUtils#egg=clu',\n        'flax @ git+https://github.com/google/flax#egg=flax',\n        'gin-config',\n        f'jax >= {_jax_version}',\n        f'jaxlib >= {_jaxlib_version}',\n        'numpy',\n        'seqio-nightly',\n        't5',\n        'tensorflow',\n        'tensorstore',\n    ],\n    extras_require={\n        'gcp': [\n            'gevent', 'google-api-python-client', 'google-compute-engine',\n            'google-cloud-storage', 'oauth2client'\n        ],\n        'test': ['pytest'],\n\n        # Cloud TPU requirements.\n        'tpu': [f'jax[tpu] >= {_jax_version}'],\n    },\n    classifiers=[\n        'Development Status :: 4 - Beta',\n        'Intended Audience :: Developers',\n        'Intended Audience :: Science/Research',\n        'License :: OSI Approved :: Apache Software License',\n        'Topic :: Scientific/Engineering :: Artificial Intelligence',\n    ],\n    keywords='text nlp machinelearning',\n)\n", "path": "setup.py"}], "after_files": [{"content": "# Copyright 2022 The T5X Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Install T5X.\"\"\"\n\nimport os\nimport sys\nimport setuptools\n\n# To enable importing version.py directly, we add its path to sys.path.\nversion_path = os.path.join(os.path.dirname(__file__), 't5x')\nsys.path.append(version_path)\nfrom version import __version__  # pylint: disable=g-import-not-at-top\n\n# Get the long description from the README file.\nwith open('README.md') as fp:\n  _LONG_DESCRIPTION = fp.read()\n\n_jax_version = '0.2.27'\n_jaxlib_version = '0.1.76'\n\nsetuptools.setup(\n    name='t5x',\n    version=__version__,\n    description='T5-eXtended in JAX',\n    long_description=_LONG_DESCRIPTION,\n    long_description_content_type='text/markdown',\n    author='Google Inc.',\n    author_email='no-reply@google.com',\n    url='http://github.com/google-research/t5x',\n    license='Apache 2.0',\n    packages=setuptools.find_packages(),\n    package_data={\n        '': ['**/*.gin'],  # not all subdirectories may have __init__.py.\n    },\n    scripts=[],\n    install_requires=[\n        'absl-py',\n        'cached_property',\n        # TODO(adarob): Replace with 'clu' once >0.0.6 is released.\n        'clu @ git+https://github.com/google/CommonLoopUtils#egg=clu',\n        'flax @ git+https://github.com/google/flax#egg=flax',\n        'gin-config',\n        f'jax >= {_jax_version}',\n        f'jaxlib >= {_jaxlib_version}',\n        'numpy',\n        'seqio-nightly',\n        't5',\n        'tensorflow',\n        'tensorstore >= 0.1.20',\n    ],\n    extras_require={\n        'gcp': [\n            'gevent', 'google-api-python-client', 'google-compute-engine',\n            'google-cloud-storage', 'oauth2client'\n        ],\n        'test': ['pytest'],\n\n        # Cloud TPU requirements.\n        'tpu': [f'jax[tpu] >= {_jax_version}'],\n    },\n    classifiers=[\n        'Development Status :: 4 - Beta',\n        'Intended Audience :: Developers',\n        'Intended Audience :: Science/Research',\n        'License :: OSI Approved :: Apache Software License',\n        'Topic :: Scientific/Engineering :: Artificial Intelligence',\n    ],\n    keywords='text nlp machinelearning',\n)\n", "path": "setup.py"}]}