Fix GPU CI (keras-team#20637)

* Fix GPU CI

* Fix dtype issue

* Remove duplicate tests

keras/src/layers/preprocessing/image_preprocessing/bounding_boxes/bounding_box.py

@@ -419,11 +419,11 @@ def _compute_inverse_affine_matrix(
         shear_x = -shear_x
         shear_y = -shear_y
 
-        cx = center_x * (width - 1)
-        cy = center_y * (height - 1)
+        cx = ops.numpy.multiply(center_x, (width - 1))
+        cy = ops.numpy.multiply(center_y, (height - 1))
         rot = ops.numpy.multiply(angle, 1.0 / 180.0 * math.pi)
-        tx = -translate_x * (width - 1)
-        ty = -translate_y * (height - 1)
+        tx = ops.numpy.multiply(-translate_x, (width - 1))
+        ty = ops.numpy.multiply(-translate_y, (height - 1))
         sx = ops.numpy.multiply(shear_x, 1.0 / 180.0 * math.pi)
         sy = ops.numpy.multiply(shear_y, 1.0 / 180.0 * math.pi)
 
@@ -442,10 +442,10 @@ def _compute_inverse_affine_matrix(
 
         # Inverted rotation matrix with scale and shear
         # det([[a, b], [c, d]]) == 1, since det(rotation) = 1 and det(shear) = 1
-        a0 = d * scale
-        a1 = -b * scale
-        b0 = -c * scale
-        b1 = a * scale
+        a0 = ops.numpy.multiply(d, scale)
+        a1 = ops.numpy.multiply(-b, scale)
+        b0 = ops.numpy.multiply(-c, scale)
+        b1 = ops.numpy.multiply(a, scale)
         a2 = cx - a0 * cx_plus_tx - a1 * cy_plus_ty
         b2 = cy - b0 * cx_plus_tx - b1 * cy_plus_ty
 

keras/src/layers/preprocessing/image_preprocessing/bounding_boxes/converters_test.py

@@ -45,7 +45,7 @@ def setUp(self):
             [[[0.01, 0.02, 0.1, 0.1], [0.02, 0.03, 0.1, 0.1]]], dtype="float32"
         )
 
-        self.images = np.ones([2, 1000, 1000, 3])
+        self.images = np.ones([2, 1000, 1000, 3], dtype="float32")
         self.height = 1000
         self.width = 1000
 
@@ -131,12 +131,12 @@ def test_affine_identity(self):
         batch_size = self.boxes["xyxy"].shape[0]
         transformed_boxes = affine_transform(
             boxes=self.boxes["xyxy"],
-            angle=np.zeros([batch_size]),
-            translate_x=np.zeros([batch_size]),
-            translate_y=np.zeros([batch_size]),
-            scale=np.ones([batch_size]),
-            shear_x=np.zeros([batch_size]),
-            shear_y=np.zeros([batch_size]),
+            angle=np.zeros([batch_size], dtype="float32"),
+            translate_x=np.zeros([batch_size], dtype="float32"),
+            translate_y=np.zeros([batch_size], dtype="float32"),
+            scale=np.ones([batch_size], dtype="float32"),
+            shear_x=np.zeros([batch_size], dtype="float32"),
+            shear_y=np.zeros([batch_size], dtype="float32"),
             height=self.height,
             width=self.width,
         )

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

@@ -17,6 +17,8 @@ def test_layer(self):
             input_shape=(8, 3, 4, 3),
             supports_masking=False,
             expected_output_shape=(8, 3, 4, 3),
+            # StatelessRandomGammaV3 is not supported on XLA_GPU_JIT
+            run_training_check=not testing.tensorflow_uses_gpu(),
         )
 
     def test_mix_up_basic_functionality(self):

keras/src/trainers/trainer_test.py

@@ -2745,50 +2745,3 @@ def predict_step(self, *args):
             verbose=0,
         )
         self.assertLessEqual(tracing_count[0], 2)
-
-
-class TrainerDistributeTest(testing.TestCase):
-    @pytest.mark.skipif(
-        backend.backend() != "tensorflow", reason="Requires tf.distribute"
-    )
-    def test_end_to_end_tf_distribute(self):
-        import tensorflow as tf
-        from tensorflow.python.eager import context
-
-        context._reset_context()
-        cpus = tf.config.list_physical_devices("CPU")
-        tf.config.set_logical_device_configuration(
-            cpus[0],
-            [
-                tf.config.LogicalDeviceConfiguration(),
-                tf.config.LogicalDeviceConfiguration(),
-            ],
-        )
-        strategy = tf.distribute.MirroredStrategy(["CPU:0", "CPU:1"])
-        with strategy.scope():
-            model = keras.Sequential(
-                [
-                    keras.Input((2,)),
-                    keras.layers.Dense(
-                        2,
-                        activation="softmax",
-                        use_bias=False,
-                        kernel_initializer="ones",
-                    ),
-                ]
-            )
-            model.compile(
-                optimizer="sgd",
-                loss="sparse_categorical_crossentropy",
-                metrics=["sparse_categorical_accuracy"],
-            )
-        x = (np.arange(512) / 128).reshape((256, 2))
-        y = (np.arange(256) % 2).reshape((256, 1))
-        out_fit = model.fit(x, y)
-        self.assertLess(out_fit.history["sparse_categorical_accuracy"][0], 0.6)
-        out_eval = model.evaluate(x, y)
-        self.assertLess(out_eval[1], 0.6)
-        out_predict = model.predict(x)
-        self.assertEqual(out_predict.shape, (256, 2))
-
-        context._reset_context()