[Pallas] Fix lowering tests for reduction ops

Remove unnecessary skip statements. Also added tests for bf16 types.

PiperOrigin-RevId: 707739536

Author: ayaka14732

tests/pallas/ops_test.py

@@ -1895,35 +1895,35 @@ def reduce(x_ref, y_ref):
       for axis in [0, 1, (1,), (0, 1)]
       for dtype in [
           "float16",
+          "bfloat16",
           "float32",
           "float64",
           "int32",
           "int64",
           "uint32",
           "uint64",
       ]
-      if isinstance(axis, int) or "arg" not in op_name
   ])
   def test_array_reduce(self, op, dtype, axis):
-    if jtu.test_device_matches(["tpu"]) and jnp.dtype(dtype).itemsize == 2:
-      self.skipTest("16-bit types are not supported on TPU")
+    if not isinstance(axis, int):
+      self.skipTest("TODO: tuple axes are not yet supported")
 
     if not jax.config.x64_enabled and jnp.dtype(dtype).itemsize == 8:
       self.skipTest("64-bit types require x64_enabled")
 
+    # The Pallas TPU lowering currently supports only blocks of rank >= 1
+    if jtu.test_device_matches(["tpu"]):
+      self.skipTest("Not implemented on TPU")
+
     # Skip argmin/argmax on GPU in 64-bit mode because Pallas expects
     # `index_type` to be i32
     if (
         jax.config.x64_enabled
         and jtu.test_device_matches(["gpu"])
-        and op in {jnp.argmin, jnp.argmax}
+        and op in (jnp.argmin, jnp.argmax)
     ):
       self.skipTest("Not supported on GPU in 64-bit mode")
 
-    # The Pallas TPU lowering currently supports only blocks of rank >= 1
-    if jtu.test_device_matches(["tpu"]):
-      self.skipTest("Not supported on TPU")
-
     m, n = 32, 8
 
     def make_x(key):
@@ -1955,7 +1955,7 @@ def reduce(x_ref, y_ref):
       x = make_x(key)
       y = reduce(x)
       y_ref = op(x, axis=axis)
-      np.testing.assert_allclose(y, y_ref, atol=1e-2, rtol=1e-2, err_msg=i)
+      self.assertAllClose(y, y_ref, atol=1e-2, rtol=1e-2, err_msg=i)
 
   @parameterized.product(
       axis=[0, 1],