Allow jnp.bfloat16 arrays to be correctly recognized as floats.

PiperOrigin-RevId: 492022573

Author: ColCarroll

tensorflow_probability/python/internal/dtype_util.py

@@ -47,6 +47,13 @@
 SKIP_DTYPE_CHECKS = False
 
 
+_issubdtype = np.issubdtype
+if JAX_MODE:
+  # jnp.issubdtype handles custom JAX types like bfloat16
+  import jax.numpy as jnp  # pylint: disable=g-import-not-at-top
+  _issubdtype = jnp.issubdtype
+
+
 def is_numpy_compatible(dtype):
   """Returns if dtype has a corresponding NumPy dtype."""
   if JAX_MODE or NUMPY_MODE:
@@ -270,23 +277,23 @@ def is_complex(dtype):
   dtype = tf.as_dtype(dtype)
   if hasattr(dtype, 'is_complex'):
     return dtype.is_complex
-  return np.issubdtype(np.dtype(dtype), np.complexfloating)
+  return _issubdtype(np.dtype(dtype), np.complexfloating)
 
 
 def is_floating(dtype):
   """Returns whether this is a (non-quantized, real) floating point type."""
   dtype = tf.as_dtype(dtype)
   if hasattr(dtype, 'is_floating'):
     return dtype.is_floating
-  return np.issubdtype(np.dtype(dtype), np.floating)
+  return _issubdtype(np.dtype(dtype), np.floating)
 
 
 def is_integer(dtype):
   """Returns whether this is a (non-quantized) integer type."""
   dtype = tf.as_dtype(dtype)
   if hasattr(dtype, 'is_integer') and not callable(dtype.is_integer):
     return dtype.is_integer
-  return np.issubdtype(np.dtype(dtype), np.integer)
+  return _issubdtype(np.dtype(dtype), np.integer)
 
 
 def max(dtype):  # pylint: disable=redefined-builtin

tensorflow_probability/python/internal/dtype_util_test.py

@@ -25,6 +25,7 @@
 from tensorflow_probability.python.internal import test_util
 
 
+NUMPY_MODE = False
 JAX_MODE = False
 
 
@@ -281,6 +282,14 @@ def test_size(self):
     self.assertEqual(dtype_util.size(np.float32), 4)
     self.assertEqual(dtype_util.size(np.float64), 8)
 
+  @parameterized.named_parameters(
+      ('float32', tf.float32, True),
+      ('bfloat16', 'bfloat16', True),
+      ('not_int8', tf.int8, False))
+  def test_is_floating(self, dtype, expected):
+    if NUMPY_MODE and dtype == 'bfloat16':
+      self.skipTest('No bfloat16 in numpy')
+    self.assertEqual(dtype_util.is_floating(dtype), expected)
 
 if __name__ == '__main__':
   test_util.main()