Saving a Keras model that includes a mixture layer (e.g. MixtureNormal) might be unsafe

Hi Team,

I am trying to save the weights of a Keras model that end with a MixtureNormal layer. And I get the following warning:

.../keras/src/saving/serialization_lib.py:(line) 159: UserWarning:The object being serialized includes a `lambda`. This is unsafe. In order to reload the object, you will have to pass `safe_mode=False` to the loading function. Please avoid using `lambda` in the future, and use named Python functions instead. This is the `lambda` being serialized:         lambda t: MixtureNormal.new(  # pylint: disable=g-long-lambda
            t, num_components, event_shape, validate_args),

I know that tfp mixture layers have a strong commitment to lambda functions so you can pass tfp distributions. Is this something you plan to work on?

Start with keras/src/saving/serialization_lib.py around line 159, then trace how the MixtureNormal layer is serialized and reloaded. Check the interaction with safe_mode=False and determine whether the reported warning can be avoided or needs explicit documentation; done means the save/reload behavior is resolved or clearly supported.