Work-around for protobuf init crash in sentencepiece (#9693)

A work-around for the protobuf init crash in sentencepiece. The code
within the try/except block triggers sentencepiece loading path. This
has to be done between ``import torch`` and ``import _XLAC`` and works
even though the loading of "" triggers file not found error.

#9691

Author: jeffhataws

.github/workflows/_test.yml

@@ -107,6 +107,7 @@ 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,6 +266,10 @@ 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

@@ -0,0 +1,19 @@
+import unittest
+
+
+class SentencepieceInterop(unittest.TestCase):
+
+  def test_sentencepiece_interop(self):
+    import os
+    if not os.path.exists("/tmp/test_model.model"):
+      import urllib.request
+      urllib.request.urlretrieve(
+          "https://github.com/google/sentencepiece/raw/refs/heads/master/python/test/test_model.model",
+          "/tmp/test_model.model")
+    import torch_xla
+    import sentencepiece as spm
+    sp_model = spm.SentencePieceProcessor("/tmp/test_model.model")
+
+
+if __name__ == "__main__":
+  unittest.main()

torch_xla/__init__.py

@@ -7,6 +7,15 @@
 
 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__

torchax/dev-requirements.txt

@@ -2,4 +2,5 @@
 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