[jax2tf] Fix bfloat16 bug in select_and_gather_add conversion. (#4058)

* [jax2tf] Fix bfloat16 bug in select_and_gather_add conversion.

This fix makes it possible to run bfloat16 tests for the jax2tf
conversion of select_and_gather_add.

jax/experimental/jax2tf/jax2tf.py

@@ -799,8 +799,8 @@ def _select_and_gather_add(tangents: TfVal,
                            padding: Sequence[Tuple[int, int]]):
   # Note: this function follows the pattern in
   # jax.lax._select_and_gather_add_translation.
-  dtype = to_jax_dtype(operand.dtype)
-  nbits = dtypes.finfo(dtype).bits
+  dtype = operand.dtype
+  nbits = dtypes.finfo(dtype.as_numpy_dtype).bits
 
   # Specializing the function for 64 bits. Only up to 32 bits are supported on TPU,
   # we thus intend to let the code throw a different exception on this platform.

jax/experimental/jax2tf/tests/primitives_test.py

@@ -265,9 +265,6 @@ def _reconstruct_operand(result, is_tf: bool):
   def test_select_and_gather_add(self, harness: primitive_harness.Harness):
     dtype = harness.params["dtype"]
 
-    if dtype is dtypes.bfloat16:
-      raise unittest.SkipTest("bfloat16 not implemented")
-
     max_bits = 64
     if jtu.device_under_test() == "tpu":
       max_bits = 32