Enable dtype_util.size(...) to work in the numpy / jax backend (allow…

…ing Categorical with float64 parameters). PiperOrigin-RevId: 275943193
JulianTakealot · Oct 21, 2019 · fe80271 · fe80271
1 parent 748e66b
commit fe80271
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 2 deletions.
diff --git a/tensorflow_probability/python/internal/dtype_util.py b/tensorflow_probability/python/internal/dtype_util.py
@@ -150,7 +150,7 @@ def name(dtype):
 def size(dtype):
   """Returns the number of bytes to represent this `dtype`."""
   dtype = tf.as_dtype(dtype)
-  if hasattr(dtype, 'size'):
+  if hasattr(dtype, 'size') and hasattr(dtype, 'as_numpy_dtype'):
     return dtype.size
   return np.dtype(dtype).itemsize
 

diff --git a/tensorflow_probability/python/internal/dtype_util_test.py b/tensorflow_probability/python/internal/dtype_util_test.py
@@ -149,7 +149,17 @@ def test_assert_same_float_dtype(self):
     self.assertRaises(ValueError, dtype_util.assert_same_float_dtype,
                       [const_int])
 
+  def test_size(self):
+    self.assertEqual(dtype_util.size(tf.int32), 4)
+    self.assertEqual(dtype_util.size(tf.int64), 8)
+    self.assertEqual(dtype_util.size(tf.float32), 4)
+    self.assertEqual(dtype_util.size(tf.float64), 8)
+
+    self.assertEqual(dtype_util.size(np.int32), 4)
+    self.assertEqual(dtype_util.size(np.int64), 8)
+    self.assertEqual(dtype_util.size(np.float32), 4)
+    self.assertEqual(dtype_util.size(np.float64), 8)
+
 
 if __name__ == '__main__':
   tf.test.main()
-