class_docstr: A wrapper class for `JointDistribution` which pins, e.g., the evidence.
  
  This object is experimental; the API may change without warning.
  
  Think of this object as `functools.partial` for joint distributions. Sampling
  trims off pinned values (after specifying them as `jd.sample(value=pins)` to
  the underlying distribution). Log-density evaluates the joint probability of
  the given event and the pinned values.
  
  Mathematically speaking, the object represents a joint probability density,
  `p(x, y)` where the `x` are pinned and the `y` are unpinned. Accordingly, it
  is also proportional to `p(y | x)`, up to a (generally) intractable
  normalizing constant `p(x)`, i.e. `p(x, y) = p(y | x) p(x)`.
  
  The measure on the unpinned values `y` that this object represents
  is no longer a probability measure, in that the integral of
  `log_prob` over the space of `y` is `p(x)`, which in general is not
  1.  As such, this is not a `tfp.distributions.Distribution`, and
  lacks a `sample` method. In its place, it provides:
  
  * `sample_unpinned`, `sample_and_log_weight`
  
  A common use-case with probabilistic inference is writing out a generative
  model to explain some observed data:
  
  ```python
  jd = tfd.JointDistributionNamed(dict(
    loc = tfd.Normal(0., 1.),
    scale = tfd.Gamma(1., 1.),
    obs = lambda loc, scale: tfd.Normal(loc, scale),
  ))
  ```
  
  Later, when we want to infer 'typical' values of `loc` and `scale` conditioned
  on some given `data`, we will often write:
  
  ```python
  def target_log_prob_fn(loc, scale):
    return jd.log_prob(loc=loc, scale=scale, obs=data)
  ```
  
  This class enables one to write instead:
  
  ```python
  partial = tfde.JointDistributionPinned(jd, obs=data)
  target_log_prob_fn = partial.log_prob
  ```
  
  Or, even more concisely `partial = jd.experimental_pin(obs=data)`.
  
  This is nice, but it wasn't too hard to write out the `target_log_prob_fn`
  function explicitly.
  
  Now, let's consider that for many inference and optimization methods, we may
  want to use a smooth change of variables to perform inference in the
  unconstrained space of real numbers. In some cases this transformation can be
  parameter-dependent. For example, if we want to unconstrain the support of
  `tfp.distributions.Uniform(-3., 2.)` to the real line, we might use
  `tfp.bijectors.Sigmoid(low=-3., high=2.)`. In support of such use cases,
  most distributions (including the `JointDistribution*` classes) provide a
  `experimental_default_event_space_bijector()` method.
  
  When these transformations may be dependent on ancestral parts of a joint
  distribution, and some of those parameters may be pinned, it is helpful to
  have a utility class to bridge the gap and provide the multi-part bijective
  transform. This is the "raison d'être" of this class.
  
  The model below is somewhat contrived, but demonstrates the use-case.
  
  ```python
  tfd = tfp.distributions
  tfde = tfp.experimental.distributions
  
  n = 75
  dim = 3
  joint = tfd.JointDistributionNamed(dict(
    upper = tfd.Uniform(.4, 1.5),
    concentration = tfd.Gamma(1., .5),
    corr = lambda concentration: tfd.CholeskyLKJ(
        dim, concentration=concentration),
    stddev = lambda upper: tfd.Sample(tfd.Uniform(.2, upper), dim),
    obs = lambda corr, stddev: tfd.Sample(
        tfd.MultivariateNormalTriL(
            loc=tf.zeros([dim]), scale_tril=corr * stddev[..., tf.newaxis]),
        n)
  ))
  fixed_upper = 1.3
  data = joint.sample(upper=fixed_upper)['obs']
  
  pinned = tfde.JointDistributionPinned(joint, upper=fixed_upper, obs=data)
  bij = pinned.experimental_default_event_space_bijector()
  pulled_back_shape = bij.inverse_event_shape(pinned.event_shape)
  
  # Fit an ensemble using SGD.
  batch = 16
  uniform_init = tf.nest.map_structure(
      lambda s: tf.random.uniform(tf.concat([[batch], s], axis=0), -2., 2.),
      pulled_back_shape)
  vars = tf.nest.map_structure(tf.Variable, uniform_init)
  
  opt = tf_keras.optimizers.Adam(.01)
  
  @tf.function(autograph=False)
  def one_step():
    with tf.GradientTape() as tape:
      lp = pinned.log_prob(bij.forward(vars))
    gradients = tape.gradient(lp, vars)
    opt.apply_gradients(zip(gradients.values(), vars.values()))
  
  for _ in range(100):
    one_step()
  
  # Alternatively, sample using MCMC (currently aspirational):
  initial_state = bij.forward(uniform_init)
  
  kernel = tfp.mcmc.HamiltonianMonteCarlo(
      target_log_prob_fn=pinned.log_prob,
      step_size=.5, num_leapfrog_steps=4)
  # **This line is currently aspirational**, to demonstrate the use-case.
  kernel = tfp.mcmc.TransformedTransitionKernel(kernel, bij)
  tfp.mcmc.sample_chain(10, kernel=kernel, current_state=initial_state)
  ```
