Finish adding SavedModel support

examples/keras_io/tensorflow/keras_recipes/tf_serving.py

@@ -364,7 +364,7 @@ channel = grpc.insecure_channel("localhost:8500")
 
 """
 ```shell
-!pip install -q tensorflow_serving_api
+pip install -q tensorflow_serving_api
 ```
 
 ```python
@@ -378,8 +378,6 @@ stub = prediction_service_pb2_grpc.PredictionServiceStub(channel)
 
 # Get the serving_input key
 loaded_model = tf.saved_model.load(model_export_path)
-print(loaded_model)
-print(loaded_model.signatures)
 input_name = list(
     loaded_model.signatures["serving_default"]
     .structured_input_signature[1]

keras_core/backend/tensorflow/layer.py

@@ -6,17 +6,43 @@ class TFLayer(tf.__internal__.tracking.AutoTrackable):
     def _default_save_signature(self):
         """For SavedModel support: returns the default serving signature."""
 
-        shapes_dict = self._build_shapes_dict
-        if len(shapes_dict) == 1:
-            input_shape = tuple(shapes_dict.values())[0]
-            input_signature = [tf.TensorSpec(input_shape, self.compute_dtype)]
-        else:
+        from keras_core.models.functional import Functional
+        from keras_core.models.model import Model
+        from keras_core.models.sequential import Sequential
+
+        if not isinstance(self, Model):
+            return None
+
+        inputs = None
+        if (
+            isinstance(self, Sequential)
+            and getattr(self, "_functional", None) is not None
+        ):
+            inputs = self._functional.input
+        elif isinstance(self, Functional):
+            inputs = self.input
+
+        if inputs is not None:
             input_signature = [
                 tf.nest.map_structure(
                     lambda x: tf.TensorSpec(x.shape, self.compute_dtype),
-                    shapes_dict,
+                    inputs,
                 )
             ]
+        else:
+            shapes_dict = self._build_shapes_dict
+            if len(shapes_dict) == 1:
+                input_shape = tuple(shapes_dict.values())[0]
+                input_signature = [
+                    tf.TensorSpec(input_shape, self.compute_dtype)
+                ]
+            else:
+                input_signature = [
+                    tf.nest.map_structure(
+                        lambda x: tf.TensorSpec(x.shape, self.compute_dtype),
+                        shapes_dict,
+                    )
+                ]
 
         @tf.function(input_signature=input_signature)
         def serving_default(inputs):

keras_core/layers/pooling/base_pooling.py

@@ -1,9 +1,8 @@
-from keras_core import backend
 from keras_core import operations as ops
+from keras_core.backend import image_data_format
 from keras_core.layers.input_spec import InputSpec
 from keras_core.layers.layer import Layer
 from keras_core.operations.operation_utils import compute_pooling_output_shape
-from keras_core.utils import argument_validation
 
 
 class BasePooling(Layer):
@@ -22,16 +21,13 @@ class BasePooling(Layer):
     ):
         super().__init__(name=name, **kwargs)
 
-        self.pool_size = argument_validation.standardize_tuple(
-            pool_size, pool_dimensions, "pool_size"
-        )
-        strides = pool_size if strides is None else strides
-        self.strides = argument_validation.standardize_tuple(
-            strides, pool_dimensions, "strides", allow_zero=True
-        )
+        self.pool_size = pool_size
+        self.strides = pool_size if strides is None else strides
         self.pool_mode = pool_mode
         self.padding = padding
-        self.data_format = backend.standardize_data_format(data_format)
+        self.data_format = (
+            image_data_format() if data_format is None else data_format
+        )
 
         self.input_spec = InputSpec(ndim=pool_dimensions + 2)
 

keras_core/operations/numpy.py

@@ -1514,71 +1514,33 @@ def full_like(x, fill_value, dtype=None):
 
 class GetItem(Operation):
     def call(self, x, key):
+        if not isinstance(key, int):
+            # TODO: support slicing.
+            raise ValueError(
+                "Only scalar int keys are supported at this time. Cannot "
+                f"process key {key}"
+            )
         return x[key]
 
     def compute_output_spec(self, x, key):
-        remaining_shape = list(x.shape)
-        new_shape = []
-        if isinstance(key, int):
-            remaining_key = [key]
-        elif isinstance(key, tuple):
-            remaining_key = list(key)
-        else:
+        if not isinstance(key, int):
+            # TODO: support slicing.
             raise ValueError(
-                f"Unsupported key type for array slice. Recieved: `{key}`"
+                "Only scalar int keys are supported at this time. Cannot "
+                f"process key {key}"
             )
-        num_ellipses = remaining_key.count(Ellipsis)
-        if num_ellipses > 1:
+        if len(x.shape) == 0:
             raise ValueError(
-                f"Slice should only have one ellipsis. Recieved: `{key}`"
+                f"Too many indices for array: array is scalar "
+                f"but index {key} was requested. A scalar array "
+                "cannot be indexed."
             )
-        elif num_ellipses == 0:
-            # Add an implicit final ellipsis.
-            remaining_key.append(Ellipsis)
-        # Consume slice key element by element.
-        while True:
-            if not remaining_key:
-                break
-            subkey = remaining_key.pop(0)
-            # Check for `newaxis` and `Ellipsis`.
-            if subkey == Ellipsis:
-                # Keep as many slices remain in our key, omitting `newaxis`.
-                needed = len(remaining_key) - remaining_key.count(np.newaxis)
-                consumed = len(remaining_shape) - needed
-                new_shape += remaining_shape[:consumed]
-                remaining_shape = remaining_shape[consumed:]
-                continue
-            # All frameworks follow numpy for newaxis. `np.newaxis == None`.
-            if subkey == np.newaxis:
-                new_shape.append(1)
-                continue
-            # At this point, we need to consume a new axis from the shape.
-            if not remaining_shape:
-                raise ValueError(
-                    f"Array has shape {x.shape} but slice "
-                    f"has to many indices. Recieved: `{key}`"
-                )
-            length = remaining_shape.pop(0)
-            if isinstance(subkey, int):
-                if length is not None:
-                    index = subkey if subkey >= 0 else subkey + length
-                    if index < 0 or index >= length:
-                        raise ValueError(
-                            f"Array has shape {x.shape} but out-of-bounds "
-                            f"index {key} was requested."
-                        )
-            elif isinstance(subkey, slice):
-                if length is not None:
-                    # python3 friendly way to compute a slice length.
-                    new_length = len(range(*subkey.indices(length)))
-                    new_shape.append(new_length)
-                else:
-                    new_shape.append(length)
-            else:
-                raise ValueError(
-                    f"Unsupported key type for array slice. Recieved: `{key}`"
-                )
-        return KerasTensor(tuple(new_shape), dtype=x.dtype)
+        if x.shape[0] is not None and key >= x.shape[0]:
+            raise ValueError(
+                f"Array has shape {x.shape} "
+                f"but out-of-bound index {key} was requested."
+            )
+        return KerasTensor(x.shape[1:], dtype=x.dtype)
 
 
 @keras_core_export(
@@ -1587,6 +1549,12 @@ class GetItem(Operation):
 def get_item(x, key):
     if any_symbolic_tensors((x,)):
         return GetItem().symbolic_call(x, key)
+    if not isinstance(key, int):
+        # TODO: support slicing.
+        raise ValueError(
+            "Only scalar int keys are supported at this time. Cannot "
+            f"process key {key}"
+        )
     return x[key]
 
 

keras_core/operations/numpy_test.py

@@ -882,26 +882,8 @@ class NumpyOneInputOpsDynamicShapeTest(testing.TestCase):
         self.assertEqual(knp.floor(x).shape, (None, 3))
 
     def test_get_item(self):
-        x = KerasTensor([None, 5, 16])
-        # Simple slice.
-        sliced = knp.get_item(x, 5)
-        self.assertEqual(sliced.shape, (5, 16))
-        # Ellipsis slice.
-        sliced = knp.get_item(x, np.s_[..., -1])
-        self.assertEqual(sliced.shape, (None, 5))
-        # `newaxis` slice.
-        sliced = knp.get_item(x, np.s_[:, np.newaxis, ...])
-        self.assertEqual(sliced.shape, (None, 1, 5, 16))
-        # Strided slice.
-        sliced = knp.get_item(x, np.s_[:5, 3:, 3:12:2])
-        self.assertEqual(sliced.shape, (None, 2, 5))
-        # Error states.
-        with self.assertRaises(ValueError):
-            sliced = knp.get_item(x, np.s_[:, 17, :])
-        with self.assertRaises(ValueError):
-            sliced = knp.get_item(x, np.s_[..., 5, ...])
-        with self.assertRaises(ValueError):
-            sliced = knp.get_item(x, np.s_[:, :, :, :])
+        x = KerasTensor([None, None])
+        self.assertEqual(knp.get_item(x, 5).shape, (None,))
 
     def test_hstack(self):
         x = KerasTensor([None, 3])

keras_core/operations/symbolic_arguments.py

@@ -1,5 +1,3 @@
-import typing
-
 from tensorflow import nest
 
 from keras_core.backend import KerasTensor
@@ -43,10 +41,9 @@ class SymbolicArguments:
             return (tensor_dict[self._single_positional_tensor],), {}
 
         def switch_fn(x):
-            if isinstance(x, typing.Hashable):
-                val = tensor_dict.get(x, None)
-                if val is not None:
-                    return val
+            val = tensor_dict.get(x, None)
+            if val is not None:
+                return val
             return x
 
         return self.convert(switch_fn)