fix(export+models): Enhance support for dictionary-based model input signatures in TensorFlow and JAX #20842

harshaljanjani · 2025-02-02T10:18:26Z

This fix goes beyond the requirements of the issue and adds support for handling Keras models with dictionary-based inputs, particularly when exporting to the TFSavedModel format for both the TensorFlow and JAX backends. Previously, models with dictionary inputs would fail during export with ValueErrors related to input structure mismatches.

Key changes:

Added proper handling of model._input_names for dictionary-based inputs in Functional and Model classes
Enhanced input signature generation to properly handle dictionary input structures
Added test coverage for dictionary input model export
Added debug logging throughout the input signature handling flow

This PR aims to fix #20835 where models with dictionary inputs would fail to export properly to SavedModel format.

Example of fixed functionality:

inputs = {
  "foo": layers.Input(shape=()),
  "bar": layers.Input(shape=()),
}
outputs = layers.Add()([inputs["foo"], inputs["bar"]])
model = models.Model(inputs, outputs)
ref_input = {"foo": tf.constant([1.0]), "bar": tf.constant([2.0])}
ref_output = model(ref_input)
model.export('/test', format="tf_saved_model")
revived_model = tf.saved_model.load(temp_filepath)
revived_output = revived_model.serve(ref_input)

…ased models

…l_test.py

…cases) - Improves input structure validation in Model and Functional classes - Adds strict validation with clear error messages for mismatches

codecov-commenter · 2025-02-03T03:09:40Z

Codecov Report

Attention: Patch coverage is 59.61538% with 42 lines in your changes missing coverage. Please review.

Project coverage is 82.20%. Comparing base (e045b6a) to head (72371d7).
Report is 20 commits behind head on master.

Files with missing lines	Patch %	Lines
keras/src/models/functional.py	55.76%	12 Missing and 11 partials ⚠️
keras/src/models/model.py	21.05%	10 Missing and 5 partials ⚠️
keras/src/export/export_utils.py	71.42%	2 Missing and 2 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master   #20842      +/-   ##
==========================================
- Coverage   82.26%   82.20%   -0.06%     
==========================================
  Files         561      561              
  Lines       52693    53035     +342     
  Branches     8146     8228      +82     
==========================================
+ Hits        43347    43600     +253     
- Misses       7344     7391      +47     
- Partials     2002     2044      +42

Flag	Coverage Δ
keras	`82.02% <59.61%> (-0.06%)`	⬇️
keras-jax	`64.14% <54.80%> (-0.11%)`	⬇️
keras-numpy	`58.89% <23.07%> (-0.10%)`	⬇️
keras-openvino	`32.36% <1.92%> (-0.19%)`	⬇️
keras-tensorflow	`64.60% <46.15%> (-0.24%)`	⬇️
keras-torch	`64.16% <39.42%> (-0.12%)`	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

- Make the linter happy :)

- Added shape logging for preprocessed and crossed features - Added debug messages in Functional model processing (To be removed)

harshaljanjani · 2025-02-03T18:46:52Z

TODO: Remove debugging prints.
Edit: Done and ready for review.

harshaljanjani · 2025-02-04T02:43:55Z

Ping: @mattdangerw @fchollet.

jeffcarp

Thanks for the contribution! Left some comments.

keras/src/export/saved_model_test.py

keras/src/layers/preprocessing/feature_space_test.py

keras/src/models/functional_test.py

keras/src/models/model.py

keras/src/models/functional.py

harshaljanjani · 2025-02-06T14:42:42Z

Thanks for the review! Will fix it at the earliest.

harshaljanjani · 2025-02-08T19:39:09Z

Please do look into the changes and let me know if it's the intended behavior. Looking forward to your guidance, thanks!

harshaljanjani · 2025-02-17T02:01:59Z

Ping: @fchollet, looking forward to an update on this PR. Thanks!

fchollet · 2025-02-17T02:51:54Z

Thanks for the update!

This branch has conflicts that must be resolved

Please resolve merge conflicts.

@jeffcarp does the PR look good?

harshaljanjani · 2025-02-17T05:15:49Z

@fchollet
Resolved merge conflicts, thanks!

jeffcarp

LGTM % some nits

keras/src/models/functional.py

harshaljanjani · 2025-02-19T01:59:45Z

@fchollet @jeffcarp
Fixed! Thanks for reviewing the PR.

jeffcarp

Thanks!

harshaljanjani · 2025-02-22T13:10:09Z

