Ensure JDPinned produces types consistent with underlying JD

PiperOrigin-RevId: 339919800

tensorflow_probability/python/experimental/distributions/joint_distribution_pinned.py

@@ -355,16 +355,17 @@ def _model_flatten(self, x):
 
   def _model_unflatten(self, xs):
     """Unflattens `xs` to a structure like-typed to `self.distribution`."""
-    model = self.distribution.dtype
-    if isinstance(model, dict):
+    # Use the underlying JD dtype to infer model structure.
+    dtype = self.distribution.dtype
+    if isinstance(dtype, dict):
       ks = self._flat_resolve_names()
       if len(ks) != len(xs):
         raise ValueError('Invalid xs length {}, ks={}'.format(len(xs), ks))
-      return type(model)(zip(ks, xs))
-    if hasattr(model, '_fields') and hasattr(model, '_asdict'):
-      ks = [k for k in model._fields if k not in self.pins]
+      return type(dtype)(zip(ks, xs))
+    if hasattr(dtype, '_fields') and hasattr(dtype, '_asdict'):
+      ks = [k for k in dtype._fields if k not in self.pins]
       return structural_tuple.structtuple(ks)(*xs)
-    return tuple(xs)
+    return type(dtype)(xs)
 
   def _prune(self, xs, retain=None):
     """Drops fields from `xs`, retaining those specified by `retain`.

tensorflow_probability/python/experimental/distributions/joint_distribution_pinned_test.py

@@ -324,7 +324,9 @@ def test_bijector(self):
         lambda a: tfd.Uniform(a + tf.ones_like(a), a + tf.constant(2, a.dtype)),
         lambda b, a: tfd.Uniform(a, b, name='c')])
     bij = jd.experimental_default_event_space_bijector(a=-.5, b=1.)
-    self.assertAllClose((2/3,), tf.math.sigmoid(bij.inverse((0.5,))))
+    test_input = (0.5,)
+    self.assertIs(type(jd.dtype), type(bij.inverse(test_input)))
+    self.assertAllClose((2/3,), tf.math.sigmoid(bij.inverse(test_input)))
 
     @tfd.JointDistributionCoroutine
     def model():