Skip some tests that are broken with tf-nightly / ai-edge-litert-nightly.

Along the way, remove some skipTest calls for TF 2.10: the minimum
required version has been TF2.12 for a long time now.

PiperOrigin-RevId: 799138439

Author: arnoegw

tensorflow_gnn/graph/graph_tensor_ops_test.py

@@ -25,9 +25,11 @@
 from tensorflow_gnn.graph import graph_tensor_ops as ops
 from tensorflow_gnn.graph import pool_ops
 from tensorflow_gnn.graph import readout
-# pylint: disable=g-direct-tensorflow-import
-from ai_edge_litert import interpreter as tfl_interpreter
-# pylint: enable=g-direct-tensorflow-import
+# pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
+if not tf.__version__.startswith('2.20.'):  # TODO: b/441006328 - Remove this.
+  # The following import crashes with tf-nightly~=2.20.0.
+  from ai_edge_litert import interpreter as tfl_interpreter
+# pylint: enable=g-direct-tensorflow-import,g-import-not-at-top
 
 as_tensor = tf.convert_to_tensor
 as_ragged = tf.ragged.constant
@@ -639,12 +641,11 @@ def testTFLite(self):
     model = tf.keras.Model(inputs, outputs)
     expected = model(test_graph_dict).numpy()
 
-    # TODO(b/276291104): Remove when TF 2.11+ is required by all of TFGNN
-    if tf.__version__.startswith('2.10.'):
-      self.skipTest('GNN models are unsupported in TFLite until TF 2.11 but '
-                    f'got TF {tf.__version__}')
     converter = tf.lite.TFLiteConverter.from_keras_model(model)
     model_content = converter.convert()
+    if tf.__version__.startswith('2.20.'):
+      self.skipTest('TODO: b/441006328 - tfl_interpreter cannot be imported '
+                    f'next to tf-nightly~=2.20.0; got TF {tf.__version__}')
     interpreter = tfl_interpreter.Interpreter(model_content=model_content)
     signature_runner = interpreter.get_signature_runner('serving_default')
     obtained = signature_runner(**test_graph_dict)['final_edge_adjacency']
@@ -1305,12 +1306,11 @@ def testTFLite(self):
     model = tf.keras.Model(inputs, outputs)
     expected = model(test_graph_dict).numpy()
 
-    # TODO(b/276291104): Remove when TF 2.11+ is required by all of TFGNN
-    if tf.__version__.startswith('2.10.'):
-      self.skipTest('GNN models are unsupported in TFLite until TF 2.11 but '
-                    f'got TF {tf.__version__}')
     converter = tf.lite.TFLiteConverter.from_keras_model(model)
     model_content = converter.convert()
+    if tf.__version__.startswith('2.20.'):
+      self.skipTest('TODO: b/441006328 - tfl_interpreter cannot be imported '
+                    f'next to tf-nightly~=2.20.0; got TF {tf.__version__}')
     interpreter = tfl_interpreter.Interpreter(model_content=model_content)
     signature_runner = interpreter.get_signature_runner('serving_default')
     obtained = signature_runner(**test_graph_dict)['final_edge_adjacency']
@@ -1677,12 +1677,11 @@ def testTFLite(self):
     model = tf.keras.Model(inputs, outputs)
     expected = model(test_graph_dict).numpy()
 
-    # TODO(b/276291104): Remove when TF 2.11+ is required by all of TFGNN
-    if tf.__version__.startswith('2.10.'):
-      self.skipTest('GNN models are unsupported in TFLite until TF 2.11 but '
-                    f'got TF {tf.__version__}')
     converter = tf.lite.TFLiteConverter.from_keras_model(model)
     model_content = converter.convert()
+    if tf.__version__.startswith('2.20.'):
+      self.skipTest('TODO: b/441006328 - tfl_interpreter cannot be imported '
+                    f'next to tf-nightly~=2.20.0; got TF {tf.__version__}')
     interpreter = tfl_interpreter.Interpreter(model_content=model_content)
     signature_runner = interpreter.get_signature_runner('serving_default')
     obtained = signature_runner(**test_graph_dict)['final_edge_adjacency']

tensorflow_gnn/graph/graph_tensor_test.py

@@ -24,10 +24,12 @@
 from tensorflow_gnn.graph import graph_tensor as gt
 from tensorflow_gnn.graph import graph_tensor_test_utils as tu
 