@fchollet, just checking in on this PR. Since the changes are approved, is there anything else needed before merging? Looking forward to your feedback, thanks!

hertschuh

Thank you for your work on this. I can see that a lot of effort was put into this. However, I believe this can be simplified quite a bit.

There is a lot of logic around creating, maintaining, restoring self._input_names. However, it is only actually needed in export_utils.get_input_signature, and it's only used for one purpose, which is to know the order of inputs. But the order of inputs is already known, right? It's in self._input_struct. So you can just use self._input_struct in export_utils.get_input_signature and remove self._input_names.

Additionally, I don't think dicts constitute a special case in most cases. Layers can have 1 or more positional arguments as inputs, and each one of those is a nested structure that can have dicts at any layer.

Out of curiosity, I added the following 2 tests in saved_model_test.py, and they don't pass:

    def test_export_with_two_dict_inputs_functional(self):
        temp_filepath = os.path.join(self.get_temp_dir(), "exported_model")
        inputs1 = {
            "foo": layers.Input(shape=()),
            "bar": layers.Input(shape=()),
        }
        inputs2 = {
            "baz": layers.Input(shape=()),
        }
        outputs = layers.Add()([inputs1["foo"], inputs1["bar"], inputs2["baz"]])
        model = models.Model((inputs1, inputs2), outputs)
        ref_input = (
            {"foo": tf.constant([1.0]), "bar": tf.constant([2.0])},
            {"baz": tf.constant([2.0])},
        )
        ref_output = model(ref_input)
        model.export(temp_filepath, format="tf_saved_model")
        revived_model = tf.saved_model.load(temp_filepath)
        revived_output = revived_model.serve(ref_input)
        self.assertAllClose(ref_output, revived_output)

    def test_export_with_two_dict_inputs_subclass(self):
        class TwoDictModel(models.Model):
            def call(self, inputs1, inputs2):
                return ops.add(
                    ops.add(inputs1["foo"], inputs1["bar"]), inputs2["baz"]
                )

        temp_filepath = os.path.join(self.get_temp_dir(), "exported_model")
        model = TwoDictModel()
        ref_input1 = {"foo": tf.constant([1.0]), "bar": tf.constant([2.0])}
        ref_input2 = {"baz": tf.constant([2.0])}
        ref_output = model(ref_input1, ref_input2)
        model.export(temp_filepath, format="tf_saved_model")
        revived_model = tf.saved_model.load(temp_filepath)
        revived_output = revived_model.serve(ref_input1, ref_input2)
        self.assertAllClose(ref_output, revived_output)

This shows that dicts as second arguments are not supported.

hertschuh · 2025-02-27T19:59:44Z

keras/src/backend/jax/export.py

+            # Create a simplified wrapper that handles both dict and
+            # positional args, similar to TensorFlow implementation.
+            def wrapped_fn(arg, **kwargs):
+                return fn(arg)


I'm not following

why this is needed

what tensorflow implementation it follows

All it does is keep the first position argument and drop everything else. What if you have multiple positional arguments?

hertschuh · 2025-02-27T20:02:18Z

keras/src/backend/jax/export.py

-            jax2tf_kwargs["polymorphic_shapes"] = self._to_polymorphic_shape(
-                input_signature
-            )
+            if isinstance(input_signature, dict):


This shouldn't be needed. input_signature should always be a list because the model can have multiple positional arguments. If the inputs are a dict, input_signature should be a list of one dict already. If you need to add the outer list here, it means that an incorrect input_signature was passed in the first place.

hertschuh · 2025-02-27T20:03:29Z

keras/src/backend/tensorflow/export.py

-        decorated_fn = tf.function(
-            fn, input_signature=input_signature, autograph=False
-        )
+        if isinstance(input_signature, dict):


Same comment about why this shouldn't be needed. input_signature should always be a list because the model can have multiple positional arguments. If the inputs are a dict, input_signature should be a list of one dict already. If you need to add the outer list here, it means that an incorrect input_signature was passed in the first place.

hertschuh · 2025-02-27T20:04:04Z

keras/src/backend/tensorflow/export.py

+        if isinstance(input_signature, dict):
+            # Create a simplified wrapper that handles both dict and
+            # positional args.
+            def wrapped_fn(arg, **kwargs):


Same comment about why this is needed and dropping all but the first positional argument.

hertschuh · 2025-02-27T21:41:08Z

keras/src/export/export_utils.py

