Fix code formatting and remove unused variable

- Remove unused 'result' variable in _reconstruct_state_tree_with_values
- Fix long comment line in test file
- Apply code formatting changes

Author: amitsrivastava78

keras/src/backend/jax/__init__.py

@@ -24,8 +24,8 @@
 from keras.src.backend.jax.core import shape
 from keras.src.backend.jax.core import stop_gradient
 from keras.src.backend.jax.core import vectorized_map
+from keras.src.backend.jax.distribution_lib import process_id
 from keras.src.backend.jax.rnn import cudnn_ok
 from keras.src.backend.jax.rnn import gru
 from keras.src.backend.jax.rnn import lstm
 from keras.src.backend.jax.rnn import rnn
-from keras.src.backend.jax.distribution_lib import process_id

keras/src/backend/numpy/__init__.py

@@ -20,8 +20,8 @@
 from keras.src.backend.numpy.core import random_seed_dtype
 from keras.src.backend.numpy.core import shape
 from keras.src.backend.numpy.core import vectorized_map
+from keras.src.backend.numpy.distribution_lib import process_id
 from keras.src.backend.numpy.rnn import cudnn_ok
 from keras.src.backend.numpy.rnn import gru
 from keras.src.backend.numpy.rnn import lstm
 from keras.src.backend.numpy.rnn import rnn
-from keras.src.backend.numpy.distribution_lib import process_id

keras/src/backend/numpy/distribution_lib.py

@@ -3,4 +3,4 @@
 
 def process_id():
     """Return the current process ID for the distribution setting."""
-    return 0
+    return 0

keras/src/backend/openvino/__init__.py

@@ -20,8 +20,8 @@
 from keras.src.backend.openvino.core import random_seed_dtype
 from keras.src.backend.openvino.core import shape
 from keras.src.backend.openvino.core import vectorized_map
+from keras.src.backend.openvino.distribution_lib import process_id
 from keras.src.backend.openvino.rnn import cudnn_ok
 from keras.src.backend.openvino.rnn import gru
 from keras.src.backend.openvino.rnn import lstm
 from keras.src.backend.openvino.rnn import rnn
-from keras.src.backend.openvino.distribution_lib import process_id

keras/src/backend/openvino/distribution_lib.py

@@ -3,4 +3,4 @@
 
 def process_id():
     """Return the current process ID for the distribution setting."""
-    return 0
+    return 0

keras/src/backend/tensorflow/__init__.py

@@ -23,8 +23,8 @@
 from keras.src.backend.tensorflow.core import shape
 from keras.src.backend.tensorflow.core import stop_gradient
 from keras.src.backend.tensorflow.core import vectorized_map
+from keras.src.backend.tensorflow.distribution_lib import process_id
 from keras.src.backend.tensorflow.rnn import cudnn_ok
 from keras.src.backend.tensorflow.rnn import gru
 from keras.src.backend.tensorflow.rnn import lstm
 from keras.src.backend.tensorflow.rnn import rnn
-from keras.src.backend.tensorflow.distribution_lib import process_id

keras/src/backend/torch/__init__.py

@@ -39,8 +39,8 @@
 from keras.src.backend.torch.core import stop_gradient
 from keras.src.backend.torch.core import to_torch_dtype
 from keras.src.backend.torch.core import vectorized_map
+from keras.src.backend.torch.distribution_lib import process_id
 from keras.src.backend.torch.rnn import cudnn_ok
 from keras.src.backend.torch.rnn import gru
 from keras.src.backend.torch.rnn import lstm
 from keras.src.backend.torch.rnn import rnn
-from keras.src.backend.torch.distribution_lib import process_id

keras/src/backend/torch/distribution_lib.py

@@ -10,4 +10,4 @@ def process_id():
             return dist.get_rank()
         return 0
     except (ImportError, AttributeError):
-        return 0
+        return 0

keras/src/callbacks/orbax_checkpoint.py

@@ -41,7 +41,7 @@
 def _get_state_tree(model):
     """Get the complete model state as a nested tree structure."""
     state_tree = model.get_state_tree(value_format="numpy_array")
-    
+
     # Convert numpy scalar types to Python types for Orbax compatibility
     def convert_scalars(obj):
         if isinstance(obj, np.ndarray) and obj.ndim == 0:
@@ -54,29 +54,30 @@ def convert_scalars(obj):
             return {k: convert_scalars(v) for k, v in obj.items()}
         else:
             return obj
-    
+
     return convert_scalars(state_tree)
 
 
 def _flatten_state_tree_values(state_tree):
     """Flatten nested state tree into a list of values in consistent order."""
     values = []
+
     def _flatten(obj):
         if isinstance(obj, dict):
             for key in sorted(obj.keys()):  # Sort for consistent ordering
                 _flatten(obj[key])
         else:
             # Save any non-dict value (numpy arrays, lists, scalars, etc.)
             values.append(obj)
