Stabilize container weight paths across compatible layer subclasses

keras/src/layers/rnn/lstm_test.py

@@ -1,13 +1,48 @@
+import os
+
 import numpy as np
 import pytest
 from absl.testing import parameterized
 
 from keras.src import initializers
 from keras.src import layers
+from keras.src import models
 from keras.src import testing
 
 
 class LSTMTest(testing.TestCase):
+    def test_load_weights_from_custom_lstm_into_base_lstm(self):
+        class MyCustomLSTM(layers.LSTM):
+            pass
+
+        def make_model(use_custom_lstm):
+            inputs = layers.Input(shape=(None, 4), name="inputs")
+            lstm_cls = MyCustomLSTM if use_custom_lstm else layers.LSTM
+            outputs = lstm_cls(
+                units=8,
+                return_sequences=True,
+                name="my_lstm",
+            )(inputs)
+            return models.Model(inputs=inputs, outputs=outputs)
+
+        x = np.random.random((2, 3, 4)).astype("float32")
+        weights_path = os.path.join(
+            self.get_temp_dir(), "custom_lstm.weights.h5"
+        )
+
+        source_model = make_model(use_custom_lstm=True)
+        source_model(x)
+        source_model.save_weights(weights_path)
+
+        dest_model = make_model(use_custom_lstm=False)
+        dest_model(x)
+        dest_model.load_weights(weights_path)
+
+        for source_var, dest_var in zip(
+            source_model.weights, dest_model.weights
+        ):
+            self.assertAllClose(source_var, dest_var)
+
     def test_basics(self):
         self.run_layer_test(
             layers.LSTM,

keras/src/saving/saving_lib.py

@@ -916,6 +916,62 @@ def _get_unique_name(name, used_names):
     return name
 
 
+def _container_saveable_base_name(saveable):
+    # For user-defined subclasses, use the first Keras class in the MRO
+    # so container paths stay compatible with the built-in class hierarchy.
+    for cls in saveable.__class__.__mro__:
+        if cls is object:
+            continue
+        module = cls.__module__
+        if not module.startswith("keras.src"):
+            continue
+        # Classes declared in test modules are custom subclasses in this
+        # repository context; skip them and use the nearest built-in parent.
+        if module.endswith("_test") or ".tests." in module:
+            continue
+        if cls.__name__ in _GENERIC_SAVEABLE_BASE_CLASSES:
+            break
+        return naming.to_snake_case(cls.__name__)
+    return naming.to_snake_case(saveable.__class__.__name__)
+
+
+def _container_saveable_name(saveable, used_names):
+    base_name = _container_saveable_base_name(saveable)
+    return _get_unique_name(base_name, used_names)
+
+
+_GENERIC_SAVEABLE_BASE_CLASSES = {
+    "Layer",
+    "Model",
+    "Loss",
+    "Metric",
+    "Optimizer",
+    "Operation",
+}
+
+
+def _container_alias_names(saveable, suffix):
+    alias_roots = []
+    for cls in saveable.__class__.__mro__:
+        if cls is object:
+            continue
+        if not cls.__module__.startswith("keras.src"):
+            continue
+        if cls.__name__ in _GENERIC_SAVEABLE_BASE_CLASSES:
+            break
+        root = naming.to_snake_case(cls.__name__)
+        if root not in alias_roots:
+            alias_roots.append(root)
+
+    aliases = []
+    for root in alias_roots:
+        if suffix > 0:
+            aliases.append(f"{root}_{suffix}")
+        else:
+            aliases.append(root)
+    return aliases
+
+
 def _container_path_present(weights_store, assets_store, path):
     """True if either store has anything under `path`.
 
@@ -950,13 +1006,7 @@ def _save_container_state(
 
     for saveable in container:
         if isinstance(saveable, KerasSaveable):
-            # Do NOT address the saveable via `saveable.name`, since
-            # names are usually autogenerated and thus not reproducible
-            # (i.e. they may vary across two instances of the same model).
-            name = _get_unique_name(
-                naming.to_snake_case(saveable.__class__.__name__),
-                used_names,
-            )
+            name = _container_saveable_name(saveable, used_names)
             _save_state(
                 saveable,
                 weights_store,
@@ -1001,20 +1051,54 @@ def _load_container_state(
 
     for saveable in container:
         if isinstance(saveable, KerasSaveable):
-            name = _get_unique_name(
-                naming.to_snake_case(saveable.__class__.__name__),
-                used_names,
-            )
-            _load_state(
-                saveable,
-                weights_store,
-                assets_store,
-                inner_path=file_utils.join(inner_path, name).replace("\\", "/"),
-                skip_mismatch=skip_mismatch,
-                visited_saveables=visited_saveables,
-                failed_saveables=failed_saveables,
-                error_msgs=error_msgs,
-            )
+            name = _container_saveable_name(saveable, used_names)
+            name_suffix = used_names[_container_saveable_base_name(saveable)]
+            candidate_names = [name]
+            for alias in _container_alias_names(saveable, name_suffix):
+                if alias not in candidate_names:
+                    candidate_names.append(alias)
+            attempted_candidate = False
+            for candidate in candidate_names:
+                candidate_path = file_utils.join(inner_path, candidate).replace(
+                    "\\", "/"
+                )
+                if not _container_path_present(
+                    weights_store, assets_store, candidate_path
+                ):
+                    continue
+
+                attempted_candidate = True
+                _load_state(
+                    saveable,
+                    weights_store,
+                    assets_store,
+                    inner_path=candidate_path,
+                    skip_mismatch=skip_mismatch,
+                    visited_saveables=visited_saveables,
+                    failed_saveables=failed_saveables,
+                    error_msgs=error_msgs,
+                )
+                # If this candidate loaded cleanly (or if failures aren't being
+                # tracked), keep it and stop trying aliases.
+                if (
+                    failed_saveables is None
+                    or id(saveable) not in failed_saveables
+                ):
+                    break
+
+            if not attempted_candidate:
+                _load_state(
+                    saveable,
+                    weights_store,
+                    assets_store,
+                    inner_path=file_utils.join(inner_path, name).replace(
+                        "\\", "/"
+                    ),
+                    skip_mismatch=skip_mismatch,
+                    visited_saveables=visited_saveables,
+                    failed_saveables=failed_saveables,
+                    error_msgs=error_msgs,
+                )
         elif isinstance(saveable, (list, dict, tuple)):
             name = _get_unique_name("container", used_names)
             nested_path = file_utils.join(inner_path, name).replace("\\", "/")