-        tensor_spec = x
-    else:
+        return x
+    if isinstance(x, dict):


make_tf_tensor_spec is always called (and should always be called) within keras.tree.map_structure, so you'll never handles dicts here. Remove this case.

It is indeed not covered, showing it's not used:
https://app.codecov.io/gh/keras-team/keras/pull/20842?src=pr&el=tree&filepath=keras%2Fsrc%2Fexport%2Fexport_utils.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=keras-team#4690cab8cef00d17f2a67990022886ff-R103

hertschuh · 2025-02-27T22:18:55Z

keras/src/models/model.py

@@ -399,6 +413,15 @@ def quantize(self, mode, **kwargs):
    def build_from_config(self, config):
        if not config:
            return
+        # Fetch the input structure from config if available.


Lines 416-424 is all dead code. "input_names" is never in the config because nothing adds it there. Remove.

Coverage confirms it:
https://app.codecov.io/gh/keras-team/keras/pull/20842?src=pr&el=tree&filepath=keras%2Fsrc%2Fmodels%2Fmodel.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=keras-team#b343e88c9bde68f6f2b9d00f13fa225e-R417

hertschuh · 2025-02-27T22:20:54Z

keras/src/models/model.py

@@ -204,6 +213,11 @@ def get_layer(self, name=None, index=None):
                return self.layers[index]

        if name is not None:
+            # Check if the name matches any of the input names.


I don't think this is needed, the fallback should work for this use case.

Also, coverage shows it's never used:
https://app.codecov.io/gh/keras-team/keras/pull/20842?src=pr&el=tree&filepath=keras%2Fsrc%2Fmodels%2Fmodel.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=keras-team#b343e88c9bde68f6f2b9d00f13fa225e-R218

hertschuh · 2025-02-27T22:23:06Z

keras/src/models/model.py

@@ -156,6 +156,15 @@ def __init__(self, *args, **kwargs):
            functional.Functional.__init__(self, *args, **kwargs)
        else:
            Layer.__init__(self, *args, **kwargs)
+            if args:
+                inputs = args[0]


This assumes there is only one input argument. That's true for Functional models, that is not necessarily true for subclass model.

hertschuh · 2025-02-27T22:23:43Z

keras/src/models/model.py

@@ -156,6 +156,15 @@ def __init__(self, *args, **kwargs):
            functional.Functional.__init__(self, *args, **kwargs)
        else:
            Layer.__init__(self, *args, **kwargs)
+            if args:
+                inputs = args[0]
+            # Only set _input_names if not already initialized by Functional.


You're in an else, you know for sure it's not a Functional model if you're here.

hertschuh · 2025-02-27T22:38:57Z

keras/src/export/export_utils.py

-        if isinstance(input_signature, list) and len(input_signature) > 1:
-            input_signature = [input_signature]
+        if hasattr(model, "_input_names") and model._input_names:
+            if isinstance(model._inputs_struct, dict):


Why would _inputs_struct not have the right order?

harshaljanjani · 2025-02-28T01:56:26Z

This is lovely! Thanks for sharing your insights, @hertschuh. I believe I was missing the fact that I could make much better use of the existing _input_struct. I’ll try to understand the problem better based on your comments and make the necessary changes!

fix(export+models): Enhance input signature handling for dictionary-b…

5a84fc5

…ased models

google-ml-butler bot added the size:M label Feb 2, 2025

google-ml-butler bot assigned gbaned Feb 2, 2025

harshaljanjani added 2 commits February 2, 2025 21:06

fix: Fix ExportArchiveTest::test_export_with_dict_input in saved_mode…

d903b41

…l_test.py

fix: ExportArchiveTest, FunctionalTest (most cases) & ModelTest (all …

9e38bd8

…cases) - Improves input structure validation in Model and Functional classes - Adds strict validation with clear error messages for mismatches

harshaljanjani added 2 commits February 3, 2025 14:04

refactor: Fix test cases and input error formatting.

1d969c3

- Make the linter happy :)

fix: Fix spurious change and test cases

299b4ba

- Added shape logging for preprocessed and crossed features - Added debug messages in Functional model processing (To be removed)

fix: Remove debug print statements from modules

0c65c02

harshaljanjani marked this pull request as ready for review February 3, 2025 19:54

jeffcarp reviewed Feb 5, 2025

View reviewed changes

harshaljanjani added 8 commits February 8, 2025 10:19

refactor: Reintroduce debugging prints

52d7fb3