class_docstr_tok: ["A","wrapper","class","for","`","JointDistribution","`","which","pins",",","e.g.",",","the","evidence",".","This","object","is","experimental",";","the","API","may","change","without","warning",".","Think","of","this","object","as","`","functools.partial","`","for","joint","distributions",".","Sampling","trims","off","pinned","values","(","after","specifying","them","as","`","jd.sample","(","value=pins",")","`","to","the","underlying","distribution",")",".","Log-density","evaluates","the","joint","probability","of","the","given","event","and","the","pinned","values",".","Mathematically","speaking",",","the","object","represents","a","joint","probability","density",",","`","p","(","x",",","y",")","`","where","the","`","x","`","are","pinned","and","the","`","y","`","are","unpinned",".","Accordingly",",","it","is","also","proportional","to","`","p","(","y","|","x",")","`",",","up","to","a","(","generally",")","intractable","normalizing","constant","`","p","(","x",")","`",",","i.e",".","`","p","(","x",",","y",")","=","p","(","y","|","x",")","p","(","x",")","`",".","The","measure","on","the","unpinned","values","`","y","`","that","this","object","represents","is","no","longer","a","probability","measure",",","in","that","the","integral","of","`","log_prob","`","over","the","space","of","`","y","`","is","`","p","(","x",")","`",",","which","in","general","is","not","1",".","As","such",",","this","is","not","a","`","tfp.distributions.Distribution","`",",","and","lacks","a","`","sample","`","method",".","In","its","place",",","it","provides",":","*","`","sample_unpinned","`",",","`","sample_and_log_weight","`","A","common","use-case","with","probabilistic","inference","is","writing","out","a","generative","model","to","explain","some","observed","data",":","``","`","python","jd","=","tfd.JointDistributionNamed","(","dict","(","loc","=","tfd.Normal","(","0.",",","1",".",")",",","scale","=","tfd.Gamma","(","1.",",","1",".",")",",","obs","=","lambda","loc",",","scale",":","tfd.Normal","(","loc",",","scale",")",",",")",")","``","`","Later",",","when","we","want","to","infer","'typical","'","values","of","`","loc","`","and","`","scale","`","conditioned","on","some","given","`","data","`",",","we","will","often","write",":","``","`","python","def","target_log_prob_fn","(","loc",",","scale",")",":","return","jd.log_prob","(","loc=loc",",","scale=scale",",","obs=data",")","``","`","This","class","enables","one","to","write","instead",":","``","`","python","partial","=","tfde.JointDistributionPinned","(","jd",",","obs=data",")","target_log_prob_fn","=","partial.log_prob","``","`","Or",",","even","more","concisely","`","partial","=","jd.experimental_pin","(","obs=data",")","`",".","This","is","nice",",","but","it","was","n't","too","hard","to","write","out","the","`","target_log_prob_fn","`","function","explicitly",".","Now",",","let","'s","consider","that","for","many","inference","and","optimization","methods",",","we","may","want","to","use","a","smooth","change","of","variables","to","perform","inference","in","the","unconstrained","space","of","real","numbers",".","In","some","cases","this","transformation","can","be","parameter-dependent",".","For","example",",","if","we","want","to","unconstrain","the","support","of","`","tfp.distributions.Uniform","(","-3.",",","2",".",")","`","to","the","real","line",",","we","might","use","`","tfp.bijectors.Sigmoid","(","low=-3.",",","high=2.",")","`",".","In","support","of","such","use","cases",",","most","distributions","(","including","the","`","JointDistribution","*","`","classes",")","provide","a","`","experimental_default_event_space_bijector","(",")","`","method",".","When","these","transformations","may","be","dependent","on","ancestral","parts","of","a","joint","distribution",",","and","some","of","those","parameters","may","be","pinned",",","it","is","helpful","to","have","a","utility","class","to","bridge","the","gap","and","provide","the","multi-part","bijective","transform",".","This","is","the","``","raison","d'\u00eatre","''","of","this","class",".","The","model","below","is","somewhat","contrived",",","but","demonstrates","the","use-case",".","``","`","python","tfd","=","tfp.distributions","tfde","=","tfp.experimental.distributions","n","=","75","dim","=","3","joint","=","tfd.JointDistributionNamed","(","dict","(","upper","=","tfd.Uniform","(",".4",",","1.5",")",",","concentration","=","tfd.Gamma","(","1.",",",".5",")",",","corr","=","lambda","concentration",":","tfd.CholeskyLKJ","(","dim",",","concentration=concentration",")",",","stddev","=","lambda","upper",":","tfd.Sample","(","tfd.Uniform","(",".2",",","upper",")",",","dim",")",",","obs","=","lambda","corr",",","stddev",":","tfd.Sample","(","tfd.MultivariateNormalTriL","(","loc=tf.zeros","(","[","dim","]",")",",","scale_tril=corr","*","stddev","[","...",",","tf.newaxis","]",")",",","n",")",")",")","fixed_upper","=","1.3","data","=","joint.sample","(","upper=fixed_upper",")","[","'obs","'","]","pinned","=","tfde.JointDistributionPinned","(","joint",",","upper=fixed_upper",",","obs=data",")","bij","=","pinned.experimental_default_event_space_bijector","(",")","pulled_back_shape","=","bij.inverse_event_shape","(","pinned.event_shape",")","#","Fit","an","ensemble","using","SGD",".","batch","=","16","uniform_init","=","tf.nest.map_structure","(","lambda","s",":","tf.random.uniform","(","tf.concat","(","[","[","batch","]",",","s","]",",","axis=0",")",",","-2.",",","2",".",")",",","pulled_back_shape",")","vars","=","tf.nest.map_structure","(","tf.Variable",",","uniform_init",")","opt","=","tf_keras.optimizers.Adam","(",".01",")","@","tf.function","(","autograph=False",")","def","one_step","(",")",":","with","tf.GradientTape","(",")","as","tape",":","lp","=","pinned.log_prob","(","bij.forward","(","vars",")",")","gradients","=","tape.gradient","(","lp",",","vars",")","opt.apply_gradients","(","zip","(","gradients.values","(",")",",","vars.values","(",")",")",")","for","_","in","range","(","100",")",":","one_step","(",")","#","Alternatively",",","sample","using","MCMC","(","currently","aspirational",")",":","initial_state","=","bij.forward","(","uniform_init",")","kernel","=","tfp.mcmc.HamiltonianMonteCarlo","(","target_log_prob_fn=pinned.log_prob",",","step_size=.5",",","num_leapfrog_steps=4",")","#","*","*","This","line","is","currently","aspirational","*","*",",","to","demonstrate","the","use-case",".","kernel","=","tfp.mcmc.TransformedTransitionKernel","(","kernel",",","bij",")","tfp.mcmc.sample_chain","(","10",",","kernel=kernel",",","current_state=initial_state",")","``","`"]
docstr: Constructs a `JointDistributionPinned`.
  
  ### Examples:
  
  ```
  JointDistributionPinned(
     tfd.JointDistributionSequential([
         tfd.Gamma(1., 1.), lambda scale: tfd.Normal(0., scale, name='x')]),
     1.)  # Pins the `scale` part.
  
  JointDistributionPinned(
     tfd.JointDistributionSequential([
         tfd.Gamma(1., 1.), lambda scale: tfd.Normal(0., scale, name='x')]),
     [1.])  # Pins the `scale` part to scalar `1.`.
  
  JointDistributionPinned(
     tfd.JointDistributionSequential([
         tfd.Gamma(1., 1.), lambda scale: tfd.Normal(0., scale, name='x')]),
     None, 1.)  # Pins the `x` part.
  
  JointDistributionPinned(
     tfd.JointDistributionSequential([
         tfd.Gamma(1., 1.), lambda scale: tfd.Normal(0., scale, name='x')]),
     [None, 1.])  # Pins the `x` part.
  
  JointDistributionPinned(
     tfd.JointDistributionSequential([
         tfd.Gamma(1., 1.), lambda scale: tfd.Normal(0., scale, name='x')]),
     scale=1.)  # Pins the `scale` part.
  
  JointDistributionPinned(
     tfd.JointDistributionSequential([
         tfd.Gamma(1., 1.), lambda scale: tfd.Normal(0., scale, name='x')]),
     x=1.)  # Pins the `x` part.
  ```
  
  Args:
    distribution: A `tfp.distributions.JointDistribution`.
    *pins: A single object like the `value` argument that may be passed into
      `JointDistribution.sample` (some parts may be `None`), or a sequence of
      objects similar to such sequence as might be passed to
      `JointDistribution.log_prob`, but with the difference that some parts
      may be `None` (`log_prob` would require all parts be specified).
      More precisely, the user may pass (A) a single argument specifiying pins
      of one or more of the parts of the underlying `distribution` either by
      name (i.e. a `dict`, `namedtuple`) or by sequence ordering (`tuple`,
      `list`), or (B) a sequence of arguments which align with the model of
      the underlying distribution (which must be ordered). It is an error to
      use an unordered sequence of pins with an unordered model, e.g. a
      `tfp.distributions.JointDistributionNamed` constructed with a `dict`
      model (`collections.OrderedDict` is allowed).
    name: Python `str` name for this distribution. If `None`, defaults to
      'Pinned{distribution.name}'.
      Default value: `None`.
    **named_pins: Named elements to pin. The names given must align with the
      part names defined by `distribution._flat_resolve_names()`, i.e. either
      the explicitly named parts of `tfp.distributions.JointDistributionNamed`
      or the `name` parameters passed to distributions constructed by the
      model given to `JointDistribution*`.
