Make vectorized_map op serializable. (#21597)

All ops that take tensors as inputs are in scope for serialization. They get serialized if they appear as part of a functional model.

The `VectorizedMap` op class needs a custom `get_config()` and `from_config()` to handle the `function` parameter, which is not a simple type. The fallback `get_config` logic only handles simple types.

Note that some other ops will need a similar change, which will come in a separate PR because it requires a another change.

Author: hertschuh

keras/src/backend/openvino/excluded_concrete_tests.txt

@@ -186,6 +186,7 @@ CoreOpsCorrectnessTest::test_scatter
 CoreOpsCorrectnessTest::test_switch 
 CoreOpsCorrectnessTest::test_unstack
 CoreOpsCorrectnessTest::test_vectorized_map
+CoreOpsBehaviorTests::test_vectorized_map_serialization
 ExtractSequencesOpTest::test_extract_sequences_call
 InTopKTest::test_in_top_k_call
 MathOpsCorrectnessTest::test_erfinv_operation_basic

keras/src/ops/core.py

@@ -8,6 +8,7 @@
 from keras.src.backend import any_symbolic_tensors
 from keras.src.backend.common.backend_utils import slice_along_axis
 from keras.src.ops.operation import Operation
+from keras.src.saving import serialization_lib
 from keras.src.utils import traceback_utils
 
 
@@ -1105,6 +1106,19 @@ def append_batch_axis(t):
         y = tree.map_structure(append_batch_axis, y)
         return y
 
+    def get_config(self):
+        config = super().get_config()
+        config.update({"function": self.function})
+        return config
+
+    @classmethod
+    def from_config(cls, config):
+        config = config.copy()
+        config["function"] = serialization_lib.deserialize_keras_object(
+            config["function"]
+        )
+        return cls(**config)
+
 
 @keras_export("keras.ops.vectorized_map")
 def vectorized_map(function, elements):

keras/src/ops/core_test.py

@@ -16,7 +16,9 @@
 from keras.src import tree
 from keras.src.backend.common import dtypes
 from keras.src.backend.common.keras_tensor import KerasTensor
+from keras.src.layers.core import input_layer
 from keras.src.ops import core
+from keras.src.saving import object_registration
 from keras.src.testing.test_utils import named_product
 
 
@@ -1622,6 +1624,18 @@ def test_stop_gradient_compute_output_spec(self):
         self.assertEqual(output_spec.shape, variable.shape)
         self.assertEqual(output_spec.dtype, variable.dtype)
 
+    def test_vectorized_map_serialization(self):
+        @object_registration.register_keras_serializable()
+        def f(x):
+            return x + x
+
+        inputs = input_layer.Input((10,), dtype="float32")
+        outputs = core.vectorized_map(f, inputs)
+        model = models.Functional(inputs, outputs)
+        reloaded_model = model.from_config(model.get_config())
+        x = np.random.rand(5, 10).astype("float32")
+        self.assertAllClose(model(x), reloaded_model(x))
+
     def test_while_loop_output_spec(self):
         # Define dummy cond and body functions
         def cond(x):