fix: correct test cases for TensorFlow backend compatibility

18adae3

fix: correct failing two test cases in ExportArchiveTest

5178e44

refactor: Make linter happy :)

5de08b4

TODO: Inference support (Tensorflow) and support for JAX

b74a295

fix: Inference support for Tensorflow backend

4ed5607

fix: Inference support for JAX backend

0659804

end of sprint: Remove debugging prints

c32c02c

harshaljanjani requested a review from jeffcarp February 8, 2025 17:10

google-ml-butler bot added the awaiting review label Feb 8, 2025

harshaljanjani changed the title ~~fix(export+models): Enhance input signature handling for dictionary-based models~~ fix(export+models): Enhance support for dictionary-based model input signatures in TensorFlow and JAX Feb 9, 2025

Merge branch 'master' into model-export-input-signature-fix

9e7f3c8

jeffcarp approved these changes Feb 18, 2025

View reviewed changes

keras/src/models/functional.py Outdated Show resolved Hide resolved

keras/src/models/functional.py Outdated Show resolved Hide resolved

google-ml-butler bot added kokoro:force-run ready to pull Ready to be merged into the codebase labels Feb 18, 2025

kokoro-team removed the kokoro:force-run label Feb 18, 2025

fix: Fix inconsistent spacing and formatting

72371d7

google-ml-butler bot removed the ready to pull Ready to be merged into the codebase label Feb 19, 2025

jeffcarp approved these changes Feb 19, 2025

View reviewed changes

google-ml-butler bot added kokoro:force-run ready to pull Ready to be merged into the codebase labels Feb 19, 2025

kokoro-team removed the kokoro:force-run label Feb 19, 2025

hertschuh requested changes Feb 27, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

fix(export+models): Enhance support for dictionary-based model input signatures in TensorFlow and JAX #20842

fix(export+models): Enhance support for dictionary-based model input signatures in TensorFlow and JAX #20842

harshaljanjani commented Feb 2, 2025 •

edited

Loading

codecov-commenter commented Feb 3, 2025 •

edited

Loading

harshaljanjani commented Feb 3, 2025 •

edited

Loading

harshaljanjani commented Feb 4, 2025

jeffcarp left a comment

harshaljanjani commented Feb 6, 2025

harshaljanjani commented Feb 8, 2025

harshaljanjani commented Feb 17, 2025

fchollet commented Feb 17, 2025

harshaljanjani commented Feb 17, 2025

jeffcarp left a comment

harshaljanjani commented Feb 19, 2025

jeffcarp left a comment

harshaljanjani commented Feb 22, 2025 •

edited

Loading

hertschuh left a comment

hertschuh Feb 27, 2025

hertschuh Feb 27, 2025

hertschuh Feb 27, 2025

hertschuh Feb 27, 2025

hertschuh Feb 27, 2025

hertschuh Feb 27, 2025

hertschuh Feb 27, 2025

hertschuh Feb 27, 2025

hertschuh Feb 27, 2025

hertschuh Feb 27, 2025

harshaljanjani commented Feb 28, 2025

fix(export+models): Enhance support for dictionary-based model input signatures in TensorFlow and JAX #20842

Are you sure you want to change the base?

fix(export+models): Enhance support for dictionary-based model input signatures in TensorFlow and JAX #20842

Conversation

harshaljanjani commented Feb 2, 2025 • edited Loading

codecov-commenter commented Feb 3, 2025 • edited Loading

Codecov Report

harshaljanjani commented Feb 3, 2025 • edited Loading

harshaljanjani commented Feb 4, 2025

jeffcarp left a comment

Choose a reason for hiding this comment

harshaljanjani commented Feb 6, 2025

harshaljanjani commented Feb 8, 2025

harshaljanjani commented Feb 17, 2025

fchollet commented Feb 17, 2025

harshaljanjani commented Feb 17, 2025

jeffcarp left a comment

Choose a reason for hiding this comment

harshaljanjani commented Feb 19, 2025

jeffcarp left a comment

Choose a reason for hiding this comment

harshaljanjani commented Feb 22, 2025 • edited Loading

hertschuh left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

harshaljanjani commented Feb 28, 2025

harshaljanjani commented Feb 2, 2025 •

edited

Loading

codecov-commenter commented Feb 3, 2025 •

edited

Loading

harshaljanjani commented Feb 3, 2025 •

edited

Loading

harshaljanjani commented Feb 22, 2025 •

edited

Loading