Fix a bug that caused user-defined subclasses of JointDistributions to error in __new__.

PiperOrigin-RevId: 474175053

Author: emilyfertig

tensorflow_probability/python/distributions/joint_distribution_auto_batched.py

@@ -589,16 +589,17 @@ class JointDistributionSequentialAutoBatched(
 
   def __new__(cls, *args, **kwargs):
     """Maybe returns a `_JointDistributionSequentialAutobatched`."""
-    if args:
-      model = args[0]
-    else:
-      model = kwargs.get('model')
-
-    # Return a `_JointDistributionSequentialAutoBatched` instance if `model`
-    # contains distributions that are not CompositeTensors.
-    if not all(isinstance(d, tf.__internal__.CompositeTensor) or callable(d)
-               for d in model):
-      return _JointDistributionSequentialAutoBatched(*args, **kwargs)
+    if cls is JointDistributionSequentialAutoBatched:
+      if args:
+        model = args[0]
+      else:
+        model = kwargs.get('model')
+
+      # Return a `_JointDistributionSequentialAutoBatched` instance if `model`
+      # contains distributions that are not CompositeTensors.
+      if not all(isinstance(d, tf.__internal__.CompositeTensor) or callable(d)
+                 for d in model):
+        return _JointDistributionSequentialAutoBatched(*args, **kwargs)
     return super(JointDistributionSequentialAutoBatched, cls).__new__(cls)
 
   @property
@@ -625,16 +626,17 @@ class JointDistributionNamedAutoBatched(
 
   def __new__(cls, *args, **kwargs):
     """Maybe returns a `_JointDistributionNamedAutoBatched`."""
-    if args:
-      model = args[0]
-    else:
-      model = kwargs.get('model')
-
-    # Return a `_JointDistributionNamedAutoBatched` instance if `model` contains
-    # distributions that are not CompositeTensors.
-    if not all(isinstance(d, tf.__internal__.CompositeTensor) or callable(d)
-               for d in tf.nest.flatten(model)):
-      return _JointDistributionNamedAutoBatched(*args, **kwargs)
+    if cls is JointDistributionNamedAutoBatched:
+      if args:
+        model = args[0]
+      else:
+        model = kwargs.get('model')
+
+      # Return a `_JointDistributionNamedAutoBatched` instance if `model`
+      # contains distributions that are not CompositeTensors.
+      if not all(isinstance(d, tf.__internal__.CompositeTensor) or callable(d)
+                 for d in tf.nest.flatten(model)):
+        return _JointDistributionNamedAutoBatched(*args, **kwargs)
     return super(JointDistributionNamedAutoBatched, cls).__new__(cls)
 
   @property

tensorflow_probability/python/distributions/joint_distribution_named.py

@@ -464,14 +464,15 @@ class JointDistributionNamed(_JointDistributionNamed,
 
   def __new__(cls, *args, **kwargs):
     """Returns a `_JointDistributionNamed` if `model` contains non-CT dists."""
-    if args:
-      model = args[0]
-    else:
-      model = kwargs.get('model')
+    if cls is JointDistributionNamed:
+      if args:
+        model = args[0]
+      else:
+        model = kwargs.get('model')
 
-    if not all(isinstance(d, tf.__internal__.CompositeTensor) or callable(d)
-               for d in tf.nest.flatten(model)):
-      return _JointDistributionNamed(*args, **kwargs)
+      if not all(isinstance(d, tf.__internal__.CompositeTensor) or callable(d)
+                 for d in tf.nest.flatten(model)):
+        return _JointDistributionNamed(*args, **kwargs)
     return super(JointDistributionNamed, cls).__new__(cls)
 
   @property

tensorflow_probability/python/distributions/joint_distribution_sequential.py

@@ -728,14 +728,15 @@ class JointDistributionSequential(_JointDistributionSequential,
 
   def __new__(cls, *args, **kwargs):
     """Returns a `_JointDistributionSequential` if `model` has non-CT dists."""
-    if args:
-      model = args[0]
-    else:
-      model = kwargs.get('model')
-
-    if not all(isinstance(d, tf.__internal__.CompositeTensor) or callable(d)
-               for d in model):
-      return _JointDistributionSequential(*args, **kwargs)
+    if cls is JointDistributionSequential:
+      if args:
+        model = args[0]
+      else:
+        model = kwargs.get('model')
+
+      if not all(isinstance(d, tf.__internal__.CompositeTensor) or callable(d)
+                 for d in model):
+        return _JointDistributionSequential(*args, **kwargs)
     return super(JointDistributionSequential, cls).__new__(cls)
 
   @property