-# pylint: disable=g-direct-tensorflow-import
-from ai_edge_litert import interpreter as tfl_interpreter
+# pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
+if not tf.__version__.startswith('2.20.'):  # TODO: b/441006328 - Remove this.
+  # The following import crashes with tf-nightly~=2.20.0.
+  from ai_edge_litert import interpreter as tfl_interpreter
 from tensorflow.python.framework import type_spec
-# pylint: enable=g-direct-tensorflow-import
+# pylint: enable=g-import-not-at-top,g-direct-tensorflow-import
 
 as_tensor = tf.convert_to_tensor
 as_ragged = tf.ragged.constant
@@ -1544,12 +1546,11 @@ def testTFLite(self):
     # The other unit tests should verify that this is correct
     expected = [0, 1, 2]
 
-    # TODO(b/276291104): Remove when TF 2.11+ is required by all of TFGNN
-    if tf.__version__.startswith('2.10.'):
-      self.skipTest('GNN models are unsupported in TFLite until TF 2.11 but '
-                    f'got TF {tf.__version__}')
     converter = tf.lite.TFLiteConverter.from_keras_model(model)
     model_content = converter.convert()
+    if tf.__version__.startswith('2.20.'):
+      self.skipTest('TODO: b/441006328 - tfl_interpreter cannot be imported '
+                    f'next to tf-nightly~=2.20.0; got TF {tf.__version__}')
     interpreter = tfl_interpreter.Interpreter(model_content=model_content)
     signature_runner = interpreter.get_signature_runner('serving_default')
     obtained = signature_runner(
@@ -1753,6 +1754,9 @@ def testTFLite(self):
 
     converter = tf.lite.TFLiteConverter.from_keras_model(model)
     model_content = converter.convert()
+    if tf.__version__.startswith('2.20.'):
+      self.skipTest('TODO: b/441006328 - tfl_interpreter cannot be imported '
+                    f'next to tf-nightly~=2.20.0; got TF {tf.__version__}')
     interpreter = tfl_interpreter.Interpreter(model_content=model_content)
     signature_runner = interpreter.get_signature_runner('serving_default')
     obtained = signature_runner(

tensorflow_gnn/keras/layers/graph_ops_test.py

@@ -23,9 +23,11 @@
 from tensorflow_gnn.graph import graph_constants as const
 from tensorflow_gnn.graph import graph_tensor as gt
 from tensorflow_gnn.keras.layers import graph_ops
-# pylint: disable=g-direct-tensorflow-import
-from ai_edge_litert import interpreter as tfl_interpreter
-# pylint: enable=g-direct-tensorflow-import
+# pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
+if not tf.__version__.startswith("2.20."):  # TODO: b/441006328 - Remove this.
+  # The following import crashes with tf-nightly~=2.20.0.
+  from ai_edge_litert import interpreter as tfl_interpreter
+# pylint: enable=g-direct-tensorflow-import,g-import-not-at-top
 
 
 class ReadoutTest(tf.test.TestCase, parameterized.TestCase):
@@ -168,12 +170,11 @@ def testTFLite(self, location):
     model = tf.keras.Model(inputs, outputs)
     expected = model(test_graph_134_dict)
 
-    # TODO(b/276291104): Remove when TF 2.11+ is required by all of TFGNN
-    if tf.__version__.startswith("2.10."):
-      self.skipTest("GNN models are unsupported in TFLite until TF 2.11 but "
-                    f"got TF {tf.__version__}")
     converter = tf.lite.TFLiteConverter.from_keras_model(model)
     model_content = converter.convert()
+    if tf.__version__.startswith("2.20."):
+      self.skipTest("TODO: b/441006328 - tfl_interpreter cannot be imported "
+                    f"next to tf-nightly~=2.20.0; got TF {tf.__version__}")
     interpreter = tfl_interpreter.Interpreter(model_content=model_content)
     signature_runner = interpreter.get_signature_runner("serving_default")
     obtained = signature_runner(**test_graph_134_dict)["test_readout"]
@@ -303,11 +304,11 @@ def testTFLite(self):
     expected = model(test_graph_22_dict)
 
     # TODO(b/276291104): Remove when TF 2.11+ is required by all of TFGNN
-    if tf.__version__.startswith("2.10."):
-      self.skipTest("GNN models are unsupported in TFLite until TF 2.11 but "
-                    f"got TF {tf.__version__}")
     converter = tf.lite.TFLiteConverter.from_keras_model(model)
     model_content = converter.convert()
+    if tf.__version__.startswith("2.20."):
+      self.skipTest("TODO: b/441006328 - tfl_interpreter cannot be imported "
+                    f"next to tf-nightly~=2.20.0; got TF {tf.__version__}")
     interpreter = tfl_interpreter.Interpreter(model_content=model_content)
     signature_runner = interpreter.get_signature_runner("serving_default")
     obtained = signature_runner(**test_graph_22_dict)["test_readout_first"]
@@ -435,6 +436,9 @@ def testTFLite(self):
 
     converter = tf.lite.TFLiteConverter.from_keras_model(model)
     model_content = converter.convert()
+    if tf.__version__.startswith("2.20."):
+      self.skipTest("TODO: b/441006328 - tfl_interpreter cannot be imported "
+                    f"next to tf-nightly~=2.20.0; got TF {tf.__version__}")
     interpreter = tfl_interpreter.Interpreter(model_content=model_content)
     signature_runner = interpreter.get_signature_runner("serving_default")
     actual = signature_runner(
@@ -569,6 +573,9 @@ def testTFLite(self):
 
     converter = tf.lite.TFLiteConverter.from_keras_model(model)
     model_content = converter.convert()
+    if tf.__version__.startswith("2.20."):
+      self.skipTest("TODO: b/441006328 - tfl_interpreter cannot be imported "
+                    f"next to tf-nightly~=2.20.0; got TF {tf.__version__}")
     interpreter = tfl_interpreter.Interpreter(model_content=model_content)
     signature_runner = interpreter.get_signature_runner("serving_default")
     actual = signature_runner(
@@ -632,6 +639,9 @@ def testTFLite(self):
 
     converter = tf.lite.TFLiteConverter.from_keras_model(model)
     model_content = converter.convert()
+    if tf.__version__.startswith("2.20."):
+      self.skipTest("TODO: b/441006328 - tfl_interpreter cannot be imported "
+                    f"next to tf-nightly~=2.20.0; got TF {tf.__version__}")
     interpreter = tfl_interpreter.Interpreter(model_content=model_content)
     signature_runner = interpreter.get_signature_runner("serving_default")
     actual = signature_runner(
@@ -748,12 +758,11 @@ def testTFLite(self):
     model = tf.keras.Model(inputs, outputs)
     expected = model(test_graph_134_dict).numpy()
 
-    # TODO(b/276291104): Remove when TF 2.11+ is required by all of TFGNN
-    if tf.__version__.startswith("2.10."):
-      self.skipTest("GNN models are unsupported in TFLite until TF 2.11 but "
-                    f"got TF {tf.__version__}")
     converter = tf.lite.TFLiteConverter.from_keras_model(model)
     model_content = converter.convert()
+    if tf.__version__.startswith("2.20."):
+      self.skipTest("TODO: b/441006328 - tfl_interpreter cannot be imported "
+                    f"next to tf-nightly~=2.20.0; got TF {tf.__version__}")
     interpreter = tfl_interpreter.Interpreter(model_content=model_content)
     signature_runner = interpreter.get_signature_runner("serving_default")
     obtained = signature_runner(**test_graph_134_dict)["final_edge_states"]
@@ -950,12 +959,11 @@ def testTFLite(self, tag, location):
     model = tf.keras.Model(inputs, outputs)
     expected = model(test_values)
 
-    # TODO(b/276291104): Remove when TF 2.11+ is required by all of TFGNN
-    if tf.__version__.startswith("2.10."):
-      self.skipTest("GNN models are unsupported in TFLite until TF 2.11 but "
-                    f"got TF {tf.__version__}")
     converter = tf.lite.TFLiteConverter.from_keras_model(model)
     model_content = converter.convert()
+    if tf.__version__.startswith("2.20."):
+      self.skipTest("TODO: b/441006328 - tfl_interpreter cannot be imported "
+                    f"next to tf-nightly~=2.20.0; got TF {tf.__version__}")
     interpreter = tfl_interpreter.Interpreter(model_content=model_content)
     signature_runner = interpreter.get_signature_runner("serving_default")
     obtained = signature_runner(**test_values)["test_broadcast"]
@@ -1258,12 +1266,11 @@ def testTFLite(self, tag, location, reduce_type):
     model = tf.keras.Model(inputs, outputs)
     expected = model(test_values)
 
-    # TODO(b/276291104): Remove when TF 2.11+ is required by all of TFGNN
-    if tf.__version__.startswith("2.10."):
-      self.skipTest("GNN models are unsupported in TFLite until TF 2.11 but "
-                    f"got TF {tf.__version__}")
     converter = tf.lite.TFLiteConverter.from_keras_model(model)
     model_content = converter.convert()
+    if tf.__version__.startswith("2.20."):
+      self.skipTest("TODO: b/441006328 - tfl_interpreter cannot be imported "
+                    f"next to tf-nightly~=2.20.0; got TF {tf.__version__}")
     interpreter = tfl_interpreter.Interpreter(model_content=model_content)
     signature_runner = interpreter.get_signature_runner("serving_default")
     obtained = signature_runner(**test_values)["test_pool"]

tensorflow_gnn/keras/layers/next_state_test.py

@@ -19,9 +19,11 @@
 from tensorflow_gnn.graph import graph_constants as const
 from tensorflow_gnn.keras.layers import next_state as next_state_lib
 from tensorflow_gnn.utils import tf_test_utils as tftu
-# pylint: disable=g-direct-tensorflow-import
-from ai_edge_litert import interpreter as tfl_interpreter
-# pylint: enable=g-direct-tensorflow-import
+# pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
+if not tf.__version__.startswith("2.20."):  # TODO: b/441006328 - Remove this.
+  # The following import crashes with tf-nightly~=2.20.0.
+  from ai_edge_litert import interpreter as tfl_interpreter
+# pylint: enable=g-direct-tensorflow-import,g-import-not-at-top
 
 
 class NextStateFromConcatTest(tf.test.TestCase, parameterized.TestCase):
@@ -182,6 +184,9 @@ def testTFLite(self):
 
     converter = tf.lite.TFLiteConverter.from_keras_model(model)
     model_content = converter.convert()
+    if tf.__version__.startswith("2.20."):
+      self.skipTest("TODO: b/441006328 - tfl_interpreter cannot be imported "
+                    f"next to tf-nightly~=2.20.0; got TF {tf.__version__}")
     interpreter = tfl_interpreter.Interpreter(model_content=model_content)
     signature_runner = interpreter.get_signature_runner("serving_default")
     obtained = signature_runner(**test_input_dict)["residual_next_state"]

tensorflow_gnn/keras/layers/padding_ops_test.py

@@ -23,9 +23,11 @@
 from tensorflow_gnn.keras import keras_tensors  # For registration. pylint: disable=unused-import
 from tensorflow_gnn.keras.layers import padding_ops
 from tensorflow_gnn.utils import tf_test_utils as tftu
-# pylint: disable=g-direct-tensorflow-import
-from ai_edge_litert import interpreter as tfl_interpreter
-# pylint: enable=g-direct-tensorflow-import
+# pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
+if not tf.__version__.startswith("2.20."):  # TODO: b/441006328 - Remove this.
+  # The following import crashes with tf-nightly~=2.20.0.
+  from ai_edge_litert import interpreter as tfl_interpreter
+# pylint: enable=g-direct-tensorflow-import,g-import-not-at-top
 
 
 class PadToTotalSizesTest(tf.test.TestCase, parameterized.TestCase):

tensorflow_gnn/models/gat_v2/layers_test.py

@@ -17,9 +17,11 @@
 import tensorflow_gnn as tfgnn
 from tensorflow_gnn.models import gat_v2
 from tensorflow_gnn.utils import tf_test_utils as tftu
-# pylint: disable=g-direct-tensorflow-import
-from ai_edge_litert import interpreter as tfl_interpreter
-# pylint: enable=g-direct-tensorflow-import
+# pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
+if not tf.__version__.startswith("2.20."):  # TODO: b/441006328 - Remove this.
+  # The following import crashes with tf-nightly~=2.20.0.
+  from ai_edge_litert import interpreter as tfl_interpreter
+# pylint: enable=g-direct-tensorflow-import,g-import-not-at-top
 
 
 class GATv2Test(tf.test.TestCase, parameterized.TestCase):
@@ -199,6 +201,8 @@ def testMultihead(self, merge_type):
       want = tf.nn.relu((want_logits[:, :4] + want_logits[:, 4:]) / 2)
     elif merge_type == "concat":
       want = tf.nn.relu(want_logits)
+    else:
+      raise AssertionError(f"Bad test set-up: unknown merge_type {merge_type}")
 
     self.assertAllClose(got, want, atol=0.0001)
 
@@ -701,12 +705,11 @@ def testBasic(self):
     # The other unit tests should verify that this is correct
     expected = model(test_graph_1_dict).numpy()
 
-    # TODO(b/276291104): Remove when TF 2.11+ is required by all of TFGNN
-    if tf.__version__.startswith("2.10."):
-      self.skipTest("GNN models are unsupported in TFLite until TF 2.11 but "
-                    f"got TF {tf.__version__}")
     converter = tf.lite.TFLiteConverter.from_keras_model(model)
     model_content = converter.convert()
+    if tf.__version__.startswith("2.20."):
+      self.skipTest("TODO: b/441006328 - tfl_interpreter cannot be imported "
+                    f"next to tf-nightly~=2.20.0; got TF {tf.__version__}")
     interpreter = tfl_interpreter.Interpreter(model_content=model_content)
     signature_runner = interpreter.get_signature_runner("serving_default")
     obtained = signature_runner(**test_graph_1_dict)["final_node_states"]

tensorflow_gnn/models/gcn/gcn_conv_test.py

@@ -19,9 +19,11 @@
 import tensorflow_gnn as tfgnn
 from tensorflow_gnn.models.gcn import gcn_conv
 from tensorflow_gnn.utils import tf_test_utils as tftu
-# pylint: disable=g-direct-tensorflow-import
-from ai_edge_litert import interpreter as tfl_interpreter
-# pylint: enable=g-direct-tensorflow-import
+# pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
+if not tf.__version__.startswith('2.20.'):  # TODO: b/441006328 - Remove this.
+  # The following import crashes with tf-nightly~=2.20.0.
+  from ai_edge_litert import interpreter as tfl_interpreter
+# pylint: enable=g-direct-tensorflow-import,g-import-not-at-top
 
 
 class GcnConvTest(tf.test.TestCase, parameterized.TestCase):
@@ -869,12 +871,11 @@ def testBasic(self, add_self_loops, edge_weight_feature_name):
     # The other unit tests should verify that this is correct
     expected = model(test_graph_1_dict).numpy()
 
-    # TODO(b/276291104): Remove when TF 2.11+ is required by all of TFGNN
-    if tf.__version__.startswith('2.10.'):
-      self.skipTest('GNN models are unsupported in TFLite until TF 2.11 but '
-                    f'got TF {tf.__version__}')
     converter = tf.lite.TFLiteConverter.from_keras_model(model)
     model_content = converter.convert()
+    if tf.__version__.startswith('2.20.'):
+      self.skipTest('TODO: b/441006328 - tfl_interpreter cannot be imported '
+                    f'next to tf-nightly~=2.20.0; got TF {tf.__version__}')
     interpreter = tfl_interpreter.Interpreter(model_content=model_content)
     signature_runner = interpreter.get_signature_runner('serving_default')
     obtained = signature_runner(**test_graph_1_dict)['final_node_states']

tensorflow_gnn/models/graph_sage/layers_test.py

@@ -19,9 +19,11 @@
 import tensorflow_gnn as tfgnn
 from tensorflow_gnn.models.graph_sage import layers as graph_sage
 from tensorflow_gnn.utils import tf_test_utils as tftu
-# pylint: disable=g-direct-tensorflow-import
-from ai_edge_litert import interpreter as tfl_interpreter
-# pylint: enable=g-direct-tensorflow-import
+# pylint: disable=g-direct-tensorflow-import,g-import-not-at-top
+if not tf.__version__.startswith("2.20."):  # TODO: b/441006328 - Remove this.
+  # The following import crashes with tf-nightly~=2.20.0.
+  from ai_edge_litert import interpreter as tfl_interpreter
+# pylint: enable=g-direct-tensorflow-import,g-import-not-at-top
 
 _FEATURE_NAME = "f"
 
@@ -620,12 +622,11 @@ def testBasic(self, use_pooling, hidden_units, combine_type):
     # The other unit tests should verify that this is correct
     expected = model(test_graph_1_dict).numpy()
 
-    # TODO(b/276291104): Remove when TF 2.11+ is required by all of TFGNN
-    if tf.__version__.startswith("2.10."):
-      self.skipTest("GNN models are unsupported in TFLite until TF 2.11 but "
-                    f"got TF {tf.__version__}")
     converter = tf.lite.TFLiteConverter.from_keras_model(model)
     model_content = converter.convert()
+    if tf.__version__.startswith("2.20."):
+      self.skipTest("TODO: b/441006328 - tfl_interpreter cannot be imported "
+                    f"next to tf-nightly~=2.20.0; got TF {tf.__version__}")
     interpreter = tfl_interpreter.Interpreter(model_content=model_content)
     signature_runner = interpreter.get_signature_runner("serving_default")
     obtained = signature_runner(**test_graph_1_dict)["final_node_states"]