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",")","``","`"]
code: def batch_shape_tensor(self):
      batch_shape = self.distribution.batch_shape_tensor()
      if tf.nest.is_nested(batch_shape):
          return self._prune(batch_shape, retain="unpinned")
      return batch_shape
  
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 16 for calling the JointDistributionPinned obj's underlying member method code with example usage: obj.batch_shape_tensor() and returns: batch_shape, self