Store defaults of optional arguments in auto-config

[vpratz]

Dear Keras team,

I'm a co-developer of the BayesFlow library. We are currently facing the challenge that we want to be able to adapt the defaults of our networks when we find better configurations, but do so without breaking model loading for models that were saved to disk with the old defaults. We currently use the auto-config functionality to keep our code simple, and ideally we would like to continue to do so. Our problem is the following:

The current auto-config implementation only stores the supplied arguments, but not the values of the optional arguments of the constructor. As a consequence, loading a model fails when the default value of an optional argument has changed, even when the rest of the code of the model remains unchanged.

This PR extends the auto-config functionality to also store default values, while otherwise being (as far as I can tell) compatible with the previous implementation. This ensures that the "old defaults" are serialized, so changing defaults is no longer a problem. Is this a change you would be willing to include, or do you see any downsides with this approach?

See also bayesflow-org/bayesflow#566.

P.S.: A minor thing I encountered when looking at the current implementation is that variable arguments cannot be handled properly, as they cannot be properly converted to a dictionary. Currently, an argument *args is ignored without a warning, which might lead users to thinking it is handled automatically.

Edit: correct description of how *args is handled by the current implementation.

[gemini-code-assist]

Summary of Changes

Hello @vpratz, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request enhances the auto-configuration mechanism in Keras to include the default values of optional arguments when generating a model's configuration. This change addresses a critical issue where models saved with older default argument values would fail to load if those defaults were subsequently modified in the code, ensuring backward compatibility and more robust model serialization.

Highlights

• Enhanced Auto-Configuration: The auto-config functionality is extended to capture and store the default values of optional arguments during object instantiation.
• Improved Model Loading Robustness: This prevents breakage when loading models that were saved with previous default configurations, even if the code's default values have since changed.
• Leverages inspect.signature: The implementation now uses inspect.signature and bound_parameters.apply_defaults() to accurately determine and include default argument values in the generated configuration.
• Handles Variable Arguments: Specific logic is added to manage *args (varargs) for backward compatibility, storing only the first argument if present, or removing the entry if empty.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request improves the auto-config functionality by storing default values for optional arguments, which is a great enhancement for model serialization and backward compatibility. The use of inspect.signature is a good move towards a more robust argument handling.

I've found a critical issue in the handling of variable keyword arguments (**kwargs) that could lead to a crash, and I've also suggested an improvement for handling variable positional arguments (*args) to prevent silent data loss. Please see my detailed comments below.

[gemini-code-assist]

If cls.__init__ does not have a variable keyword argument (**kwargs), argspec.varkw will be None. In this case, kwargs.pop(None, {}) will raise a KeyError, causing a crash. You should guard this operation with a check to ensure argspec.varkw is not None.

Suggested change

	kwargs |= kwargs.pop(argspec.varkw, {})
	if argspec.varkw is not None:
	kwargs |= kwargs.pop(argspec.varkw, {})

[vpratz]

.pop does not raise a KeyError when a default is provided

[vpratz]

@divyashreepathihalli Guarding against None is not strictly necessary here, as all keys in kwargs are strings, so kwargs.pop(None, {}) will just return {}. Would you prefer to have it in there anyway to improve clarity?

[codecov-commenter]

Codecov Report

❌ Patch coverage is 86.36364% with 3 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.49%. Comparing base (4415fcc) to head (13c4d40).

Files with missing lines	Patch %	Lines
keras/src/ops/operation.py	86.36%	2 Missing and 1 partial ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##           master   #21615   +/-   ##
=======================================
  Coverage   82.49%   82.49%           
=======================================
  Files         572      572           
  Lines       57451    57470   +19     
  Branches     8982     8987    +5     
=======================================
+ Hits        47395    47411   +16     
- Misses       7760     7762    +2     
- Partials     2296     2297    +1     

Flag	Coverage Δ
keras	82.30% <86.36%> (+<0.01%)	⬆️
keras-jax	63.52% <86.36%> (+<0.01%)	⬆️
keras-numpy	57.85% <86.36%> (+<0.01%)	⬆️
keras-openvino	34.35% <86.36%> (+0.01%)	⬆️
keras-tensorflow	64.19% <86.36%> (+<0.01%)	⬆️
keras-torch	63.76% <86.36%> (+<0.01%)	⬆️

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.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[fchollet]

The try/except is very broad, that seems dangerous?

[fchollet]

Thanks for the PR.

We are currently facing the challenge that we want to be able to adapt the defaults of our networks when we find better configurations, but do so without breaking model loading for models that were saved to disk with the old defaults

