Fix: move additional metrics from approximator to networks

[vpratz]

Supplying the additional metrics for inference and summary networks via the approximators compile method caused problems during deseralization (#497). This can be resolved nicely by moving the metrics directly to the networks' constructors, analogous to how Keras normally handles custom metrics in layers.

As summary networks and inference networks inherit from the respective base classes, this change only requires minor adaptations. Calls to layer_kwargs are now only used in classes that directly inherit from keras.Layer, and have been moved into InferenceNetwork and SummaryNetwork.

Fixes #497.

[codecov]

Codecov Report

Attention: Patch coverage is 91.56627% with 7 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
bayesflow/utils/serialization.py	92.30%	3 Missing ⚠️
bayesflow/approximators/continuous_approximator.py	60.00%	2 Missing ⚠️
...low/approximators/model_comparison_approximator.py	60.00%	2 Missing ⚠️

Files with missing lines	Coverage Δ
...ow/experimental/diffusion_model/diffusion_model.py	79.14% <100.00%> (ø)
...flow/experimental/free_form_flow/free_form_flow.py	92.59% <ø> (-0.07%)	⬇️
...w/networks/consistency_models/consistency_model.py	97.65% <100.00%> (ø)
bayesflow/networks/coupling_flow/coupling_flow.py	100.00% <ø> (ø)
bayesflow/networks/flow_matching/flow_matching.py	92.79% <ø> (-0.07%)	⬇️
bayesflow/networks/inference_network.py	79.24% <100.00%> (+5.43%)	⬆️
bayesflow/networks/mlp/mlp.py	90.90% <100.00%> (+0.16%)	⬆️
bayesflow/networks/point_inference_network.py	94.84% <100.00%> (+0.10%)	⬆️
bayesflow/networks/summary_network.py	97.56% <100.00%> (+3.27%)	⬆️
bayesflow/approximators/continuous_approximator.py	90.62% <60.00%> (+0.22%)	⬆️
... and 2 more

... and 3 files with indirect coverage changes

[vpratz]

@LarsKue @stefanradev93 @han-ol I encountered the problem that to pass metrics to the networks in their constructors, we would have to specify all get_config functions manually, as the metrics are not basic types.
I have now instead adapted the auto-config capabilities of keras.Layer to accept any object we can (de)serialize, and to use our serialization functions. This should make the auto-config functionality much more flexible. Take a look at the BaseLayer class for the implementation. It seems to me that this doesn't cause problems for now (hoping that all tests will pass).
Do you see any downsides with this approach? Or do you know something about the motivation of Keras to limit this functionality to basic types?

[stefanradev93]

I would let @LarsKue chip in on this, as I am bit concerned about having more and more "custom" variants of basic keras containers (e.g., Sequential, Layer,...).

[vpratz]

Thanks for the comment, I can understand this, relying too much on hacking/modifying Keras internals has the downside that our code might become more fragile with respect to changes in Keras. I think the ability to pass non-basic types to our inference and summary network base classes would be nice to have, the behavior in this PR is one example of this.
The options that I see are:

1. accept that we cannot pass non-basic arguments to the base classes, making implementation of some features more cumbersome (or impossible)
2. forego the auto-config from offered by keras.Layer and explicitly implement get_config in all our networks, which works but requires some redundant work and can be error-prone
3. a mechanism like the one implemented here, which offers more flexibility, but relies on Keras making not too drastic changes to the keras.Layer internals, which we cannot control

[LarsKue]

I would be on board if it was a lightweight wrapper, but it seems that there is a lot of copied code from keras, which we should avoid in general, imo.

@vpratz Can we solve this through monkey-patching somehow, like with the serialize and deserialize functions?

[stefanradev93]

If this is not breaking in any downstream task and only adds to the existing functionality, we may consider doing a PR to keras or asking them to implement it (or why they decided not to)?

[vpratz]

@LarsKue Thanks for taking a look. I'll take a look at alternative ways to achieve this behavior.

[vpratz]

@LarsKue Thanks a lot, this was a really good pointer. I have now implemented a wrapper around the constructor that is applied inside the serializable decorator. In addition, the default from_config is replaced with a variant that deserializes the config again.

There is still some copied code, but I think it is now limited to the part that is required for the feature to behave as similar to the Keras implementation as possible, apart from the desired changes (we might not need the part regarding the dtype, but I'm not sure and would leave it in there for now).

With those changes, we could remove the from_config method from most classes. Is this a change I should make?

@stefanradev93 @LarsKue Please take another look and let me know what you think (the most relevant part is in bayesflow/utils/serialization.py

[LarsKue]

Thanks, this looks much better. Since this is a sensitive change, I think we should extensively test it before rolling it out. Otherwise, green light from my side!

[vpratz]

I forgot to make the same changes in the ModelComparisonApproximator, this is now fixed. I have changed the output that is supplied to its classifier metrics from predictions to probabilities. This way, it is compatible with the metrics offered by Keras.

[paul-buerkner]

@vpratz and @LarsKue what is the status of this PR? Is this ready to be merged?

[vpratz]

@paul-buerkner I would like to get a review from @stefanradev93 before we merge this, as the changes affect multiple parts of the code. Also, I would wait until we have decided how we proceed on #525...

[vpratz]

I see some weird behavior due to the inheritance, I will have to give this another look...

[vpratz]

Ok, I think it works. The general mechanism for getting the config of a subclass MyLayer of keras.Layer is the following. For an instance layer = MyLayer(), calling layer.get_config is resolved in the "standard" Python fashion, so there are two cases:

a) The subclass overrides get_config

• MyLayer.get_config is called
• If super().get_config is called, keras.Layer.get_config provides trainable and dtype, and calls keras.Operation.get_config
• keras.Operation.get_config provides name, detects that get_config was overridden, and does not apply the auto config

b) The subclass does not override get_config, which triggers the auto-config behavior

• keras.Layer.get_config is called, which is decorated with the @keras.src.utils.python_utils.default decorator. Again, it provides trainable and dtype, and calls keras.Operation.get_config
• keras.Operation.get_config detects that get_config was not overridden by using keras.src.utils.python_utils.is_default. It reads the config from self._auto_config, adds name to it, and returns it.

The default and is_default are two simple functions, which set/read the _is_default flag on the method:

def default(method):
    """Decorates a method to detect overrides in subclasses."""
    method._is_default = True
    return method


def is_default(method):
    """Check if a method is decorated with the `default` wrapper."""
    return getattr(method, "_is_default", False)

To adapt this setup for us, we want our base classes (InferenceNetwork, SummaryNetwork, PointInferenceNetwork) to do the same thing as keras.Layer: Provide a few entries for the config, and have the auto-config used if the sub-class does not override get_config.

The easiest way to achieve this is a get_config function decorated with @python_utils.default. Note that in the auto-config case, the subclass-config is provided by the last get_config call, so we cannot filter the base_config dictionaries in this setup.

[vpratz]

@stefanradev93 @LarsKue I have fixed the problem I found (see above) and cleaned up the code a bit. Could you please take another look?