Revert "Fix dtype issues of JAX CPU in SD3. (#2338)" (#2344)

This reverts commit ecef6d1.

Author: divyashreepathihalli

keras_hub/src/models/stable_diffusion_3/stable_diffusion_3_backbone_test.py

@@ -31,26 +31,11 @@ def setUp(self):
             64,
             "quick_gelu",
             -2,
-            # TODO: JAX CPU doesn't support float16 for
-            # `nn.dot_product_attention`. We set dtype to float32 despite the
-            # model defaulting to float16.
-            dtype="float32",
+            dtype="float16",
             name="clip_l",
         )
         clip_g = CLIPTextEncoder(
-            20,
-            64,
-            64,
-            2,
-            2,
-            128,
-            "gelu",
-            -2,
-            # TODO: JAX CPU doesn't support float16 for
-            # `nn.dot_product_attention`. We set dtype to float32 despite the
-            # model defaulting to float16.
-            dtype="float32",
-            name="clip_g",
+            20, 64, 64, 2, 2, 128, "gelu", -2, dtype="float16", name="clip_g"
         )
         self.init_kwargs = {
             "mmdit_patch_size": 2,