dtype_util.common_dtype handles nested dtypes.

PiperOrigin-RevId: 476199135

Author: emilyfertig

tensorflow_probability/python/internal/dtype_util.py

@@ -84,27 +84,101 @@ def __repr__(self):
 
 
 def common_dtype(args, dtype_hint=None):
-  """Returns explict dtype from `args` if there is one."""
+  """Returns (nested) explict dtype from `args` if there is one.
+
+  Dtypes of all args, and `dtype_hint`, must have the same nested structure if
+  they are not `None`. `args` itself may be any nested structure; its
+  structure is flattened and ignored.
+
+  Args:
+    args: A nested structure of objects that may have `dtype`.
+    dtype_hint: Optional (nested) dtype containing defaults to use in place of
+      `None`. If `dtype_hint` is not nested and the common dtype of `args` is
+      nested, `dtype_hint` serves as the default for each element of the common
+      nested dtype structure.
+
+  Returns:
+    dtype: The (nested) dtype common across all elements of `args`, or `None`.
+
+  #### Examples
+
+  Usage with non-nested dtype:
+
+  ```python
+  x = tf.ones([3, 4], dtype=tf.float64)
+  y = 4.
+  z = None
+  common_dtype([x, y, z], dtype_hint=tf.float32)  # ==> tf.float64
+  common_dtype([y, z], dtype_hint=tf.float32)     # ==> tf.float32
+
+  # The arg to `common_dtype` can be an arbitrary nested structure; it is
+  # flattened, and the common dtype of its contents is returned.
+  common_dtype({'x': x, 'yz': (y, z)})
+  # ==> tf.float64
+  ```
+
+  Usage with nested dtype:
+
+  ```python
+  # Define `x` and `y` as JointDistributions with the same nested dtype.
+  x = tfd.JointDistributionNamed(
+      {'a': tfd.Uniform(np.float64(0.), 1.),
+       'b': tfd.JointDistributionSequential(
+          [tfd.Normal(0., 2.), tfd.Bernoulli(0.4)])})
+  x.dtype  # ==> {'a': tf.float64, 'b': [tf.float32, tf.int32]}
+
+  y = tfd.JointDistributionNamed(
+      {'a': tfd.LogitNormal(np.float64(0.), 1.),
+       'b': tfd.JointDistributionSequential(
+          [tfd.Normal(-1., 1.), tfd.Bernoulli(0.6)])})
+  y.dtype  # ==> {'a': tf.float64, 'b': [tf.float32, tf.int32]}
+
+  # Pack x and y into an arbitrary nested structure and pass it to
+  # `common_dtype`.
+  args0 = [x, y]
+  common_dtype(args0)  # ==> {'a': tf.float64, 'b': [tf.float32, tf.int32]}
+
+  # The nested structure of the argument to `common_dtype` is flattened and
+  # ignored; only the nested structures of the dtypes are relevant.
+  args1 = {'x': x, 'yz': {'y': y, 'z': None}}
+  common_dtype(args1)  # ==> {'a': tf.float64, 'b': [tf.float32, tf.int32]}
+  ```
+  """
+
+  def _unify_dtype(current, new):
+    if current is not None and new is not None and current != new:
+      if SKIP_DTYPE_CHECKS:
+        return (np.ones([2], dtype) + np.ones([2], dt)).dtype
+      raise TypeError
+    return new if current is None else current
+
   dtype = None
   flattened_args = tf.nest.flatten(args)
   seen = [_NOT_YET_SEEN] * len(flattened_args)
   for i, a in enumerate(flattened_args):
     if hasattr(a, 'dtype') and a.dtype:
-      dt = as_numpy_dtype(a.dtype)
+      dt = tf.nest.map_structure(
+          lambda d: d if d is None else as_numpy_dtype(d), a.dtype)
       seen[i] = dt
     else:
       seen[i] = None
       continue
     if dtype is None:
       dtype = dt
