Revert "Work-around for protobuf init crash in sentencepiece" (#9697)

Reverts #9693

Author: jeffhataws

.github/workflows/_test.yml

@@ -107,7 +107,6 @@ jobs:
                   pip install fsspec
                   pip install rich
                   pip install flax
-                  pip install sentencepiece
             - name: Extra CI deps
               if: inputs.has_code_changes == 'true'
               shell: bash

test/run_tests.sh

@@ -266,10 +266,6 @@ function run_xla_op_tests3 {
 
 function run_xla_op_tests4 {
   run_test "$_TEST_DIR/test_jax_interop.py"
-  # issue #9691: random crashes with sentencepiece protobuf; run multiple times to trigger
-  for i in $(seq 1 5); do
-    run_test "$_TEST_DIR/test_sentencepiece_interop.py"
-  done
 }
 
 function run_xla_op_tests5 {

test/test_sentencepiece_interop.py

@@ -7,15 +7,6 @@
 
 import torch
 
-# issue 9691: ensure sentencepiece protobuf init happen between
-# torch/torch-xla protobuf inits to work-around protobuf crash
-try:
-  import sentencepiece as spm
-  sp_model = spm.SentencePieceProcessor()
-  sp_model.load('')
-except:
-  pass
-
 import _XLAC
 from ._internal import tpu
 from .version import __version__

torch_xla/__init__.py

@@ -2,5 +2,4 @@
 torch==2.8.0 ; sys_platform == 'darwin'  # macOS
 torch==2.8.0+cpu; sys_platform != 'darwin' # Non-macOS (CPU-only), like on TPU
 yapf==0.40.2  # N.B.: keep in sync with `infra/ansible/config/pip.yaml`, `.github/workflows/lintercheck.yml`
-jax==0.7.2  # N.B.: torchax breaks on newer JAX versions that would be pulled from `flax` dependencies
 flax==0.10.6