Disable JAX serialization tests that fail after a JAX refactoring (https://github.com/google/jax/pull/10232).
JAX does not promise anything in particular about the contents of the HLO it produces, so tests that pattern match on that HLO may be broken at any time.
# limitations under the License.
import collections
+import unittest
from absl.testing import absltest
import jax
@@ -29,6 +30,8 @@ from tensorflow_federated.python.core.impl.xla_context import xla_serialization
class JaxSerializationTest(absltest.TestCase):
+ @unittest.skip('HLO pattern matching broken by '
+ 'https://github.com/google/jax/pull/10232')
def test_serialize_jax_with_noarg_to_int32(self):
def traced_fn():
@@ -49,6 +52,8 @@ class JaxSerializationTest(absltest.TestCase):
self.assertEqual(str(comp_pb.xla.parameter), '')
self.assertEqual(str(comp_pb.xla.result), 'tensor {\n' ' index: 0\n' '}\n')
+ @unittest.skip('HLO pattern matching broken by '
+ 'https://github.com/google/jax/pull/10232')
def test_serialize_jax_with_int32_to_int32(self):
def traced_fn(x):
@@ -68,6 +73,8 @@ class JaxSerializationTest(absltest.TestCase):
self.assertEqual(str(comp_pb.xla.result), str(comp_pb.xla.parameter))
self.assertEqual(str(comp_pb.xla.result), 'tensor {\n' ' index: 0\n' '}\n')
+ @unittest.skip('HLO pattern matching broken by '
+ 'https://github.com/google/jax/pull/10232')
def test_serialize_jax_with_2xint32_to_2xint32(self):
def traced_fn(x):
@@ -112,6 +119,8 @@ class JaxSerializationTest(absltest.TestCase):
' }\n'
'}\n')
+ @unittest.skip('HLO pattern matching broken by '
+ 'https://github.com/google/jax/pull/10232')
def test_serialize_jax_with_two_args(self):
def traced_fn(x, y):
