class_docstr: Creates a funnel-shaped distribution.
  
  This distribution was first described in [1]. The distribution is constructed
  by transforming a N-D gaussian with scale [3, 1, ...] by scaling all but the
  first dimensions by `exp(x0 / 2)` where  `x0` is the value of the first
  dimension.
  
  ```none
  x[0] ~ Normal(loc=0, scale=3)
  for i in range(1, ndims):
    x[i] ~ Normal(loc=0, scale=exp(x[0] / 2))
  ```
  
  This distribution is notable for having a relatively narrow "neck" region
  which is challenging for HMC to explore. This distribution resembles the
  posteriors of centrally parameterized hierarchical models.
  
  #### References
  
  1. Neal, R. M. (2003). Slice sampling. Annals of Statistics, 31(3), 705-767.
docstr: Construct the Neal's funnel model.
  
  Args:
    ndims: Python integer. Dimensionality of the distribution. Must be at
      least 2.
    name: Python `str` name prefixed to Ops created by this class.
    pretty_name: A Python `str`. The pretty name of this model.
  
  Raises:
    ValueError: If ndims < 2.
docstr_tok: ["Construct","the","Neal","'s","funnel","model",".","Args",":","ndims",":","Python","integer",".","Dimensionality","of","the","distribution",".","Must","be","at","least","2.","name",":","Python","`","str","`","name","prefixed","to","Ops","created","by","this","class",".","pretty_name",":","A","Python","`","str","`",".","The","pretty","name","of","this","model",".","Raises",":","ValueError",":","If","ndims","<","2","."]
code: def __init__(
      self,
      ndims=10,
      name="neals_funnel",
      pretty_name="Neal's Funnel",
  ):
      """Construct the Neal's funnel model.
  
      Args:
        ndims: Python integer. Dimensionality of the distribution. Must be at
          least 2.
        name: Python `str` name prefixed to Ops created by this class.
        pretty_name: A Python `str`. The pretty name of this model.
  
      Raises:
        ValueError: If ndims < 2.
      """
      if ndims < 2:
          raise ValueError(
              "ndims must be at least 2, saw: {}".format(ndims)
          )
  
      with tf.name_scope(name):
  
          def bijector_fn(x):
              """Funnel transform."""
              batch_shape = ps.shape(x)[:-1]
              scale = tf.concat(
                  [
                      tf.ones(ps.concat([batch_shape, [1]], axis=0)),
                      tf.exp(x[..., :1] / 2)
                      * tf.ones(
                          ps.concat([batch_shape, [ndims - 1]], axis=0)
                      ),
                  ],
                  axis=-1,
              )
              return tfb.Scale(scale)
  
          mg = tfd.MultivariateNormalDiag(
              loc=tf.zeros(ndims),
              scale_diag=[3.0] + [1.0] * (ndims - 1),
          )
          funnel = tfd.TransformedDistribution(
              mg,
              bijector=tfb.MaskedAutoregressiveFlow(
                  bijector_fn=bijector_fn
              ),
          )
  
          sample_transformations = {
              "identity": model.Model.SampleTransformation(
                  fn=lambda params: params,
                  pretty_name="Identity",
                  # The trailing dimensions come from a product distribution of
                  # independent standard normal and a log-normal with a scale of
                  # 3 / 2. See
                  # https://en.wikipedia.org/wiki/Product_distribution for the
                  # formulas. For the mean, the formulas yield zero.
                  ground_truth_mean=np.zeros(ndims),
                  # For the standard deviation, all means are zero and standard
                  # deivations of the normals are 1, so the formula reduces to
                  # `sqrt((sigma_log_normal + mean_log_normal**2))` which
                  # reduces to `exp((sigma_log_normal)**2)`.
                  ground_truth_standard_deviation=np.array(
                      [3.0] + [np.exp((3.0 / 2) ** 2)] * (ndims - 1)
                  ),
              )
          }
  
      self._funnel = funnel
  
      super(NealsFunnel, self).__init__(
          default_event_space_bijector=tfb.Identity(),
          event_shape=funnel.event_shape,
          dtype=funnel.dtype,
          name=name,
          pretty_name=pretty_name,
          sample_transformations=sample_transformations,
      )
  
file_path: probability/spinoffs/inference_gym/inference_gym/targets/neals_funnel.py
imports: import numpy
  import tensorflow.compat.v2
  import tensorflow_probability
  from tensorflow_probability.python.internal import prefer_static
  from inference_gym.targets import model
label_desc: Use image node_id 1 to create a new NealsFunnel object from inherited base classes: model with example: obj = NealsFunnel(ndims, name, pretty_name)