class_docstr: Joint distribution parameterized by distribution-making functions.
  
  This distribution enables both sampling and joint probability computation from
  a single model specification.
  
  A joint distribution is a collection of possibly interdependent distributions.
  Like `tf_keras.Sequential`, the `JointDistributionSequential` can be specified
  via a `list` of functions (each responsible for making a
  `tfp.distributions.Distribution`-like instance).  Unlike
  `tf_keras.Sequential`, each function can depend on the output of all previous
  elements rather than only the immediately previous.
  
  #### Mathematical Details
  
  The `JointDistributionSequential` implements the chain rule of probability.
  That is, the probability function of a length-`d` vector `x` is,
  
  ```none
  p(x) = prod{ p(x[i] | x[:i]) : i = 0, ..., (d - 1) }
  ```
  
  The `JointDistributionSequential` is parameterized by a `list` comprised of
  either:
  
  1. `tfp.distributions.Distribution`-like instances or,
  2. `callable`s which return a `tfp.distributions.Distribution`-like instance.
  
  Each `list` element implements the `i`-th *full conditional distribution*,
  `p(x[i] | x[:i])`. The "conditioned on" elements are represented by the
  `callable`'s required arguments. Directly providing a `Distribution`-like
  instance is a convenience and is semantically identical a zero argument
  `callable`.
  
  Denote the `i`-th `callable`s non-default arguments as `args[i]`. Since the
  `callable` is the conditional manifest, `0 <= len(args[i]) <= i - 1`. When
  `len(args[i]) < i - 1`, the `callable` only depends on a subset of the
  previous distributions, specifically those at indexes:
  `range(i - 1, i - 1 - num_args[i], -1)`.
  (See "Examples" and "Discussion" for why the order is reversed.)
  
  **Name resolution**: `The names of `JointDistributionSequential` components
  are defined by explicit `name` arguments passed to distributions
  (`tfd.Normal(0., 1., name='x')`) and/or by the argument names in
  distribution-making functions (`lambda x: tfd.Normal(x., 1.)`). Both
  approaches may be used in the same distribution, as long as they are
  consistent; referring to a single component by multiple names will raise a
  `ValueError`. Unnamed components will be assigned a dummy name.
  
  #### Examples
  
  Consider the following generative model:
  
  ```
  e ~ Exponential(rate=[100,120])
  g ~ Gamma(concentration=e[0], rate=e[1])
  n ~ Normal(loc=0, scale=2.)
  m ~ Normal(loc=n, scale=g)
  for i = 1, ..., 12:
    x[i] ~ Bernoulli(logits=m)
  ```
  
  We can code this as:
  
  ```python
  tfd = tfp.distributions
  joint = tfd.JointDistributionSequential([
                   tfd.Exponential(rate=[100, 120]),           # e
      lambda    e: tfd.Gamma(concentration=e[0], rate=e[1]),    # g
                   tfd.Normal(loc=0, scale=2.),                 # n
      lambda n, g: tfd.Normal(loc=n, scale=g),                  # m
      lambda    m: tfd.Sample(tfd.Bernoulli(logits=m), 12)      # x
  ], batch_ndims=0, use_vectorized_map=True)
  ```
  
  Notice the 1:1 correspondence between "math" and "code".
  
  ```python
  x = joint.sample()
  # ==> A length-5 list of Tensors representing a draw/realization from each
  #     distribution.
  joint.log_prob(x)
  # ==> A scalar `Tensor` representing the total log prob under all five
  #     distributions.
  
  joint.resolve_graph()
  # ==> (('e', ()),
  #      ('g', ('e',)),
  #      ('n', ()),
  #      ('m', ('n', 'g')),
  #      ('x', ('m',)))
  ```
  
  #### Discussion
  
  `JointDistributionSequential` builds each distribution in `list` order; list
  items must be either a:
  
  1. `tfd.Distribution`-like instance (e.g., `e` and `n`), or a
  2. Python `callable` (e.g., `g`, `m`, `x`).
  
  Regarding #1, an object is deemed "`tfd.Distribution`-like" if it has a
  `sample`, `log_prob`, and distribution properties, e.g., `batch_shape`,
  `event_shape`, `dtype`.
  
  Regarding #2, in addition to using a function (or `lambda`), supplying a TFD
  "`class`" is also permissible, this also being a "Python `callable`." For
  example, instead of writing:
  `lambda loc, scale: tfd.Normal(loc=loc, scale=scale)`
  one could have simply written `tfd.Normal`.
  
  Notice that directly providing a `tfd.Distribution`-like instance means there
  cannot exist a (dynamic) dependency on other distributions; it is
  "independent" both "computationally" and "statistically." The same is
  self-evidently true of zero-argument `callable`s.
  
  A distribution instance depends on other distribution instances through the
  distribution making function's *required arguments*. If the distribution maker
  has `k` required arguments then the `JointDistributionSequential` calls the
  maker with samples produced by the previous `k` distributions.
  
  **Note**: **maker arguments are provided in reverse order** of the previous
  elements in the list. In the example, notice that `m` depends on `n` and `g`
  in this order. The order is reversed for convenience. We reverse arguments
  under the heuristic observation that many graphical models have chain-like
  dependencies which are self-evidently topologically sorted from the human
  cognitive process of postulating the generative process. By feeding the
  previous num required args in reverse order we (often) enable a simpler maker
  function signature. If the maker needs to depend on distribution previous to
  one which is not a dependency, one must use a dummy arg, to "gobble up" the
  unused dependency, e.g.,
  `lambda _ignore, actual_dependency: SomeDistribution(actual_dependency)`.
  
  **Note**: unlike other non-`JointDistribution` distributions in
  `tfp.distributions`, `JointDistribution.sample` (and subclasses) return a
  structure of  `Tensor`s rather than a `Tensor`.  A structure can be anything
  which is `list`-like, e.g., a `list` or `tuple` of `distribution` makers.
  Accordingly `joint.event_shape` returns a `list`-like structure of
  `TensorShape`s for each of the distributions' event shapes and
  `joint.event_shape_tensor()` returns a `list`-like structure of `Tensor`s for
  each of the distributions' event shapes.
  
  **Note**: If a `JointDistributionSequential` instance contains a `callable`
  that closes over a `Tensor`, the `JointDistributionSequential` cannot cross
  the boundary of a `tf.function`. (If this behavior is necessary, an instance
  of `_JointDistributionSequential` may be used instead, at the expense of extra
  `tf.function` retracing.)
  
  
  #### Vectorized sampling and model evaluation
  
  When a joint distribution's `sample` method  is called with
  a `sample_shape` (or the `log_prob` method is called on an input with
  multiple sample dimensions) the model must be equipped to handle
  additional batch dimensions. This may be done manually, or automatically
  by passing `use_vectorized_map=True`. Manual vectorization has historically
  been the default, but we now recommend that most users enable automatic
  vectorization unless they are affected by a specific issue; some
  known issues are listed below.
  
  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 as required using `tf.vectorized_map` (`vmap` in JAX). In many
  cases this allows for significant simplications. For example, the following
  manually-vectorized `tfd.JointDistributionSequential` model:
  
  ```python
  model = tfd.JointDistributionSequential([
      tfd.Normal(0., tf.ones([3])),
      tfd.Normal(0., 1.),
      lambda y, x: tfd.Normal(x[..., :2] + y[..., tf.newaxis], 1.)
    ])
  ```
  
  can be written in auto-vectorized form as
  
  ```python
  model = tfd.JointDistributionSequential([
      tfd.Normal(0., tf.ones([3])),
      tfd.Normal(0., 1.),
      lambda y, x: tfd.Normal(x[:2] + y, 1.)
    ],
    use_vectorized_map=True)
  ```
  
  in which we were able to avoid explicitly accounting for batch dimensions
  when indexing and slicing computed quantities in the third line.
  
  **Known limitations of automatic vectorization:**
  - A small fraction of TensorFlow ops are unsupported; models that use an
    unsupported op will raise an error and must be manually vectorized.
  - Sampling large batches may be slow under automatic vectorization because
    TensorFlow's stateless samplers are currently converted using a
    non-vectorized `while_loop`. This limitation applies only in TensorFlow;
    vectorized samplers in JAX should be approximately as fast as manually
    vectorized code.
  - Calling `sample_distributions` with nontrivial `sample_shape` will raise
    an error if the model contains any distributions that are not registered as
    CompositeTensors (TFP's basic distributions are usually fine, but support
    for wrapper distributions like `tfd.Sample` is a work in progress).
  
  #### Batch semantics and (log-)densities
  
  **tl;dr:** pass `batch_ndims=0` unless you have a good reason not to.
  
  Joint distributions now support 'auto-batching' semantics, in which
  the distribution's batch shape is derived by broadcasting the leftmost
  `batch_ndims` dimensions of its components' batch shapes. All remaining
  dimensions are considered to form a single 'event' of the joint distribution.
  If `batch_ndims==0`, then the joint distribution has batch shape `[]`, and all
  component dimensions are treated as event shape. For example, the model
  
  ```python
  jd = tfd.JointDistributionSequential([
      tfd.Normal(0., tf.ones([3])),
      lambda x: tfd.Normal(x[..., tf.newaxis], tf.ones([3, 2]))
    ],
    batch_ndims=0)
  ```
  
  creates a joint distribution with batch shape `[]` and event shape
  `([3], [3, 2])`. The log-density of a sample always has shape
  `batch_shape`, so this guarantees that
  `jd.log_prob(jd.sample())` will evaluate to a scalar value. We could
  alternately construct a joint distribution with batch shape `[3]` and event
  shape `([], [2])` by setting `batch_ndims=1`, in which case
  `jd.log_prob(jd.sample())` would evaluate to a value of shape `[3]`.
  
  Setting `batch_ndims=None` recovers the 'classic' batch semantics (currently
  still the default for backwards-compatibility reasons), in which the joint
  distribution's `log_prob` is computed by naively summing log densities from
  the component distributions. Since these component densities have shapes equal
  to the batch shapes of the individual components, to avoid broadcasting
  errors it is usually necessary to construct the components with identical
  batch shapes. For example, the component distributions in the model above
  have batch shapes of `[3]` and `[3, 2]` respectively, which would raise an
  error if summed directly, but can be aligned by wrapping with
  `tfd.Independent`, as in this model:
  
  ```python
  jd = tfd.JointDistributionSequential([
      tfd.Normal(0., tf.ones([3])),
      lambda x: tfd.Independent(tfd.Normal(x[..., tf.newaxis], tf.ones([3, 2])),
                                reinterpreted_batch_ndims=1)
    ],
    batch_ndims=None)
  ```
  
  Here the components both have batch shape `[3]`, so
  `jd.log_prob(jd.sample())` returns a value of shape `[3]`, just as in the
  `batch_ndims=1` case above. In fact, auto-batching semantics are equivalent to
  implicitly wrapping each component `dist` as `tfd.Independent(dist,
  reinterpreted_batch_ndim=(dist.batch_shape.ndims - jd.batch_ndims))`; the only
  vestigial difference is that under auto-batching semantics, the joint
  distribution has a single batch shape `[3]`, while under the classic semantics
  the value of `jd.batch_shape` is a *structure* of the component batch shapes
  `([3], [3])`. Such structured batch shapes will be deprecated in the future,
  since they are inconsistent with the definition of batch shapes used
  elsewhere in TFP.
  
  #### References
  
  [1] Dan Piponi, Dave Moore, and Joshua V. Dillon. Joint distributions for
      TensorFlow Probability. _arXiv preprint arXiv:2001.11819__,
      2020. https://arxiv.org/abs/2001.11819
class_docstr_tok: ["Joint","distribution","parameterized","by","distribution-making","functions",".","This","distribution","enables","both","sampling","and","joint","probability","computation","from","a","single","model","specification",".","A","joint","distribution","is","a","collection","of","possibly","interdependent","distributions",".","Like","`","tf_keras.Sequential","`",",","the","`","JointDistributionSequential","`","can","be","specified","via","a","`","list","`","of","functions","(","each","responsible","for","making","a","`","tfp.distributions.Distribution","`","-like","instance",")",".","Unlike","`","tf_keras.Sequential","`",",","each","function","can","depend","on","the","output","of","all","previous","elements","rather","than","only","the","immediately","previous",".","#","#","#","#","Mathematical","Details","The","`","JointDistributionSequential","`","implements","the","chain","rule","of","probability",".","That","is",",","the","probability","function","of","a","length-","`","d","`","vector","`","x","`","is",",","``","`","none","p","(","x",")","=","prod","{","p","(","x","[","i","]","|","x","[",":","i","]",")",":","i","=","0",",","...",",","(","d","-","1",")","}","``","`","The","`","JointDistributionSequential","`","is","parameterized","by","a","`","list","`","comprised","of","either",":","1",".","`","tfp.distributions.Distribution","`","-like","instances","or",",","2",".","`","callable","`","s","which","return","a","`","tfp.distributions.Distribution","`","-like","instance",".","Each","`","list","`","element","implements","the","`","i","`","-th","*","full","conditional","distribution","*",",","`","p","(","x","[","i","]","|","x","[",":","i","]",")","`",".","The","``","conditioned","on","''","elements","are","represented","by","the","`","callable","`","'s","required","arguments",".","Directly","providing","a","`","Distribution","`","-like","instance","is","a","convenience","and","is","semantically","identical","a","zero","argument","`","callable","`",".","Denote","the","`","i","`","-th","`","callable","`","s","non-default","arguments","as","`","args","[","i","]","`",".","Since","the","`","callable","`","is","the","conditional","manifest",",","`","0","<","=","len","(","args","[","i","]",")","<","=","i","-","1","`",".","When","`","len","(","args","[","i","]",")","<","i","-","1","`",",","the","`","callable","`","only","depends","on","a","subset","of","the","previous","distributions",",","specifically","those","at","indexes",":","`","range","(","i","-","1",",","i","-","1","-","num_args","[","i","]",",","-1",")","`",".","(","See","``","Examples","''","and","``","Discussion","''","for","why","the","order","is","reversed",".",")","*","*","Name","resolution","*","*",":","`","The","names","of","`","JointDistributionSequential","`","components","are","defined","by","explicit","`","name","`","arguments","passed","to","distributions","(","`","tfd.Normal","(","0.",",","1.",",","name=","'","x","'",")","`",")","and\/or","by","the","argument","names","in","distribution-making","functions","(","`","lambda","x",":","tfd.Normal","(","x.",",","1.",")","`",")",".","Both","approaches","may","be","used","in","the","same","distribution",",","as","long","as","they","are","consistent",";","referring","to","a","single","component","by","multiple","names","will","raise","a","`","ValueError","`",".","Unnamed","components","will","be","assigned","a","dummy","name",".","#","#","#","#","Examples","Consider","the","following","generative","model",":","``","`","e","~","Exponential","(","rate=","[","100,120","]",")","g","~","Gamma","(","concentration=e","[","0","]",",","rate=e","[","1","]",")","n","~","Normal","(","loc=0",",","scale=2",".",")","m","~","Normal","(","loc=n",",","scale=g",")","for","i","=","1",",","...",",","12",":","x","[","i","]","~","Bernoulli","(","logits=m",")","``","`","We","can","code","this","as",":","``","`","python","tfd","=","tfp.distributions","joint","=","tfd.JointDistributionSequential","(","[","tfd.Exponential","(","rate=","[","100",",","120","]",")",",","#","e","lambda","e",":","tfd.Gamma","(","concentration=e","[","0","]",",","rate=e","[","1","]",")",",","#","g","tfd.Normal","(","loc=0",",","scale=2",".",")",",","#","n","lambda","n",",","g",":","tfd.Normal","(","loc=n",",","scale=g",")",",","#","m","lambda","m",":","tfd.Sample","(","tfd.Bernoulli","(","logits=m",")",",","12",")","#","x","]",",","batch_ndims=0",",","use_vectorized_map=True",")","``","`","Notice","the","1:1","correspondence","between","``","math","''","and","``","code","''",".","``","`","python","x","=","joint.sample","(",")","#","==",">","A","length-5","list","of","Tensors","representing","a","draw\/realization","from","each","#","distribution",".","joint.log_prob","(","x",")","#","==",">","A","scalar","`","Tensor","`","representing","the","total","log","prob","under","all","five","#","distributions",".","joint.resolve_graph","(",")","#","==",">","(","(","'","e","'",",","(",")",")",",","#","(","'","g","'",",","(","'","e","'",",",")",")",",","#","(","'n","'",",","(",")",")",",","#","(","'m","'",",","(","'n","'",",","'","g","'",")",")",",","#","(","'","x","'",",","(","'m","'",",",")",")",")","``","`","#","#","#","#","Discussion","`","JointDistributionSequential","`","builds","each","distribution","in","`","list","`","order",";","list","items","must","be","either","a",":","1",".","`","tfd.Distribution","`","-like","instance","(","e.g.",",","`","e","`","and","`","n","`",")",",","or","a","2",".","Python","`","callable","`","(","e.g.",",","`","g","`",",","`","m","`",",","`","x","`",")",".","Regarding","#","1",",","an","object","is","deemed","``","`","tfd.Distribution","`","-like","''","if","it","has","a","`","sample","`",",","`","log_prob","`",",","and","distribution","properties",",","e.g.",",","`","batch_shape","`",",","`","event_shape","`",",","`","dtype","`",".","Regarding","#","2",",","in","addition","to","using","a","function","(","or","`","lambda","`",")",",","supplying","a","TFD","''","`","class","`","``","is","also","permissible",",","this","also","being","a","``","Python","`","callable","`",".","''","For","example",",","instead","of","writing",":","`","lambda","loc",",","scale",":","tfd.Normal","(","loc=loc",",","scale=scale",")","`","one","could","have","simply","written","`","tfd.Normal","`",".","Notice","that","directly","providing","a","`","tfd.Distribution","`","-like","instance","means","there","can","not","exist","a","(","dynamic",")","dependency","on","other","distributions",";","it","is","''","independent","''","both","``","computationally","''","and","``","statistically",".","''","The","same","is","self-evidently","true","of","zero-argument","`","callable","`","s",".","A","distribution","instance","depends","on","other","distribution","instances","through","the","distribution","making","function","'s","*","required","arguments","*",".","If","the","distribution","maker","has","`","k","`","required","arguments","then","the","`","JointDistributionSequential","`","calls","the","maker","with","samples","produced","by","the","previous","`","k","`","distributions",".","*","*","Note","*","*",":","*","*","maker","arguments","are","provided","in","reverse","order","*","*","of","the","previous","elements","in","the","list",".","In","the","example",",","notice","that","`","m","`","depends","on","`","n","`","and","`","g","`","in","this","order",".","The","order","is","reversed","for","convenience",".","We","reverse","arguments","under","the","heuristic","observation","that","many","graphical","models","have","chain-like","dependencies","which","are","self-evidently","topologically","sorted","from","the","human","cognitive","process","of","postulating","the","generative","process",".","By","feeding","the","previous","num","required","args","in","reverse","order","we","(","often",")","enable","a","simpler","maker","function","signature",".","If","the","maker","needs","to","depend","on","distribution","previous","to","one","which","is","not","a","dependency",",","one","must","use","a","dummy","arg",",","to","``","gobble","up","''","the","unused","dependency",",","e.g.",",","`","lambda","_ignore",",","actual_dependency",":","SomeDistribution","(","actual_dependency",")","`",".","*","*","Note","*","*",":","unlike","other","non-","`","JointDistribution","`","distributions","in","`","tfp.distributions","`",",","`","JointDistribution.sample","`","(","and","subclasses",")","return","a","structure","of","`","Tensor","`","s","rather","than","a","`","Tensor","`",".","A","structure","can","be","anything","which","is","`","list","`","-like",",","e.g.",",","a","`","list","`","or","`","tuple","`","of","`","distribution","`","makers",".","Accordingly","`","joint.event_shape","`","returns","a","`","list","`","-like","structure","of","`","TensorShape","`","s","for","each","of","the","distributions","'","event","shapes","and","`","joint.event_shape_tensor","(",")","`","returns","a","`","list","`","-like","structure","of","`","Tensor","`","s","for","each","of","the","distributions","'","event","shapes",".","*","*","Note","*","*",":","If","a","`","JointDistributionSequential","`","instance","contains","a","`","callable","`","that","closes","over","a","`","Tensor","`",",","the","`","JointDistributionSequential","`","can","not","cross","the","boundary","of","a","`","tf.function","`",".","(","If","this","behavior","is","necessary",",","an","instance","of","`","_JointDistributionSequential","`","may","be","used","instead",",","at","the","expense","of","extra","`","tf.function","`","retracing",".",")","#","#","#","#","Vectorized","sampling","and","model","evaluation","When","a","joint","distribution","'s","`","sample","`","method","is","called","with","a","`","sample_shape","`","(","or","the","`","log_prob","`","method","is","called","on","an","input","with","multiple","sample","dimensions",")","the","model","must","be","equipped","to","handle","additional","batch","dimensions",".","This","may","be","done","manually",",","or","automatically","by","passing","`","use_vectorized_map=True","`",".","Manual","vectorization","has","historically","been","the","default",",","but","we","now","recommend","that","most","users","enable","automatic","vectorization","unless","they","are","affected","by","a","specific","issue",";","some","known","issues","are","listed","below",".","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","as","required","using","`","tf.vectorized_map","`","(","`","vmap","`","in","JAX",")",".","In","many","cases","this","allows","for","significant","simplications",".","For","example",",","the","following","manually-vectorized","`","tfd.JointDistributionSequential","`","model",":","``","`","python","model","=","tfd.JointDistributionSequential","(","[","tfd.Normal","(","0.",",","tf.ones","(","[","3","]",")",")",",","tfd.Normal","(","0.",",","1",".",")",",","lambda","y",",","x",":","tfd.Normal","(","x","[","...",",",":2","]","+","y","[","...",",","tf.newaxis","]",",","1",".",")","]",")","``","`","can","be","written","in","auto-vectorized","form","as","``","`","python","model","=","tfd.JointDistributionSequential","(","[","tfd.Normal","(","0.",",","tf.ones","(","[","3","]",")",")",",","tfd.Normal","(","0.",",","1",".",")",",","lambda","y",",","x",":","tfd.Normal","(","x","[",":2","]","+","y",",","1",".",")","]",",","use_vectorized_map=True",")","``","`","in","which","we","were","able","to","avoid","explicitly","accounting","for","batch","dimensions","when","indexing","and","slicing","computed","quantities","in","the","third","line",".","*","*","Known","limitations","of","automatic","vectorization",":","*","*","-","A","small","fraction","of","TensorFlow","ops","are","unsupported",";","models","that","use","an","unsupported","op","will","raise","an","error","and","must","be","manually","vectorized",".","-","Sampling","large","batches","may","be","slow","under","automatic","vectorization","because","TensorFlow","'s","stateless","samplers","are","currently","converted","using","a","non-vectorized","`","while_loop","`",".","This","limitation","applies","only","in","TensorFlow",";","vectorized","samplers","in","JAX","should","be","approximately","as","fast","as","manually","vectorized","code",".","-","Calling","`","sample_distributions","`","with","nontrivial","`","sample_shape","`","will","raise","an","error","if","the","model","contains","any","distributions","that","are","not","registered","as","CompositeTensors","(","TFP","'s","basic","distributions","are","usually","fine",",","but","support","for","wrapper","distributions","like","`","tfd.Sample","`","is","a","work","in","progress",")",".","#","#","#","#","Batch","semantics","and","(","log-",")","densities","*","*","tl",";","dr",":","*","*","pass","`","batch_ndims=0","`","unless","you","have","a","good","reason","not","to",".","Joint","distributions","now","support","'auto-batching","'","semantics",",","in","which","the","distribution","'s","batch","shape","is","derived","by","broadcasting","the","leftmost","`","batch_ndims","`","dimensions","of","its","components","'","batch","shapes",".","All","remaining","dimensions","are","considered","to","form","a","single","'event","'","of","the","joint","distribution",".","If","`","batch_ndims==0","`",",","then","the","joint","distribution","has","batch","shape","`","[","]","`",",","and","all","component","dimensions","are","treated","as","event","shape",".","For","example",",","the","model","``","`","python","jd","=","tfd.JointDistributionSequential","(","[","tfd.Normal","(","0.",",","tf.ones","(","[","3","]",")",")",",","lambda","x",":","tfd.Normal","(","x","[","...",",","tf.newaxis","]",",","tf.ones","(","[","3",",","2","]",")",")","]",",","batch_ndims=0",")","``","`","creates","a","joint","distribution","with","batch","shape","`","[","]","`","and","event","shape","`","(","[","3","]",",","[","3",",","2","]",")","`",".","The","log-density","of","a","sample","always","has","shape","`","batch_shape","`",",","so","this","guarantees","that","`","jd.log_prob","(","jd.sample","(",")",")","`","will","evaluate","to","a","scalar","value",".","We","could","alternately","construct","a","joint","distribution","with","batch","shape","`","[","3","]","`","and","event","shape","`","(","[","]",",","[","2","]",")","`","by","setting","`","batch_ndims=1","`",",","in","which","case","`","jd.log_prob","(","jd.sample","(",")",")","`","would","evaluate","to","a","value","of","shape","`","[","3","]","`",".","Setting","`","batch_ndims=None","`","recovers","the","'classic","'","batch","semantics","(","currently","still","the","default","for","backwards-compatibility","reasons",")",",","in","which","the","joint","distribution","'s","`","log_prob","`","is","computed","by","naively","summing","log","densities","from","the","component","distributions",".","Since","these","component","densities","have","shapes","equal","to","the","batch","shapes","of","the","individual","components",",","to","avoid","broadcasting","errors","it","is","usually","necessary","to","construct","the","components","with","identical","batch","shapes",".","For","example",",","the","component","distributions","in","the","model","above","have","batch","shapes","of","`","[","3","]","`","and","`","[","3",",","2","]","`","respectively",",","which","would","raise","an","error","if","summed","directly",",","but","can","be","aligned","by","wrapping","with","`","tfd.Independent","`",",","as","in","this","model",":","``","`","python","jd","=","tfd.JointDistributionSequential","(","[","tfd.Normal","(","0.",",","tf.ones","(","[","3","]",")",")",",","lambda","x",":","tfd.Independent","(","tfd.Normal","(","x","[","...",",","tf.newaxis","]",",","tf.ones","(","[","3",",","2","]",")",")",",","reinterpreted_batch_ndims=1",")","]",",","batch_ndims=None",")","``","`","Here","the","components","both","have","batch","shape","`","[","3","]","`",",","so","`","jd.log_prob","(","jd.sample","(",")",")","`","returns","a","value","of","shape","`","[","3","]","`",",","just","as","in","the","`","batch_ndims=1","`","case","above",".","In","fact",",","auto-batching","semantics","are","equivalent","to","implicitly","wrapping","each","component","`","dist","`","as","`","tfd.Independent","(","dist",",","reinterpreted_batch_ndim=","(","dist.batch_shape.ndims","-","jd.batch_ndims",")",")","`",";","the","only","vestigial","difference","is","that","under","auto-batching","semantics",",","the","joint","distribution","has","a","single","batch","shape","`","[","3","]","`",",","while","under","the","classic","semantics","the","value","of","`","jd.batch_shape","`","is","a","*","structure","*","of","the","component","batch","shapes","`","(","[","3","]",",","[","3","]",")","`",".","Such","structured","batch","shapes","will","be","deprecated","in","the","future",",","since","they","are","inconsistent","with","the","definition","of","batch","shapes","used","elsewhere","in","TFP",".","#","#","#","#","References","[","1","]","Dan","Piponi",",","Dave","Moore",",","and","Joshua","V.","Dillon",".","Joint","distributions","for","TensorFlow","Probability",".","_arXiv","preprint","arXiv:2001.11819__",",","2020.","https",":","\/\/arxiv.org\/abs\/2001.11819"]
code: def model(self):
      return self._model
  
imports: import collections
  import functools
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import joint_distribution
  from tensorflow_probability.python.distributions import kullback_leibler
  from tensorflow_probability.python.internal import auto_composite_tensor
  from tensorflow_probability.python.internal import distribution_util
  from tensorflow.python.util import tf_inspect
label_desc: Use image node_id 2 for calling the _JointDistributionSequential obj's underlying member method code with example usage: obj.model() and returns: self