In general the default get_config implementation is very basic (not long ago it would just have raised NotImplemented...) so my top recommendation would be to implement get_config on your custom layers/models. That is the only saving/loading mechanism that will be reliable. You can also bake backwards compatibility mechanisms at that level if you know you're going to change, say, argument names or something.

[vpratz]

Thanks for the review and your assessment. I agree that this approach would give us the most control and reliability (given that we don't make errors when manually specifying the config). Our concern is that it is harder to maintain (e.g., if someone introduces a new argument, but forgets to update the get_config, leading to subtle deserialization bugs), and might introduce a barrier for external contributors. I will discuss with the others how we want to proceed.

Independent of how we decide there, would you be in favor of updating the auto-config functionality, or would you prefer to leave it as is for now? The two main things I see are:

• Varargs are silently ignored (I previously misinterpreted the code and thought the first value is used, this is not the case). This could be done by raising a NotImplementedError for that case, or by properly supporting it, which would require adapting __new__ and from_config.
• Default values are not stored, which might be useful for the reasons outlined in my initial post.

Let me know what you think. If you think updating this would be useful, I will adapt the code in the PR. If not, feel free to close the PR.

[fchollet]

Sounds good, we can add this feature. Please add a unit test, to make sure we don't accidentally break it in the future.

[divyashreepathihalli]

/gemini review

[gemini-code-assist]

Code Review

This pull request is a great initiative to improve the auto-config functionality by including default values of optional arguments, which will certainly enhance model serialization and backward compatibility. The approach of using inspect.signature is robust and modern. I have identified a few areas for improvement concerning error handling and edge cases in function signatures, which I've detailed in my comments. Addressing these points will make the implementation even more solid.

[vpratz]

Thanks for the response. After some consideration, I think adding support for variadic positional arguments is not worth the effort, as it breaks up the simple notion of the config being a dictionary that can be passed via cls(**config). Also, being positional only, the same argument name may be used inside **kwargs as well, leading to possible collisions if we would expand kwargs to store all arguments in the same flat dictionary.

Therefore, I have implemented the following changes:

• add default values of arguments to the config
• raise a meaningful exception in get_config for signatures that do not work with the default from_config. Those are signatures that contain variadic positional args, and, for completeness, signatures that take positional only arguments. Both cannot be passed via cls(**config).
• minor fix: the error message contained single curly brackets instead of double curly brackets in the format string, leading to an error.

Tests:

• add a optional parameter with a default value to the signatures of the test operators, check its presence in the config
• add tests for the two failure cases: non-serializable arguments and variadic positional arguments

I have also tested the changes in the BayesFlow test suite and did not encounter any problems.
Let me know what you think, and feel free to make adaptations if necessary so the code fits the coding style of the project.

[fchollet]

Wouldn't loss of support for positional-only arguments be a regression?

[vpratz]

Thanks for taking a look. No, positional-only arguments are also not supported with the current implementation, as calling cls(**config) will pass the items of config as keyword arguments, which is not allowed for positional-only arguments. This will raise an error <name>.__init__() got some positional-only arguments passed as keyword arguments. As far as I can tell, positional-only arguments are very rare, I have never really encountered them in the wild.

This would be the corresponding failing test on master in keras/src/ops/operation_test.py if you want to try it out:

class OpWithCustomConstructorNoNamePositionalOnlyArgs(operation.Operation):
    def __init__(self, alpha, /):
        super().__init__()
        self.alpha = alpha

    def call(self, x):
        return self.alpha * x

    def compute_output_spec(self, x):
        return keras_tensor.KerasTensor(x.shape, x.dtype)


class OperationTest(testing.TestCase):
    # [...]
    
    def test_serialization_custom_constructor_with_no_name_pos_only_auto_config(
        self,
    ):
        # Auto generated name is not serialized.
        op = OpWithCustomConstructorNoNamePositionalOnlyArgs(0.2)
        config = op.get_config()
        self.assertEqual(config, {"alpha": 0.2})
        revived = OpWithCustomConstructorNoNamePositionalOnlyArgs.from_config(
            config
        )
        self.assertEqual(revived.get_config(), config)

which gives the output (on master)

FAILED keras/src/ops/operation_test.py::OperationTest::test_serialization_custom_constructor_with_no_name_pos_only_auto_config - 
TypeError: Error when deserializing class 'OpWithCustomConstructorNoNamePositionalOnlyArgs'
using config={'alpha': 0.2}.

Exception encountered: OpWithCustomConstructorNoNamePositionalOnlyArgs.__init__()
got some positional-only arguments passed as keyword arguments: 'alpha'