Reinstate the tests that were skipped for TF/Keras [nightly] 2.20 by

commit f9b8f5c

PiperOrigin-RevId: 803321430

Author: arnoegw

tensorflow_gnn/graph/graph_tensor_ops_test.py

@@ -25,11 +25,9 @@
 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,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
+# pylint: disable=g-direct-tensorflow-import
+from ai_edge_litert import interpreter as tfl_interpreter
+# pylint: enable=g-direct-tensorflow-import
 
 as_tensor = tf.convert_to_tensor
 as_ragged = tf.ragged.constant
@@ -643,9 +641,6 @@ 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_graph_dict)['final_edge_adjacency']
@@ -1308,9 +1303,6 @@ 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_graph_dict)['final_edge_adjacency']
@@ -1679,9 +1671,6 @@ 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_graph_dict)['final_edge_adjacency']

tensorflow_gnn/graph/graph_tensor_test.py

@@ -24,12 +24,10 @@
 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,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: disable=g-direct-tensorflow-import
+from ai_edge_litert import interpreter as tfl_interpreter
 from tensorflow.python.framework import type_spec
-# pylint: enable=g-import-not-at-top,g-direct-tensorflow-import
+# pylint: enable=g-direct-tensorflow-import
 
 as_tensor = tf.convert_to_tensor
 as_ragged = tf.ragged.constant
@@ -1548,9 +1546,6 @@ 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(
@@ -1754,9 +1749,6 @@ 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,11 +23,9 @@
 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,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
+# pylint: disable=g-direct-tensorflow-import
+from ai_edge_litert import interpreter as tfl_interpreter
+# pylint: enable=g-direct-tensorflow-import
 
 
 class ReadoutTest(tf.test.TestCase, parameterized.TestCase):
@@ -172,9 +170,6 @@ def testTFLite(self, location):
 
     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,12 +298,8 @@ def testTFLite(self):
     model = tf.keras.Model(inputs, outputs)
     expected = model(test_graph_22_dict)
 
-    # TODO(b/276291104): Remove when TF 2.11+ is required by all of TFGNN
     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"]
@@ -436,9 +427,6 @@ 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(
@@ -573,9 +561,6 @@ 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(
@@ -639,9 +624,6 @@ 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(
@@ -760,9 +742,6 @@ 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_graph_134_dict)["final_edge_states"]
@@ -961,9 +940,6 @@ def testTFLite(self, tag, location):
 
     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"]
@@ -1268,9 +1244,6 @@ def testTFLite(self, tag, location, reduce_type):
 
     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,11 +19,9 @@
 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,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
+# pylint: disable=g-direct-tensorflow-import
+from ai_edge_litert import interpreter as tfl_interpreter
+# pylint: enable=g-direct-tensorflow-import
 
 
 class NextStateFromConcatTest(tf.test.TestCase, parameterized.TestCase):
@@ -184,9 +182,6 @@ 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,11 +23,9 @@
 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,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
+# pylint: disable=g-direct-tensorflow-import
+from ai_edge_litert import interpreter as tfl_interpreter
+# pylint: enable=g-direct-tensorflow-import
 
 
 class PadToTotalSizesTest(tf.test.TestCase, parameterized.TestCase):

tensorflow_gnn/models/gat_v2/layers_test.py

@@ -17,11 +17,9 @@
 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,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
+# pylint: disable=g-direct-tensorflow-import
+from ai_edge_litert import interpreter as tfl_interpreter
+# pylint: enable=g-direct-tensorflow-import
 
 
 class GATv2Test(tf.test.TestCase, parameterized.TestCase):
@@ -707,9 +705,6 @@ def testBasic(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_graph_1_dict)["final_node_states"]

tensorflow_gnn/models/gcn/gcn_conv_test.py

@@ -19,11 +19,9 @@
 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,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
+# pylint: disable=g-direct-tensorflow-import
+from ai_edge_litert import interpreter as tfl_interpreter
+# pylint: enable=g-direct-tensorflow-import
 
 
 class GcnConvTest(tf.test.TestCase, parameterized.TestCase):
@@ -873,9 +871,6 @@ def testBasic(self, add_self_loops, edge_weight_feature_name):
 
     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,11 +19,9 @@
 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,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
+# pylint: disable=g-direct-tensorflow-import
+from ai_edge_litert import interpreter as tfl_interpreter
+# pylint: enable=g-direct-tensorflow-import
 
 _FEATURE_NAME = "f"
 
@@ -624,9 +622,6 @@ def testBasic(self, use_pooling, hidden_units, combine_type):
 
     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/hgt/layers_test.py

@@ -19,11 +19,9 @@
 import tensorflow_gnn as tfgnn
 from tensorflow_gnn.models.hgt import layers
 from tensorflow_gnn.utils import tf_test_utils as tftu
-# 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
+# pylint: disable=g-direct-tensorflow-import
+from ai_edge_litert import interpreter as tfl_interpreter
+# pylint: enable=g-direct-tensorflow-import
 
 
 def _homogeneous_cycle_graph(node_state, edge_state=None):
@@ -810,9 +808,6 @@ def testBasic(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_graph_1_dict)["final_engine_states"]

tensorflow_gnn/models/mt_albis/layers_test.py

@@ -22,11 +22,9 @@
 import tensorflow_gnn as tfgnn
 from tensorflow_gnn.models.mt_albis import layers
 from tensorflow_gnn.utils import tf_test_utils as tftu
-# 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
+# pylint: disable=g-direct-tensorflow-import
+from ai_edge_litert import interpreter as tfl_interpreter
+# pylint: enable=g-direct-tensorflow-import
 
 
 class MtAlbisNextNodeStateTest(tf.test.TestCase, parameterized.TestCase):
@@ -383,9 +381,6 @@ def test(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_graph_1_dict)["final_node_states"]