Remove ancient XLA:TPU workaround

We've supported S16 for a long time now PiperOrigin-RevId: 573260447
google · Oct 13, 2023 · e088a8e · e088a8e
1 parent 39a3cfc
commit e088a8e
Showing 1 changed file with 0 additions and 6 deletions.
diff --git a/jax/_src/numpy/ufuncs.py b/jax/_src/numpy/ufuncs.py
@@ -477,12 +477,6 @@ def signbit(x: ArrayLike, /) -> Array:
     raise ValueError(
         "jax.numpy.signbit is not well defined for %s" % dtype)
 
-  # TPU supports BF16 but not S16 types, so as a workaround, convert BF16 to
-  # F32.
-  if dtype == dtypes.bfloat16:
-    dtype = np.dtype('float32')
-    x = lax.convert_element_type(x, dtype)
-
   info = dtypes.finfo(dtype)
   if info.bits not in _INT_DTYPES:
     raise NotImplementedError(