catch optree exception when changing backends

Author: t-kalinowski

keras/src/tree/optree_impl.py

@@ -13,26 +13,32 @@ def register_tree_node_class(cls):
     from tensorflow.python.trackable.data_structures import ListWrapper
     from tensorflow.python.trackable.data_structures import _DictWrapper
 
-    optree.register_pytree_node(
-        ListWrapper,
-        lambda x: (x, None),
-        lambda metadata, children: ListWrapper(list(children)),
-        namespace="keras",
-    )
+    try:
+        optree.register_pytree_node(
+            ListWrapper,
+            lambda x: (x, None),
+            lambda metadata, children: ListWrapper(list(children)),
+            namespace="keras",
+        )
+
+        def sorted_keys_and_values(d):
+            keys = sorted(list(d.keys()))
+            values = [d[k] for k in keys]
+            return values, keys, keys
+
+        optree.register_pytree_node(
+            _DictWrapper,
+            sorted_keys_and_values,
+            lambda metadata, children: _DictWrapper(
+                {key: child for key, child in zip(metadata, children)}
+            ),
+            namespace="keras",
+        )
+    except ValueError:
+        # optree raises a ValueError if the class is already registered.
+        # Triggered if config.set_backend() is called multiple times.
+        pass
 
-    def sorted_keys_and_values(d):
-        keys = sorted(list(d.keys()))
-        values = [d[k] for k in keys]
-        return values, keys, keys
-
-    optree.register_pytree_node(
-        _DictWrapper,
-        sorted_keys_and_values,
-        lambda metadata, children: _DictWrapper(
-            {key: child for key, child in zip(metadata, children)}
-        ),
-        namespace="keras",
-    )
 
 
 def is_nested(structure):

keras/src/utils/tracking.py

@@ -27,6 +27,15 @@ def wrapper(*args, **kwargs):
 
     return wrapper
 
+def safe_register_tree_node_class(cls):
+    try:
+        return tree.register_tree_node_class(cls)
+    except ValueError:
+        # optree raises a ValueError if the class is already registered.
+        # Triggered if config.set_backend() is called multiple times.
+        return cls
+
+
 
 class Tracker:
     """Attribute tracker, used for e.g. Variable tracking.
@@ -133,7 +142,7 @@ def replace_tracked_value(self, store_name, old_value, new_value):
         self.stored_ids[store_name].add(id(new_value))
 
 
-@tree.register_tree_node_class
+@safe_register_tree_node_class
 class TrackedList(list):
     def __init__(self, values=None, tracker=None):
         self.tracker = tracker
@@ -194,7 +203,7 @@ def tree_unflatten(cls, metadata, children):
         return cls(children)
 
 
-@tree.register_tree_node_class
+@safe_register_tree_node_class
 class TrackedDict(dict):
     def __init__(self, values=None, tracker=None):
         self.tracker = tracker
@@ -245,7 +254,7 @@ def tree_unflatten(cls, keys, values):
         return cls(zip(keys, values))
 
 
-@tree.register_tree_node_class
+@safe_register_tree_node_class
 class TrackedSet(set):
     def __init__(self, values=None, tracker=None):
         self.tracker = tracker