docstr_tok: ["Constructs","a","`","JointDistributionPinned","`",".","#","#","#","Examples",":","``","`","JointDistributionPinned","(","tfd.JointDistributionSequential","(","[","tfd.Gamma","(","1.",",","1",".",")",",","lambda","scale",":","tfd.Normal","(","0.",",","scale",",","name=","'","x","'",")","]",")",",","1",".",")","#","Pins","the","`","scale","`","part",".","JointDistributionPinned","(","tfd.JointDistributionSequential","(","[","tfd.Gamma","(","1.",",","1",".",")",",","lambda","scale",":","tfd.Normal","(","0.",",","scale",",","name=","'","x","'",")","]",")",",","[","1",".","]",")","#","Pins","the","`","scale","`","part","to","scalar","`","1.","`",".","JointDistributionPinned","(","tfd.JointDistributionSequential","(","[","tfd.Gamma","(","1.",",","1",".",")",",","lambda","scale",":","tfd.Normal","(","0.",",","scale",",","name=","'","x","'",")","]",")",",","None",",","1",".",")","#","Pins","the","`","x","`","part",".","JointDistributionPinned","(","tfd.JointDistributionSequential","(","[","tfd.Gamma","(","1.",",","1",".",")",",","lambda","scale",":","tfd.Normal","(","0.",",","scale",",","name=","'","x","'",")","]",")",",","[","None",",","1",".","]",")","#","Pins","the","`","x","`","part",".","JointDistributionPinned","(","tfd.JointDistributionSequential","(","[","tfd.Gamma","(","1.",",","1",".",")",",","lambda","scale",":","tfd.Normal","(","0.",",","scale",",","name=","'","x","'",")","]",")",",","scale=1",".",")","#","Pins","the","`","scale","`","part",".","JointDistributionPinned","(","tfd.JointDistributionSequential","(","[","tfd.Gamma","(","1.",",","1",".",")",",","lambda","scale",":","tfd.Normal","(","0.",",","scale",",","name=","'","x","'",")","]",")",",","x=1",".",")","#","Pins","the","`","x","`","part",".","``","`","Args",":","distribution",":","A","`","tfp.distributions.JointDistribution","`",".","*","pins",":","A","single","object","like","the","`","value","`","argument","that","may","be","passed","into","`","JointDistribution.sample","`","(","some","parts","may","be","`","None","`",")",",","or","a","sequence","of","objects","similar","to","such","sequence","as","might","be","passed","to","`","JointDistribution.log_prob","`",",","but","with","the","difference","that","some","parts","may","be","`","None","`","(","`","log_prob","`","would","require","all","parts","be","specified",")",".","More","precisely",",","the","user","may","pass","(","A",")","a","single","argument","specifiying","pins","of","one","or","more","of","the","parts","of","the","underlying","`","distribution","`","either","by","name","(","i.e",".","a","`","dict","`",",","`","namedtuple","`",")","or","by","sequence","ordering","(","`","tuple","`",",","`","list","`",")",",","or","(","B",")","a","sequence","of","arguments","which","align","with","the","model","of","the","underlying","distribution","(","which","must","be","ordered",")",".","It","is","an","error","to","use","an","unordered","sequence","of","pins","with","an","unordered","model",",","e.g",".","a","`","tfp.distributions.JointDistributionNamed","`","constructed","with","a","`","dict","`","model","(","`","collections.OrderedDict","`","is","allowed",")",".","name",":","Python","`","str","`","name","for","this","distribution",".","If","`","None","`",",","defaults","to","'Pinned","{","distribution.name","}","'",".","Default","value",":","`","None","`",".","*","*","named_pins",":","Named","elements","to","pin",".","The","names","given","must","align","with","the","part","names","defined","by","`","distribution._flat_resolve_names","(",")","`",",","i.e",".","either","the","explicitly","named","parts","of","`","tfp.distributions.JointDistributionNamed","`","or","the","`","name","`","parameters","passed","to","distributions","constructed","by","the","model","given","to","`","JointDistribution","*","`","."]
code: def __init__(self, distribution, *pins, name=None, **named_pins):
      """Constructs a `JointDistributionPinned`.
  
      ### Examples:
  
      ```
      JointDistributionPinned(
         tfd.JointDistributionSequential([
             tfd.Gamma(1., 1.), lambda scale: tfd.Normal(0., scale, name='x')]),
         1.)  # Pins the `scale` part.
  
      JointDistributionPinned(
         tfd.JointDistributionSequential([
             tfd.Gamma(1., 1.), lambda scale: tfd.Normal(0., scale, name='x')]),
         [1.])  # Pins the `scale` part to scalar `1.`.
  
      JointDistributionPinned(
         tfd.JointDistributionSequential([
             tfd.Gamma(1., 1.), lambda scale: tfd.Normal(0., scale, name='x')]),
         None, 1.)  # Pins the `x` part.
  
      JointDistributionPinned(
         tfd.JointDistributionSequential([
             tfd.Gamma(1., 1.), lambda scale: tfd.Normal(0., scale, name='x')]),
         [None, 1.])  # Pins the `x` part.
  
      JointDistributionPinned(
         tfd.JointDistributionSequential([
             tfd.Gamma(1., 1.), lambda scale: tfd.Normal(0., scale, name='x')]),
         scale=1.)  # Pins the `scale` part.
  
      JointDistributionPinned(
         tfd.JointDistributionSequential([
             tfd.Gamma(1., 1.), lambda scale: tfd.Normal(0., scale, name='x')]),
         x=1.)  # Pins the `x` part.
      ```
  
      Args:
        distribution: A `tfp.distributions.JointDistribution`.
        *pins: A single object like the `value` argument that may be passed into
          `JointDistribution.sample` (some parts may be `None`), or a sequence of
          objects similar to such sequence as might be passed to
          `JointDistribution.log_prob`, but with the difference that some parts
          may be `None` (`log_prob` would require all parts be specified).
          More precisely, the user may pass (A) a single argument specifiying pins
          of one or more of the parts of the underlying `distribution` either by
          name (i.e. a `dict`, `namedtuple`) or by sequence ordering (`tuple`,
          `list`), or (B) a sequence of arguments which align with the model of
          the underlying distribution (which must be ordered). It is an error to
          use an unordered sequence of pins with an unordered model, e.g. a
          `tfp.distributions.JointDistributionNamed` constructed with a `dict`
          model (`collections.OrderedDict` is allowed).
        name: Python `str` name for this distribution. If `None`, defaults to
          'Pinned{distribution.name}'.
          Default value: `None`.
        **named_pins: Named elements to pin. The names given must align with the
          part names defined by `distribution._flat_resolve_names()`, i.e. either
          the explicitly named parts of `tfp.distributions.JointDistributionNamed`
          or the `name` parameters passed to distributions constructed by the
          model given to `JointDistribution*`.
      """
      if bool(pins) == bool(named_pins):
          raise ValueError(
              "Exactly one of *pins or **named_pins should be set."
          )
  
      if name is None:
          name = "Pinned{}".format(distribution.name)
  
      self._distribution = distribution
      self._name = name
      self._pins = _to_pins(distribution, *pins, **named_pins)
      self._use_vectorized_map = getattr(
          distribution, "use_vectorized_map", False
      )
      self.allow_nan_stats = False
  
imports: import collections
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import distribution
  from tensorflow_probability.python.distributions import joint_distribution
  from tensorflow_probability.python.internal import docstring_util
  from tensorflow_probability.python.internal import structural_tuple
  from tensorflow_probability.python.internal import tensor_util
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.internal import vectorization_util
  from tensorflow.python.util import nest
label_desc: Use image node_id 1 to create a new JointDistributionPinned object from inherited base classes: object with example: obj = JointDistributionPinned(distribution)