+
     _flatten(state_tree)
     return values
 
 
 def _reconstruct_state_tree_with_values(structure, values):
     """Reconstruct state tree structure with provided values."""
-    result = {}
     value_iter = iter(values)
-    
+
     def _reconstruct(obj):
         if isinstance(obj, dict):
             new_dict = {}
@@ -109,7 +110,7 @@ def _reconstruct(obj):
                     return np.array(value, dtype=obj.dtype).reshape(obj.shape)
             else:
                 return value
-    
+
     return _reconstruct(structure)
 
 
@@ -128,15 +129,10 @@ def _restore_legacy_format(
             target_model.weights[i].assign(weight_tensor)
 
     # Restore optimizer state if available
-    if (
-        "optimizer_state" in checkpoint_data
-        and save_optimizer_state
-    ):
+    if "optimizer_state" in checkpoint_data and save_optimizer_state:
         optimizer_vars_np = checkpoint_data["optimizer_state"]
         # Only restore if the variable counts match
-        if len(optimizer_vars_np) == len(
-            target_model.optimizer.variables
-        ):
+        if len(optimizer_vars_np) == len(target_model.optimizer.variables):
             # Convert NumPy arrays back to backend tensors and assign to
             # optimizer
             for i, var_np in enumerate(optimizer_vars_np):
@@ -406,14 +402,14 @@ def _save_checkpoint(self, step, logs=None):
         trainable_values = _flatten_state_tree_values(
             state_tree["trainable_variables"]
         )
-        
+
         # Save optimizer and metrics state if requested
         optimizer_values = None
         if self.save_optimizer_state and "optimizer_variables" in state_tree:
             optimizer_values = _flatten_state_tree_values(
                 state_tree["optimizer_variables"]
             )
-            
+
         metrics_values = None
         if self.save_metrics_state and "metrics_variables" in state_tree:
             metrics_values = _flatten_state_tree_values(
@@ -423,7 +419,7 @@ def _save_checkpoint(self, step, logs=None):
         composite_state = {
             "model_weights": trainable_values,
         }
-        
+
         if optimizer_values is not None:
             composite_state["optimizer_state"] = optimizer_values
         if metrics_values is not None:
@@ -611,8 +607,9 @@ def _restore_model_state(self, checkpoint_data, model=None):
         target_model = model if model is not None else self.model
 
         # Check if this is the new flattened format
-        if ("model_weights" in checkpoint_data and
-            isinstance(checkpoint_data["model_weights"], list)):
+        if "model_weights" in checkpoint_data and isinstance(
+            checkpoint_data["model_weights"], list
+        ):
             # New format: flattened values
             return self._restore_from_flattened_values(
                 checkpoint_data, target_model
@@ -625,8 +622,10 @@ def _restore_model_state(self, checkpoint_data, model=None):
         else:
             # Fallback to legacy format
             _restore_legacy_format(
-                checkpoint_data, target_model, self.save_optimizer_state,
-                self.save_metrics_state
+                checkpoint_data,
+                target_model,
+                self.save_optimizer_state,
+                self.save_metrics_state,
             )
             return True
 
@@ -649,9 +648,9 @@ def _restore_from_flattened_values(self, checkpoint_data, target_model):
         # Restore trainable variables
         if "model_weights" in checkpoint_data:
             saved_trainable_values = checkpoint_data["model_weights"]
-            target_trainable_structure = (
-                target_state_tree["trainable_variables"]
-            )
+            target_trainable_structure = target_state_tree[
+                "trainable_variables"
+            ]
             reconstructed_state["trainable_variables"] = (
                 _reconstruct_state_tree_with_values(
                     target_trainable_structure, saved_trainable_values
@@ -665,9 +664,9 @@ def _restore_from_flattened_values(self, checkpoint_data, target_model):
             and "optimizer_variables" in target_state_tree
         ):
             saved_optimizer_values = checkpoint_data["optimizer_state"]
-            target_optimizer_structure = (
-                target_state_tree["optimizer_variables"]
-            )
+            target_optimizer_structure = target_state_tree[
+                "optimizer_variables"
+            ]
             reconstructed_state["optimizer_variables"] = (
                 _reconstruct_state_tree_with_values(
                     target_optimizer_structure, saved_optimizer_values
@@ -702,5 +701,3 @@ def _restore_from_state_tree(self, state_tree, target_model):
         if self.verbose > 0:
             print_msg("OrbaxCheckpoint: Successfully restored model state")
         return True
-
-

keras/src/callbacks/orbax_checkpoint_test.py

@@ -359,7 +359,8 @@ def test_checkpoint_error_handling(self):
         with self.assertRaises(Exception):
             callback.load_checkpoint(step=999)
 
-        # Test: Try to load latest when no checkpoints exist - should raise FileNotFoundError
+        # Test: Try to load latest when no checkpoints exist -
+        # should raise FileNotFoundError
         with self.assertRaises(FileNotFoundError):
             callback.load_latest()