class_docstr: Joint distribution parameterized by a distribution-making generator.
  
  This class provides automatic vectorization and alternative semantics for
  `tfd.JointDistributionCoroutine`, which in many cases allows for
  simplifications in the model specification.
  
  #### Automatic vectorization
  
  Auto-vectorized variants of JointDistribution allow the user to avoid
  explicitly annotating a model's vectorization semantics.
  When using manually-vectorized joint distributions, each operation in the
  model must account for the possibility of batch dimensions in Distributions
  and their samples. By contrast, auto-vectorized models need only describe
  a *single* sample from the joint distribution; any batch evaluation is
  automated using `tf.vectorized_map` as required. In many cases this
  allows for significant simplications. For example, the following
  manually-vectorized `tfd.JointDistributionCoroutine` model:
  
  ```python
  def model_fn():
    x = yield tfd.JointDistributionCoroutine.Root(
      tfd.Normal(0., tf.ones([3])))
    y = yield tfd.JointDistributionCoroutine.Root(
      tfd.Normal(0., 1.))
    z = yield tfd.Normal(x[..., :2] + y[..., tf.newaxis], 1.)
  
  can be written in auto-vectorized form as
  
  ```python
  def model_fn():
    x = yield tfd.Normal(0., tf.ones([3]))
    y = yield tfd.Normal(0., 1.)
    z = yield tfd.Normal(x[:2] + y, 1.)
  ```
  
  in which we were able to drop the specification of `Root` nodes and to
  avoid explicitly accounting for batch dimensions when indexing and slicing
  computed quantities in the third line.
  
  Note: auto-vectorization is still experimental and some TensorFlow ops may
  be unsupported. It can be disabled by setting `use_vectorized_map=False`.
  
  #### Alternative batch semantics
  
  This class also provides alternative semantics for specifying a batch of
  independent (non-identical) joint distributions.
  
  Instead of simply summing the `log_prob`s of component distributions
  (which may have different shapes), it first reduces the component `log_prob`s
  to ensure that `jd.log_prob(jd.sample())` always returns a scalar, unless
  `batch_ndims` is explicitly set to a nonzero value (in which case the result
  will have the corresponding tensor rank).
  
  The essential changes are:
  
  - An `event` of `JointDistributionCoroutineAutoBatched` is the list of
    tensors produced by `.sample()`; thus, the `event_shape` is the
    list of the shapes of sampled tensors. These combine both the event
    and batch dimensions of the component distributions. By contrast, the event
    shape of a base `JointDistribution`s does not include batch dimensions of
    component distributions.
  - The `batch_shape` is a global property of the entire model, rather
    than a per-component property as in base `JointDistribution`s.
    The global batch shape must be a prefix of the batch shapes of
    each component; the length of this prefix is specified by an optional
    argument `batch_ndims`. If `batch_ndims` is not specified, the model has
    batch shape `[]`.
  
  #### Examples
  
  A hierarchical model of Poisson log-rates, written using
  `tfd.JointDistributionCoroutineAutoBatched`:
  
  ```python
  tfd = tfp.distributions
  def model():
    global_log_rate = yield tfd.Normal(loc=0., scale=1.)
    local_log_rates = yield tfd.Normal(loc=0., scale=tf.ones([20]))
    observed_counts = yield tfd.Poisson(
      rate=tf.exp(global_log_rate + local_log_rates))
  joint = tfd.JointDistributionCoroutineAutoBatched(model)
  
  print(joint.event_shape)
  # ==> [[], [20], [20]]
  print(joint.batch_shape)
  # ==> []
  xs = joint.sample()
  print([x.shape for x in xs])
  # ==> [[], [20], [20]]
  lp = joint.log_prob(xs)
  print(lp.shape)
  # ==> []
  ```
  
  Note that the component distributions of this model would, by themselves,
  return batches of log-densities (because they are constructed with batch
  shape); the joint model implicitly sums over these to compute the single
  joint log-density.
  
  
  ```python
  ds, xs = joint.sample_distributions()
  print([d.event_shape for d in ds])
  # ==> [[], [], []] != model.event_shape
  print([d.batch_shape for d in ds])
  # ==> [[], [20], [20]] != model.batch_shape
  print([d.log_prob(x).shape for (d, x) in zip(ds, xs)])
  # ==> [[], [20], [20]]
  ```
  
  The behavior of `JointDistributionCoroutineAutoBatched` is (assuming that
  `batch_ndims` is not specified) equivalent to
  adding `tfp.distributions.Independent` wrappers to reinterpret all batch
  dimensions in a `JointDistributionCoroutine` model. That is, the model above
  would be equivalently written using `JointDistributionCoroutine` as:
  
  ```python
  def model_jdc():
    global_log_rate = yield Root(tfd.Normal(0., 1.))
    local_log_rates = yield Root(tfd.Independent(
      tfd.Normal(0., tf.ones([20])), reinterpreted_batch_ndims=1))
    observed_counts = yield Root(tfd.Independent(
      tfd.Poisson(tf.exp(global_log_rate + local_log_rates)),
      reinterpreted_batch_ndims=1))
  joint_jdc = tfd.JointDistributionCoroutine(model_jdc)
  ```
  
  To define a *batch* of joint distributions (independent, but not identical,
  joint distributions from the same family) using
  `JointDistributionCoroutineAutoBatched`, any batch dimensions must be a shared
  prefix of the batch dimensions for all components. The `batch_ndims` argument
  determines the size of the prefix to consider. For example, consider a simple
  joint model with two scalar normal random variables, where the second
  variable's mean is given by the first variable. We can write a batch of five
  such models as:
  
  ```python
  def model():
    x = yield tfd.Normal(0., scale=tf.ones([5]))
    y = yield tfd.Normal(x, scale=[3., 2., 5., 1., 6.])
  batch_joint = tfd.JointDistributionCoroutineAutoBatched(model, batch_ndims=1)
  
  print(batch_joint.event_shape)
  # ==> [[], []]
  print(batch_joint.batch_shape)
  # ==> [5]
  print(batch_joint.log_prob(batch_joint.sample()).shape)
  # ==> [5]
  ```
  
  Note that if we had not passed `batch_ndims`, this would be interpreted as a
  single model over vector-valued random variables (whose components happen to
  be independent):
  
  ```python
  alternate_joint = tfd.JointDistributionCoroutineAutoBatched(model)
  print(alternate_joint.event_shape)
  # ==> [[5], [5]]
  print(alternate_joint.batch_shape)
  # ==> []
  print(alternate_joint.log_prob(batch_joint.sample()).shape)
  # ==> []
  ```
  
  For improved readability of sampled values, the yielded distributions can also
  be named:
  
  ```python
  tfd = tfp.distributions
  def model():
    global_log_rate = yield tfd.Normal(
      loc=0., scale=1., name='global_log_rate')
    local_log_rates = yield tfd.Normal(
      loc=0., scale=tf.ones([20]), name='local_log_rates')
    observed_counts = yield tfd.Poisson(
      rate=tf.exp(global_log_rate + local_log_rates), name='observed_counts')
  joint = tfd.JointDistributionCoroutineAutoBatched(model)
  
  print(joint.event_shape)
  # ==> StructTuple(global_log_rate=[], local_log_rates=[20],
  #      observed_counts=[20])
  print(joint.batch_shape)
  # ==> []
  xs = joint.sample()
  print(['{}: {}'.format(k, x.shape) for k, x in xs._asdict().items()])
  # ==> global_log_scale: []
  #     local_log_rates: [20]
  #     observed_counts: [20]
  lp = joint.log_prob(xs)
  print(lp.shape)
  # ==> []
  
  # Passing via `kwargs` also works.
  lp = joint.log_prob(**xs._asdict())
  # Or:
  lp = joint.log_prob(
      global_log_scale=...,
      local_log_rates=...,
      observed_counts=...,
  )
  ```
  
  If any of the yielded distributions are not explicitly named, they will
  automatically be given a name of the form `var#` where `#` is the index of the
  associated distribution. E.g. the first yielded distribution will have a
  default name of `var0`.
