Fix JAX GPU CI and make formatter happy (keras-team#20749)

* Fix JAX GPU CI

* Makes formatter happy

* Makes formatter happy - 2

Author: james77777778

examples/demo_custom_torch_workflow.py

@@ -74,8 +74,8 @@ def train(model, train_loader, num_epochs, optimizer, loss_fn):
             # Print loss statistics
             if (batch_idx + 1) % 10 == 0:
                 print(
-                    f"Epoch [{epoch+1}/{num_epochs}], "
-                    f"Batch [{batch_idx+1}/{len(train_loader)}], "
+                    f"Epoch [{epoch + 1}/{num_epochs}], "
+                    f"Batch [{batch_idx + 1}/{len(train_loader)}], "
                     f"Loss: {running_loss / 10}"
                 )
                 running_loss = 0.0

examples/demo_torch_multi_gpu.py

@@ -104,8 +104,8 @@ def train(model, train_loader, num_epochs, optimizer, loss_fn):
             # Print loss statistics
             if (batch_idx + 1) % 10 == 0:
                 print(
-                    f"Epoch [{epoch+1}/{num_epochs}], "
-                    f"Batch [{batch_idx+1}/{len(train_loader)}], "
+                    f"Epoch [{epoch + 1}/{num_epochs}], "
+                    f"Batch [{batch_idx + 1}/{len(train_loader)}], "
                     f"Loss: {running_loss / 10}"
                 )
                 running_loss = 0.0

keras/src/backend/openvino/core.py

@@ -242,9 +242,9 @@ def __getitem__(self, indices):
     def __len__(self):
         ov_output = self.output
         ov_shape = ov_output.get_partial_shape()
-        assert (
-            ov_shape.rank.is_static and ov_shape.rank.get_length() > 0
-        ), "rank must be static and greater than zero"
+        assert ov_shape.rank.is_static and ov_shape.rank.get_length() > 0, (
+            "rank must be static and greater than zero"
+        )
         assert ov_shape[0].is_static, "the first dimension must be static"
         return ov_shape[0].get_length()
 
@@ -425,10 +425,10 @@ def convert_to_numpy(x):
             x = x.value
         else:
             return x.value.data
-    assert isinstance(
-        x, OpenVINOKerasTensor
-    ), "unsupported type {} for `convert_to_numpy` in openvino backend".format(
-        type(x)
+    assert isinstance(x, OpenVINOKerasTensor), (
+        "unsupported type {} for `convert_to_numpy` in openvino backend".format(
+            type(x)
+        )
     )
     try:
         ov_result = x.output

keras/src/backend/openvino/nn.py

@@ -325,9 +325,9 @@ def depthwise_conv(
     data_format = backend.standardize_data_format(data_format)
     num_spatial_dims = inputs.get_partial_shape().rank.get_length() - 2
 
-    assert (
-        data_format == "channels_last"
-    ), "`depthwise_conv` is supported only for channels_last data_format"
+    assert data_format == "channels_last", (
+        "`depthwise_conv` is supported only for channels_last data_format"
+    )
 
     strides = _adjust_strides_dilation(strides, num_spatial_dims)
     dilation_rate = _adjust_strides_dilation(dilation_rate, num_spatial_dims)

keras/src/backend/openvino/numpy.py

@@ -81,9 +81,9 @@ def mean(x, axis=None, keepdims=False):
 
 
 def max(x, axis=None, keepdims=False, initial=None):
-    assert (
-        initial is None
-    ), "`max` with not None initial is not supported by openvino backend"
+    assert initial is None, (
+        "`max` with not None initial is not supported by openvino backend"
+    )
     x = get_ov_output(x)
     reduce_axis = ov_opset.constant(axis, Type.i32).output(0)
     return OpenVINOKerasTensor(
@@ -260,9 +260,9 @@ def bincount(x, weights=None, minlength=0, sparse=False):
 
 
 def broadcast_to(x, shape):
-    assert isinstance(
-        shape, (tuple, list)
-    ), "`broadcast_to` is supported only for tuple and list `shape`"
+    assert isinstance(shape, (tuple, list)), (
+        "`broadcast_to` is supported only for tuple and list `shape`"
+    )
     target_shape = ov_opset.constant(list(shape), Type.i32).output(0)
     x = get_ov_output(x)
     return OpenVINOKerasTensor(ov_opset.broadcast(x, target_shape).output(0))

keras/src/backend/tensorflow/linalg.py

@@ -203,8 +203,7 @@ def lstsq(a, b, rcond=None):
         b = b[:, None]
     if a.ndim != 2:
         raise TypeError(
-            f"{a.ndim}-dimensional array given. "
-            "Array must be two-dimensional"
+            f"{a.ndim}-dimensional array given. Array must be two-dimensional"
         )
     if b.ndim != 2:
         raise TypeError(

keras/src/backend/tensorflow/numpy.py

@@ -2469,8 +2469,7 @@ def squeeze(x, axis=None):
         for a in axis:
             if static_shape[a] != 1:
                 raise ValueError(
-                    f"Cannot squeeze axis={a}, because the "
-                    "dimension is not 1."
+                    f"Cannot squeeze axis={a}, because the dimension is not 1."
                 )
         axis = sorted([canonicalize_axis(a, len(static_shape)) for a in axis])
     if isinstance(x, tf.SparseTensor):

keras/src/callbacks/csv_logger.py

@@ -59,7 +59,7 @@ def handle_value(k):
                 isinstance(k, collections.abc.Iterable)
                 and not is_zero_dim_ndarray
             ):
-                return f"\"[{', '.join(map(str, k))}]\""
+                return f'"[{", ".join(map(str, k))}]"'
             else:
                 return k
 

keras/src/callbacks/swap_ema_weights_test.py

@@ -53,7 +53,7 @@ def test_swap_ema_weights_with_invalid_optimizer(self):
         model = self._get_compiled_model(use_ema=False)
         with self.assertRaisesRegex(
             ValueError,
-            ("SwapEMAWeights must be used when " "`use_ema=True` is set"),
+            ("SwapEMAWeights must be used when `use_ema=True` is set"),
         ):
             model.fit(
                 self.x_train,

keras/src/export/export_utils.py

@@ -65,7 +65,7 @@ def make_input_spec(x):
     if isinstance(x, layers.InputSpec):
         if x.shape is None or x.dtype is None:
             raise ValueError(
-                "The `shape` and `dtype` must be provided. " f"Received: x={x}"
+                f"The `shape` and `dtype` must be provided. Received: x={x}"
             )
         input_spec = x
     elif isinstance(x, backend.KerasTensor):

keras/src/export/onnx.py

@@ -116,8 +116,7 @@ def _check_jax_kwargs(kwargs):
         }
     if kwargs["is_static"] is not True:
         raise ValueError(
-            "`is_static` must be `True` in `kwargs` when using the jax "
-            "backend."
+            "`is_static` must be `True` in `kwargs` when using the jax backend."
         )
     if kwargs["jax2tf_kwargs"]["enable_xla"] is not False:
         raise ValueError(

keras/src/layers/activations/leaky_relu.py

@@ -39,8 +39,7 @@ def __init__(self, negative_slope=0.3, **kwargs):
         if "alpha" in kwargs:
             negative_slope = kwargs.pop("alpha")
             warnings.warn(
-                "Argument `alpha` is deprecated. "
-                "Use `negative_slope` instead."
+                "Argument `alpha` is deprecated. Use `negative_slope` instead."
             )
         super().__init__(**kwargs)
         if negative_slope is None or negative_slope < 0:

keras/src/layers/attention/grouped_query_attention.py

@@ -112,8 +112,7 @@ def __init__(
         self.num_key_value_heads = num_key_value_heads
         if num_query_heads % num_key_value_heads != 0:
             raise ValueError(
-                "`num_query_heads` must be divisible"
-                " by `num_key_value_heads`."
+                "`num_query_heads` must be divisible by `num_key_value_heads`."
             )
         self.num_repeats = num_query_heads // num_key_value_heads
         self.dropout = dropout

keras/src/layers/attention/multi_head_attention_test.py

@@ -644,7 +644,7 @@ def test_multi_head_attention_output_shape_as_int(self):
         output = mha(query=query, value=value)
 
         assert output.shape == (2, 4, 8), (
-            f"Expected shape (2, 4, 8)," f" got {output.shape}"
+            f"Expected shape (2, 4, 8), got {output.shape}"
         )
 
     def test_multi_head_attention_output_shape_as_tuple(self):
@@ -657,7 +657,7 @@ def test_multi_head_attention_output_shape_as_tuple(self):
         output = mha(query=query, value=value)
 
         assert output.shape == (2, 4, 8, 8), (
-            f"Expected shape (2, 4, 8, 8)," f" got {output.shape}"
+            f"Expected shape (2, 4, 8, 8), got {output.shape}"
         )
 
     def test_multi_head_attention_output_shape_error(self):

keras/src/layers/convolutional/base_conv.py

@@ -282,8 +282,7 @@ def enable_lora(
             )
         if self.lora_enabled:
             raise ValueError(
-                "lora is already enabled. "
-                "This can only be done once per layer."
+                "lora is already enabled. This can only be done once per layer."
             )
         self._tracker.unlock()
         self.lora_kernel_a = self.add_weight(

keras/src/layers/core/dense.py

@@ -168,8 +168,7 @@ def enable_lora(
             )
         if self.lora_enabled:
             raise ValueError(
-                "lora is already enabled. "
-                "This can only be done once per layer."
+                "lora is already enabled. This can only be done once per layer."
             )
         self._tracker.unlock()
         self.lora_kernel_a = self.add_weight(

keras/src/layers/core/einsum_dense.py

@@ -224,8 +224,7 @@ def enable_lora(
             )
         if self.lora_enabled:
             raise ValueError(
-                "lora is already enabled. "
-                "This can only be done once per layer."
+                "lora is already enabled. This can only be done once per layer."
             )
         self._tracker.unlock()
         self.lora_kernel_a = self.add_weight(

keras/src/layers/core/embedding.py

@@ -163,8 +163,7 @@ def enable_lora(
             )
         if self.lora_enabled:
             raise ValueError(
-                "lora is already enabled. "
-                "This can only be done once per layer."
+                "lora is already enabled. This can only be done once per layer."
             )
         self._tracker.unlock()
         self.lora_embeddings_a = self.add_weight(

keras/src/layers/layer.py

@@ -1430,8 +1430,7 @@ def _build_by_run_for_kwargs(self, shapes_dict):
 
     def __repr__(self):
         return (
-            f"<{self.__class__.__name__} "
-            f"name={self.name}, built={self.built}>"
+            f"<{self.__class__.__name__} name={self.name}, built={self.built}>"
         )
 
     def __str__(self):

keras/src/layers/normalization/layer_normalization_test.py

@@ -87,10 +87,7 @@ def test_ln_basics(self):
     def test_invalid_axis(self):
         with self.assertRaisesRegex(
             TypeError,
-            (
-                "Expected an int or a list/tuple of ints for the argument "
-                "'axis'"
-            ),
+            ("Expected an int or a list/tuple of ints for the argument 'axis'"),
         ):
             layers.LayerNormalization(axis={"axis": -1})
 

keras/src/layers/preprocessing/hashed_crossing.py

@@ -90,7 +90,7 @@ def __init__(
         super().__init__(name=name, dtype=dtype)
         if sparse and backend.backend() != "tensorflow":
             raise ValueError(
-                "`sparse=True` can only be used with the " "TensorFlow backend."
+                "`sparse=True` can only be used with the TensorFlow backend."
             )
 
         argument_validation.validate_string_arg(

keras/src/layers/preprocessing/image_preprocessing/random_grayscale.py

@@ -49,8 +49,7 @@ def __init__(self, factor=0.5, data_format=None, seed=None, **kwargs):
         super().__init__(**kwargs)
         if factor < 0 or factor > 1:
             raise ValueError(
-                "`factor` should be between 0 and 1. "
-                f"Received: factor={factor}"
+                f"`factor` should be between 0 and 1. Received: factor={factor}"
             )
         self.factor = factor
         self.data_format = backend.standardize_data_format(data_format)

keras/src/layers/preprocessing/integer_lookup.py

@@ -328,7 +328,7 @@ def __init__(
             )
         if sparse and backend.backend() != "tensorflow":
             raise ValueError(
-                "`sparse=True` can only be used with the " "TensorFlow backend."
+                "`sparse=True` can only be used with the TensorFlow backend."
             )
         if vocabulary_dtype != "int64":
             raise ValueError(

keras/src/layers/preprocessing/string_lookup.py

@@ -314,7 +314,7 @@ def __init__(
             )
         if sparse and backend.backend() != "tensorflow":
             raise ValueError(
-                "`sparse=True` can only be used with the " "TensorFlow backend."
+                "`sparse=True` can only be used with the TensorFlow backend."
             )
         self.encoding = encoding
         super().__init__(

keras/src/layers/preprocessing/text_vectorization.py

@@ -226,11 +226,11 @@ def __init__(
             )
         if sparse and backend.backend() != "tensorflow":
             raise ValueError(
-                "`sparse=True` can only be used with the " "TensorFlow backend."
+                "`sparse=True` can only be used with the TensorFlow backend."
             )
         if ragged and backend.backend() != "tensorflow":
             raise ValueError(
-                "`ragged=True` can only be used with the " "TensorFlow backend."
+                "`ragged=True` can only be used with the TensorFlow backend."
             )
 
         # 'standardize' must be one of

keras/src/metrics/metric.py

@@ -247,7 +247,7 @@ def _check_super_called(self):
             )
 
     def __repr__(self):
-        return f"<{self.__class__.__name__} " f"name={self.name}>"
+        return f"<{self.__class__.__name__} name={self.name}>"
 
     def __str__(self):
         return self.__repr__()

keras/src/ops/image.py

@@ -978,8 +978,7 @@ def _pad_images(
         )
     if left_padding < 0:
         raise ValueError(
-            "left_padding must be >= 0. "
-            f"Received: left_padding={left_padding}"
+            f"left_padding must be >= 0. Received: left_padding={left_padding}"
         )
     if right_padding < 0:
         raise ValueError(
@@ -1198,8 +1197,7 @@ def _crop_images(
 
     if top_cropping < 0:
         raise ValueError(
-            "top_cropping must be >= 0. "
-            f"Received: top_cropping={top_cropping}"
+            f"top_cropping must be >= 0. Received: top_cropping={top_cropping}"
         )
     if target_height < 0:
         raise ValueError(
@@ -1213,8 +1211,7 @@ def _crop_images(
         )
     if target_width < 0:
         raise ValueError(
-            "target_width must be >= 0. "
-            f"Received: target_width={target_width}"
+            f"target_width must be >= 0. Received: target_width={target_width}"
         )
 
     # Compute start_indices and shape

keras/src/ops/linalg.py

@@ -596,12 +596,11 @@ def call(self, a, b):
     def compute_output_spec(self, a, b):
         if len(a.shape) != 2:
             raise ValueError(
-                "Expected a to have rank 2. " f"Received: a.shape={a.shape}"
+                f"Expected a to have rank 2. Received: a.shape={a.shape}"
             )
         if len(b.shape) not in (1, 2):
             raise ValueError(
-                "Expected b to have rank 1 or 2. "
-                f"Received: b.shape={b.shape}"
+                f"Expected b to have rank 1 or 2. Received: b.shape={b.shape}"
             )
         m, n = a.shape
         if b.shape[0] != m:
@@ -666,8 +665,7 @@ def _assert_1d(*arrays):
     for a in arrays:
         if a.ndim < 1:
             raise ValueError(
-                "Expected input to have rank >= 1. "
-                "Received scalar input {a}."
+                "Expected input to have rank >= 1. Received scalar input {a}."
             )
 
 

keras/src/ops/numpy.py

@@ -6754,7 +6754,7 @@ class Argpartition(Operation):
     def __init__(self, kth, axis=-1):
         super().__init__()
         if not isinstance(kth, int):
-            raise ValueError("kth must be an integer. Received:" f"kth = {kth}")
+            raise ValueError(f"kth must be an integer. Received:kth = {kth}")
         self.kth = kth
         self.axis = axis
 

keras/src/random/seed_generator.py

@@ -76,7 +76,7 @@ def __init__(self, seed=None, name=None, **kwargs):
 
         if not isinstance(seed, int):
             raise ValueError(
-                "Argument `seed` must be an integer. " f"Received: seed={seed}"
+                f"Argument `seed` must be an integer. Received: seed={seed}"
             )
 
         def seed_initializer(*args, **kwargs):

keras/src/saving/file_editor.py

@@ -500,7 +500,7 @@ def _generate_metadata_info(self, rich_style=False):
         if rich_style:
             version = f"{summary_utils.highlight_symbol(version)}"
             date = f"{summary_utils.highlight_symbol(date)}"
-        return f"Saved with Keras {version} " f"- date: {date}"
+        return f"Saved with Keras {version} - date: {date}"
 
     def _print_weights_structure(
         self, weights_dict, indent=0, is_first=True, prefix="", inner_path=""