format

Author: wenyi-guo

keras/src/layers/layer.py

@@ -1147,7 +1147,7 @@ def compute_output_spec(self, *args, **kwargs):
             )
             try:
                 output_shape = self.compute_output_shape(**shapes_dict)
-            except NotImplementedError as e:
+            except NotImplementedError:
                 return super().compute_output_spec(*args, **kwargs)
 
             if (

keras/src/utils/jax_layer.py

@@ -1,12 +1,13 @@
-import inspect
-
 import collections
 import functools
+import inspect
 import itertools
-import numpy as np
 import string
 
+import numpy as np
+
 from keras.src import backend
+from keras.src import ops
 from keras.src import random
 from keras.src import tree
 from keras.src.api_export import keras_export
@@ -16,7 +17,6 @@
 from keras.src.saving import serialization_lib
 from keras.src.utils import jax_utils
 from keras.src.utils import tracking
-from keras.src import ops
 from keras.src.utils.module_utils import jax
 from keras.src.utils.module_utils import tensorflow as tf
 

keras/src/utils/jax_layer_test.py

@@ -11,6 +11,8 @@
 from keras.src import layers
 from keras.src import metrics
 from keras.src import models
+from keras.src import ops
+from keras.src import random
 from keras.src import saving
 from keras.src import testing
 from keras.src import tree
@@ -19,8 +21,6 @@
 from keras.src.saving import object_registration
 from keras.src.utils.jax_layer import FlaxLayer
 from keras.src.utils.jax_layer import JaxLayer
-from keras.src import ops
-from keras.src import random
 
 try:
     import flax