class_docstr_tok: ["Joint","distribution","parameterized","by","a","distribution-making","generator",".","This","class","provides","automatic","vectorization","and","alternative","semantics","for","`","tfd.JointDistributionCoroutine","`",",","which","in","many","cases","allows","for","simplifications","in","the","model","specification",".","#","#","#","#","Automatic","vectorization","Auto-vectorized","variants","of","JointDistribution","allow","the","user","to","avoid","explicitly","annotating","a","model","'s","vectorization","semantics",".","When","using","manually-vectorized","joint","distributions",",","each","operation","in","the","model","must","account","for","the","possibility","of","batch","dimensions","in","Distributions","and","their","samples",".","By","contrast",",","auto-vectorized","models","need","only","describe","a","*","single","*","sample","from","the","joint","distribution",";","any","batch","evaluation","is","automated","using","`","tf.vectorized_map","`","as","required",".","In","many","cases","this","allows","for","significant","simplications",".","For","example",",","the","following","manually-vectorized","`","tfd.JointDistributionCoroutine","`","model",":","``","`","python","def","model_fn","(",")",":","x","=","yield","tfd.JointDistributionCoroutine.Root","(","tfd.Normal","(","0.",",","tf.ones","(","[","3","]",")",")",")","y","=","yield","tfd.JointDistributionCoroutine.Root","(","tfd.Normal","(","0.",",","1",".",")",")","z","=","yield","tfd.Normal","(","x","[","...",",",":2","]","+","y","[","...",",","tf.newaxis","]",",","1",".",")","can","be","written","in","auto-vectorized","form","as","``","`","python","def","model_fn","(",")",":","x","=","yield","tfd.Normal","(","0.",",","tf.ones","(","[","3","]",")",")","y","=","yield","tfd.Normal","(","0.",",","1",".",")","z","=","yield","tfd.Normal","(","x","[",":2","]","+","y",",","1",".",")","``","`","in","which","we","were","able","to","drop","the","specification","of","`","Root","`","nodes","and","to","avoid","explicitly","accounting","for","batch","dimensions","when","indexing","and","slicing","computed","quantities","in","the","third","line",".","Note",":","auto-vectorization","is","still","experimental","and","some","TensorFlow","ops","may","be","unsupported",".","It","can","be","disabled","by","setting","`","use_vectorized_map=False","`",".","#","#","#","#","Alternative","batch","semantics","This","class","also","provides","alternative","semantics","for","specifying","a","batch","of","independent","(","non-identical",")","joint","distributions",".","Instead","of","simply","summing","the","`","log_prob","`","s","of","component","distributions","(","which","may","have","different","shapes",")",",","it","first","reduces","the","component","`","log_prob","`","s","to","ensure","that","`","jd.log_prob","(","jd.sample","(",")",")","`","always","returns","a","scalar",",","unless","`","batch_ndims","`","is","explicitly","set","to","a","nonzero","value","(","in","which","case","the","result","will","have","the","corresponding","tensor","rank",")",".","The","essential","changes","are",":","-","An","`","event","`","of","`","JointDistributionCoroutineAutoBatched","`","is","the","list","of","tensors","produced","by","`",".sample","(",")","`",";","thus",",","the","`","event_shape","`","is","the","list","of","the","shapes","of","sampled","tensors",".","These","combine","both","the","event","and","batch","dimensions","of","the","component","distributions",".","By","contrast",",","the","event","shape","of","a","base","`","JointDistribution","`","s","does","not","include","batch","dimensions","of","component","distributions",".","-","The","`","batch_shape","`","is","a","global","property","of","the","entire","model",",","rather","than","a","per-component","property","as","in","base","`","JointDistribution","`","s",".","The","global","batch","shape","must","be","a","prefix","of","the","batch","shapes","of","each","component",";","the","length","of","this","prefix","is","specified","by","an","optional","argument","`","batch_ndims","`",".","If","`","batch_ndims","`","is","not","specified",",","the","model","has","batch","shape","`","[","]","`",".","#","#","#","#","Examples","A","hierarchical","model","of","Poisson","log-rates",",","written","using","`","tfd.JointDistributionCoroutineAutoBatched","`",":","``","`","python","tfd","=","tfp.distributions","def","model","(",")",":","global_log_rate","=","yield","tfd.Normal","(","loc=0.",",","scale=1",".",")","local_log_rates","=","yield","tfd.Normal","(","loc=0.",",","scale=tf.ones","(","[","20","]",")",")","observed_counts","=","yield","tfd.Poisson","(","rate=tf.exp","(","global_log_rate","+","local_log_rates",")",")","joint","=","tfd.JointDistributionCoroutineAutoBatched","(","model",")","print","(","joint.event_shape",")","#","==",">","[","[","]",",","[","20","]",",","[","20","]","]","print","(","joint.batch_shape",")","#","==",">","[","]","xs","=","joint.sample","(",")","print","(","[","x.shape","for","x","in","xs","]",")","#","==",">","[","[","]",",","[","20","]",",","[","20","]","]","lp","=","joint.log_prob","(","xs",")","print","(","lp.shape",")","#","==",">","[","]","``","`","Note","that","the","component","distributions","of","this","model","would",",","by","themselves",",","return","batches","of","log-densities","(","because","they","are","constructed","with","batch","shape",")",";","the","joint","model","implicitly","sums","over","these","to","compute","the","single","joint","log-density",".","``","`","python","ds",",","xs","=","joint.sample_distributions","(",")","print","(","[","d.event_shape","for","d","in","ds","]",")","#","==",">","[","[","]",",","[","]",",","[","]","]","!","=","model.event_shape","print","(","[","d.batch_shape","for","d","in","ds","]",")","#","==",">","[","[","]",",","[","20","]",",","[","20","]","]","!","=","model.batch_shape","print","(","[","d.log_prob","(","x",")",".shape","for","(","d",",","x",")","in","zip","(","ds",",","xs",")","]",")","#","==",">","[","[","]",",","[","20","]",",","[","20","]","]","``","`","The","behavior","of","`","JointDistributionCoroutineAutoBatched","`","is","(","assuming","that","`","batch_ndims","`","is","not","specified",")","equivalent","to","adding","`","tfp.distributions.Independent","`","wrappers","to","reinterpret","all","batch","dimensions","in","a","`","JointDistributionCoroutine","`","model",".","That","is",",","the","model","above","would","be","equivalently","written","using","`","JointDistributionCoroutine","`","as",":","``","`","python","def","model_jdc","(",")",":","global_log_rate","=","yield","Root","(","tfd.Normal","(","0.",",","1",".",")",")","local_log_rates","=","yield","Root","(","tfd.Independent","(","tfd.Normal","(","0.",",","tf.ones","(","[","20","]",")",")",",","reinterpreted_batch_ndims=1",")",")","observed_counts","=","yield","Root","(","tfd.Independent","(","tfd.Poisson","(","tf.exp","(","global_log_rate","+","local_log_rates",")",")",",","reinterpreted_batch_ndims=1",")",")","joint_jdc","=","tfd.JointDistributionCoroutine","(","model_jdc",")","``","`","To","define","a","*","batch","*","of","joint","distributions","(","independent",",","but","not","identical",",","joint","distributions","from","the","same","family",")","using","`","JointDistributionCoroutineAutoBatched","`",",","any","batch","dimensions","must","be","a","shared","prefix","of","the","batch","dimensions","for","all","components",".","The","`","batch_ndims","`","argument","determines","the","size","of","the","prefix","to","consider",".","For","example",",","consider","a","simple","joint","model","with","two","scalar","normal","random","variables",",","where","the","second","variable","'s","mean","is","given","by","the","first","variable",".","We","can","write","a","batch","of","five","such","models","as",":","``","`","python","def","model","(",")",":","x","=","yield","tfd.Normal","(","0.",",","scale=tf.ones","(","[","5","]",")",")","y","=","yield","tfd.Normal","(","x",",","scale=","[","3.",",","2.",",","5.",",","1.",",","6",".","]",")","batch_joint","=","tfd.JointDistributionCoroutineAutoBatched","(","model",",","batch_ndims=1",")","print","(","batch_joint.event_shape",")","#","==",">","[","[","]",",","[","]","]","print","(","batch_joint.batch_shape",")","#","==",">","[","5","]","print","(","batch_joint.log_prob","(","batch_joint.sample","(",")",")",".shape",")","#","==",">","[","5","]","``","`","Note","that","if","we","had","not","passed","`","batch_ndims","`",",","this","would","be","interpreted","as","a","single","model","over","vector-valued","random","variables","(","whose","components","happen","to","be","independent",")",":","``","`","python","alternate_joint","=","tfd.JointDistributionCoroutineAutoBatched","(","model",")","print","(","alternate_joint.event_shape",")","#","==",">","[","[","5","]",",","[","5","]","]","print","(","alternate_joint.batch_shape",")","#","==",">","[","]","print","(","alternate_joint.log_prob","(","batch_joint.sample","(",")",")",".shape",")","#","==",">","[","]","``","`","For","improved","readability","of","sampled","values",",","the","yielded","distributions","can","also","be","named",":","``","`","python","tfd","=","tfp.distributions","def","model","(",")",":","global_log_rate","=","yield","tfd.Normal","(","loc=0.",",","scale=1.",",","name='global_log_rate","'",")","local_log_rates","=","yield","tfd.Normal","(","loc=0.",",","scale=tf.ones","(","[","20","]",")",",","name='local_log_rates","'",")","observed_counts","=","yield","tfd.Poisson","(","rate=tf.exp","(","global_log_rate","+","local_log_rates",")",",","name='observed_counts","'",")","joint","=","tfd.JointDistributionCoroutineAutoBatched","(","model",")","print","(","joint.event_shape",")","#","==",">","StructTuple","(","global_log_rate=","[","]",",","local_log_rates=","[","20","]",",","#","observed_counts=","[","20","]",")","print","(","joint.batch_shape",")","#","==",">","[","]","xs","=","joint.sample","(",")","print","(","[","'","{","}",":","{","}","'.format","(","k",",","x.shape",")","for","k",",","x","in","xs._asdict","(",")",".items","(",")","]",")","#","==",">","global_log_scale",":","[","]","#","local_log_rates",":","[","20","]","#","observed_counts",":","[","20","]","lp","=","joint.log_prob","(","xs",")","print","(","lp.shape",")","#","==",">","[","]","#","Passing","via","`","kwargs","`","also","works",".","lp","=","joint.log_prob","(","*","*","xs._asdict","(",")",")","#","Or",":","lp","=","joint.log_prob","(","global_log_scale=","...",",","local_log_rates=","...",",","observed_counts=","...",",",")","``","`","If","any","of","the","yielded","distributions","are","not","explicitly","named",",","they","will","automatically","be","given","a","name","of","the","form","`","var","#","`","where","`","#","`","is","the","index","of","the","associated","distribution",".","E.g",".","the","first","yielded","distribution","will","have","a","default","name","of","`","var0","`","."]
docstr: Construct the `JointDistributionCoroutineAutoBatched` distribution.
  
  Args:
    model: A generator that yields a sequence of `tfd.Distribution`-like
      instances.
    sample_dtype: Samples from this distribution will be structured like
      `tf.nest.pack_sequence_as(sample_dtype, list_)`. `sample_dtype` is only
      used for `tf.nest.pack_sequence_as` structuring of outputs, never
      casting (which is the responsibility of the component distributions).
      Default value: `None` (i.e. `namedtuple`).
    batch_ndims: `int` `Tensor` number of batch dimensions. The `batch_shape`s
      of all component distributions must be such that the prefixes of
      length `batch_ndims` broadcast to a consistent joint batch shape.
      Default value: `0`.
    use_vectorized_map: Python `bool`. Whether to use `tf.vectorized_map`
      to automatically vectorize evaluation of the model. This allows the
      model specification to focus on drawing a single sample, which is often
      simpler, but some ops may not be supported.
      Default value: `True`.
    validate_args: Python `bool`.  Whether to validate input with asserts.
      If `validate_args` is `False`, and the inputs are invalid,
      correct behavior is not guaranteed.
      Default value: `False`.
    experimental_use_kahan_sum: Python `bool`. When `True`, we use Kahan
      summation to aggregate independent underlying log_prob values, which
      improves against the precision of a naive float32 sum. This can be
      noticeable in particular for large dimensions in float32. See CPU caveat
      on `tfp.math.reduce_kahan_sum`.
    name: The name for ops managed by the distribution.
      Default value: `None` (i.e., `JointDistributionCoroutine`).
