Support python 3.11

PiperOrigin-RevId: 647581034

.github/workflows/ci.yml

@@ -15,7 +15,7 @@ jobs:
 
     strategy:
       matrix:
-        python-version: ["3.10"]
+        python-version: ["3.10", "3.11"]
         os: [ubuntu-latest]
 
     steps:

requirements_tests.txt

@@ -1,4 +1,4 @@
 chex==0.1.85
 dm-haiku==0.0.11
-tensorflow==2.11.1; sys_platform != 'darwin' or platform_machine != 'arm64'
-tensorflow-macos==2.11.1; sys_platform == 'darwin' and platform_machine == 'arm64'
+tensorflow==2.16.1; sys_platform != 'darwin' or platform_machine != 'arm64'
+tensorflow-macos==2.16.1; sys_platform == 'darwin' and platform_machine == 'arm64'

setup.py

@@ -67,6 +67,7 @@ def _parse_requirements(path):
         'License :: OSI Approved :: Apache Software License',
         'Operating System :: OS Independent',
         'Programming Language :: Python :: 3.10',
+        'Programming Language :: Python :: 3.11',
         'Topic :: Scientific/Engineering :: Artificial Intelligence',
         'Topic :: Scientific/Engineering :: Mathematics',
         'Topic :: Software Development :: Libraries :: Python Modules',

test.sh

@@ -26,7 +26,7 @@ python --version
 # Install dependencies.
 pip install --upgrade pip setuptools wheel
 # See https://github.com/google/pytype/issues/1316
-pip install flake8 pytest-xdist pytype==2022.9.19 importlab==0.7 pylint pylint-exit
+pip install flake8 pytest-xdist pytype importlab pylint pylint-exit
 pip install -r requirements.txt
 pip install -r requirements_tests.txt
 

tf2jax/_src/roundtrip_test.py

@@ -69,6 +69,17 @@ def uses_native_serialization():
 
 class Jax2TfTest(test_util.TestCase):
 
+  def setUp(self):
+    super().setUp()
+    if not uses_native_serialization():
+      self._xla_op = tf2jax.ops._jax_ops.pop("XlaCallModule", None)
+
+  def tearDown(self):
+    super().tearDown()
+    if not uses_native_serialization():
+      if self._xla_op is not None:
+        tf2jax.ops._jax_ops["XlaCallModule"] = self._xla_op
+
   def _test_convert(
       self,
       jax_func,