Enable test case: subgraph sampling model, saving and loading, as tf.keras.Model.

PiperOrigin-RevId: 534966793

Author: samihaija

tensorflow_gnn/experimental/in_memory/int_arithmetic_sampler_test.py

@@ -368,7 +368,7 @@ def test_graph_sampler_constructor_without_edgesets_raises_exception(self):
 
   @parameterized.named_parameters(
       ('TFLoadSavedModel', 'TFLoadModel'),
-      # ('KerasLoadSavedGModel', 'KerasLoadModel'),  # blocked by cl/530981500
+      ('KerasLoadSavedGModel', 'KerasLoadModel'),
       ('KerasModel', 'KerasModel')
   )
   def test_smoke_integration_sampling_pipeline_model(self, operation_mode):

tensorflow_gnn/experimental/sampler/core.py

@@ -1030,6 +1030,9 @@ def get_config(self):
     return dict(validate=self._validate, **super().get_config())
 
   def call(self, inputs: GraphPieces) -> tfgnn.GraphTensor:
+    # When saving `tf.keras.Model`, the loaded model passes a tuple.
+    # Convert to namedtuple.
+    inputs = GraphPieces(*inputs)
     context: Features = inputs.context
     node_sets: Mapping[tfgnn.NodeSetName,
                        Union[Features, List[Features]]] = inputs.node_sets