docstr_tok: ["Construct","the","`","JointDistributionCoroutineAutoBatched","`","distribution",".","Args",":","model",":","A","generator","that","yields","a","sequence","of","`","tfd.Distribution","`","-like","instances",".","sample_dtype",":","Samples","from","this","distribution","will","be","structured","like","`","tf.nest.pack_sequence_as","(","sample_dtype",",","list_",")","`",".","`","sample_dtype","`","is","only","used","for","`","tf.nest.pack_sequence_as","`","structuring","of","outputs",",","never","casting","(","which","is","the","responsibility","of","the","component","distributions",")",".","Default","value",":","`","None","`","(","i.e",".","`","namedtuple","`",")",".","batch_ndims",":","`","int","`","`","Tensor","`","number","of","batch","dimensions",".","The","`","batch_shape","`","s","of","all","component","distributions","must","be","such","that","the","prefixes","of","length","`","batch_ndims","`","broadcast","to","a","consistent","joint","batch","shape",".","Default","value",":","`","0","`",".","use_vectorized_map",":","Python","`","bool","`",".","Whether","to","use","`","tf.vectorized_map","`","to","automatically","vectorize","evaluation","of","the","model",".","This","allows","the","model","specification","to","focus","on","drawing","a","single","sample",",","which","is","often","simpler",",","but","some","ops","may","not","be","supported",".","Default","value",":","`","True","`",".","validate_args",":","Python","`","bool","`",".","Whether","to","validate","input","with","asserts",".","If","`","validate_args","`","is","`","False","`",",","and","the","inputs","are","invalid",",","correct","behavior","is","not","guaranteed",".","Default","value",":","`","False","`",".","experimental_use_kahan_sum",":","Python","`","bool","`",".","When","`","True","`",",","we","use","Kahan","summation","to","aggregate","independent","underlying","log_prob","values",",","which","improves","against","the","precision","of","a","naive","float32","sum",".","This","can","be","noticeable","in","particular","for","large","dimensions","in","float32",".","See","CPU","caveat","on","`","tfp.math.reduce_kahan_sum","`",".","name",":","The","name","for","ops","managed","by","the","distribution",".","Default","value",":","`","None","`","(","i.e.",",","`","JointDistributionCoroutine","`",")","."]
code: def __init__(
      self,
      model,
      sample_dtype=None,
      batch_ndims=0,
      use_vectorized_map=True,
      validate_args=False,
      experimental_use_kahan_sum=False,
      name=None,
  ):
      """Construct the `JointDistributionCoroutineAutoBatched` distribution.
  
      Args:
        model: A generator that yields a sequence of `tfd.Distribution`-like
          instances.
        sample_dtype: Samples from this distribution will be structured like
          `tf.nest.pack_sequence_as(sample_dtype, list_)`. `sample_dtype` is only
          used for `tf.nest.pack_sequence_as` structuring of outputs, never
          casting (which is the responsibility of the component distributions).
          Default value: `None` (i.e. `namedtuple`).
        batch_ndims: `int` `Tensor` number of batch dimensions. The `batch_shape`s
          of all component distributions must be such that the prefixes of
          length `batch_ndims` broadcast to a consistent joint batch shape.
          Default value: `0`.
        use_vectorized_map: Python `bool`. Whether to use `tf.vectorized_map`
          to automatically vectorize evaluation of the model. This allows the
          model specification to focus on drawing a single sample, which is often
          simpler, but some ops may not be supported.
          Default value: `True`.
        validate_args: Python `bool`.  Whether to validate input with asserts.
          If `validate_args` is `False`, and the inputs are invalid,
          correct behavior is not guaranteed.
          Default value: `False`.
        experimental_use_kahan_sum: Python `bool`. When `True`, we use Kahan
          summation to aggregate independent underlying log_prob values, which
          improves against the precision of a naive float32 sum. This can be
          noticeable in particular for large dimensions in float32. See CPU caveat
          on `tfp.math.reduce_kahan_sum`.
        name: The name for ops managed by the distribution.
          Default value: `None` (i.e., `JointDistributionCoroutine`).
      """
      parameters = dict(locals())
      super(JointDistributionCoroutineAutoBatched, self).__init__(
          model,
          sample_dtype=sample_dtype,
          batch_ndims=batch_ndims,
          use_vectorized_map=use_vectorized_map,
          validate_args=validate_args,
          experimental_use_kahan_sum=experimental_use_kahan_sum,
          name=name or "JointDistributionCoroutineAutoBatched",
      )
      self._parameters = self._no_dependency(parameters)
  
imports: import functools
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import joint_distribution_coroutine
  from tensorflow_probability.python.distributions import joint_distribution_named
  from tensorflow_probability.python.distributions import joint_distribution_sequential
  from tensorflow_probability.python.internal import auto_composite_tensor
label_desc: Use image node_id 1 to create a new JointDistributionCoroutineAutoBatched object from inherited base classes: joint_distribution_coroutine with example: obj = JointDistributionCoroutineAutoBatched(model, sample_dtype, batch_ndims, use_vectorized_map, validate_args, experimental_use_kahan_sum, name)