-    elif dtype != dt:
-      if SKIP_DTYPE_CHECKS:
-        dtype = (np.ones([2], dtype) + np.ones([2], dt)).dtype
-      else:
-        raise TypeError(
-            'Found incompatible dtypes, {} and {}. Seen so far: {}'.format(
-                dtype, dt, tf.nest.pack_sequence_as(args, seen)))
-  return base_dtype(dtype_hint) if dtype is None else base_dtype(dtype)
+    try:
+      dtype = tf.nest.map_structure(_unify_dtype, dtype, dt)
+    except TypeError:
+      raise TypeError(
+          'Found incompatible dtypes, {} and {}. Seen so far: {}'.format(
+              dtype, dt, tf.nest.pack_sequence_as(args, seen))) from None
+  if dtype_hint is None:
+    return tf.nest.map_structure(base_dtype, dtype)
+  if dtype is None:
+    return tf.nest.map_structure(base_dtype, dtype_hint)
+  if tf.nest.is_nested(dtype) and not tf.nest.is_nested(dtype_hint):
+    dtype_hint = tf.nest.map_structure(lambda _: dtype_hint, dtype)
+  return tf.nest.map_structure(
+      lambda dt, h: base_dtype(h if dt is None else dt), dtype, dtype_hint)
 
 
 def convert_to_dtype(tensor_or_dtype, dtype=None, dtype_hint=None):

tensorflow_probability/python/internal/dtype_util_test.py

@@ -14,7 +14,7 @@
 # ============================================================================
 """Tests for dtype_util."""
 
-import collections
+import dataclasses
 
 # Dependency imports
 from absl.testing import parameterized
@@ -51,16 +51,50 @@ def testCommonDtypeAcceptsNone(self):
         tf.float16, dtype_util.common_dtype(
             [x, None], dtype_hint=tf.float32))
 
-    fake_tensor = collections.namedtuple('fake_tensor', ['dtype'])
+    fake_tensor = dataclasses.make_dataclass('fake_tensor', ['dtype'])
     self.assertEqual(
         tf.float16, dtype_util.common_dtype(
             [fake_tensor(dtype=None), None, x], dtype_hint=tf.float32))
+    self.assertEqual(
+        tf.float32, dtype_util.common_dtype(
+            [fake_tensor(dtype=tf.float32), None]))
 
   def testCommonDtypeFromLinop(self):
     x = tf.linalg.LinearOperatorDiag(tf.ones(3, tf.float16))
     self.assertEqual(
         tf.float16, dtype_util.common_dtype([x], dtype_hint=tf.float32))
 
+  def testCommonStructuredDtype(self):
+    structured_dtype_obj = dataclasses.make_dataclass(
+        'structured_dtype_obj', ['dtype'])
+    x = structured_dtype_obj({'a': tf.float32, 'b': (None, None)})
+    y = structured_dtype_obj({'a': None, 'b': (None, tf.float64)})
+    z = structured_dtype_obj({'a': None, 'b': (None, None)})
+    w = structured_dtype_obj(None)
+
+    # Check that structured dtypes unify correctly.
+    self.assertAllEqualNested(
+        dtype_util.common_dtype([w, x, y, z]),
+        {'a': tf.float32, 'b': (None, tf.float64)})
+
+    # Check that dict `args` works and that `dtype_hint` works.
+    dtype_hint = {'a': tf.int32, 'b': (tf.int32, None)}
+    self.assertAllEqualNested(
+        dtype_util.common_dtype(
+            {'x': x, 'y': y, 'z': z}, dtype_hint=dtype_hint),
+        {'a': tf.float32, 'b': (tf.int32, tf.float64)})
+    self.assertAllEqualNested(
+        dtype_util.common_dtype([w], dtype_hint=dtype_hint),
+        dtype_hint)
+
+    # Check that non-nested dtype_hint broadcasts.
+    self.assertAllEqualNested(
+        dtype_util.common_dtype([y, z], dtype_hint=tf.int32),
+        {'a': tf.int32, 'b': (tf.int32, tf.float64)})
+
+    with self.assertRaisesRegex(TypeError, 'Found incompatible dtypes'):
+      dtype_util.common_dtype([x, structured_dtype_obj(dtype_hint)])
+
   @parameterized.named_parameters(
       dict(testcase_name='Float32', dtype=tf.float32,
            expected_minval=np.float32(-3.4028235e+38)),