class_docstr: The von Mises-Fisher distribution over unit vectors on `S^{n-1}`.
  
  The von Mises-Fisher distribution is a directional distribution over vectors
  on the unit hypersphere `S^{n-1}` embedded in `n` dimensions (`R^n`).
  
  #### Mathematical details
  
  The probability density function (pdf) is,
  
  ```none
  pdf(x; mu, kappa) = C(kappa) exp(kappa * mu^T x)
  where,
  C(kappa) = (2 pi)^{-n/2} kappa^{n/2-1} / I_{n/2-1}(kappa),
  I_v(z) being the modified Bessel function of the first kind of order v
  ```
  
  where:
  * `mean_direction = mu`; a unit vector in `R^k`,
  * `concentration = kappa`; scalar real >= 0, concentration of samples around
    `mean_direction`, where 0 pertains to the uniform distribution on the
    hypersphere, and \inf indicates a delta function at `mean_direction`.
  
  
  NOTE: Currently only n in {2, 3, 4, 5} are supported. For n=5 some numerical
  instability can occur for low concentrations (<.01).
  
  NOTE: `mean_direction` is not in general the mean of the distribution. For
  spherical distributions, the mean is generally not in the support of the
  distribution.
  
  #### Examples
  
  A single instance of a vMF distribution is defined by a mean direction (or
  mode) unit vector and a scalar concentration parameter.
  
  Extra leading dimensions, if provided, allow for batches.
  
  ```python
  tfd = tfp.distributions
  
  # Initialize a single 3-dimension vMF distribution.
  mu = [0., 1, 0]
  conc = 1.
  vmf = tfd.VonMisesFisher(mean_direction=mu, concentration=conc)
  
  # Evaluate this on an observation in S^2 (in R^3), returning a scalar.
  vmf.prob([1., 0, 0])
  
  # Initialize a batch of two 3-variate vMF distributions.
  mu = [[0., 1, 0],
        [1., 0, 0]]
  conc = [1., 2]
  vmf = tfd.VonMisesFisher(mean_direction=mu, concentration=conc)
  
  # Evaluate this on two observations, each in S^2, returning a length two
  # tensor.
  x = [[0., 0, 1],
       [0., 1, 0]]
  vmf.prob(x)
  ```
class_docstr_tok: ["The","von","Mises-Fisher","distribution","over","unit","vectors","on","`","S^","{","n-1","}","`",".","The","von","Mises-Fisher","distribution","is","a","directional","distribution","over","vectors","on","the","unit","hypersphere","`","S^","{","n-1","}","`","embedded","in","`","n","`","dimensions","(","`","R^n","`",")",".","#","#","#","#","Mathematical","details","The","probability","density","function","(","pdf",")","is",",","``","`","none","pdf","(","x",";","mu",",","kappa",")","=","C","(","kappa",")","exp","(","kappa","*","mu^T","x",")","where",",","C","(","kappa",")","=","(","2","pi",")","^","{","-n\/2","}","kappa^","{","n\/2-1","}","\/","I_","{","n\/2-1","}","(","kappa",")",",","I_v","(","z",")","being","the","modified","Bessel","function","of","the","first","kind","of","order","v","``","`","where",":","*","`","mean_direction","=","mu","`",";","a","unit","vector","in","`","R^k","`",",","*","`","concentration","=","kappa","`",";","scalar","real",">","=","0",",","concentration","of","samples","around","`","mean_direction","`",",","where","0","pertains","to","the","uniform","distribution","on","the","hypersphere",",","and","\\inf","indicates","a","delta","function","at","`","mean_direction","`",".","NOTE",":","Currently","only","n","in","{","2",",","3",",","4",",","5","}","are","supported",".","For","n=5","some","numerical","instability","can","occur","for","low","concentrations","(","<",".01",")",".","NOTE",":","`","mean_direction","`","is","not","in","general","the","mean","of","the","distribution",".","For","spherical","distributions",",","the","mean","is","generally","not","in","the","support","of","the","distribution",".","#","#","#","#","Examples","A","single","instance","of","a","vMF","distribution","is","defined","by","a","mean","direction","(","or","mode",")","unit","vector","and","a","scalar","concentration","parameter",".","Extra","leading","dimensions",",","if","provided",",","allow","for","batches",".","``","`","python","tfd","=","tfp.distributions","#","Initialize","a","single","3-dimension","vMF","distribution",".","mu","=","[","0.",",","1",",","0","]","conc","=","1.","vmf","=","tfd.VonMisesFisher","(","mean_direction=mu",",","concentration=conc",")","#","Evaluate","this","on","an","observation","in","S^2","(","in","R^3",")",",","returning","a","scalar",".","vmf.prob","(","[","1.",",","0",",","0","]",")","#","Initialize","a","batch","of","two","3-variate","vMF","distributions",".","mu","=","[","[","0.",",","1",",","0","]",",","[","1.",",","0",",","0","]","]","conc","=","[","1.",",","2","]","vmf","=","tfd.VonMisesFisher","(","mean_direction=mu",",","concentration=conc",")","#","Evaluate","this","on","two","observations",",","each","in","S^2",",","returning","a","length","two","#","tensor",".","x","=","[","[","0.",",","0",",","1","]",",","[","0.",",","1",",","0","]","]","vmf.prob","(","x",")","``","`"]
code: def _sample_n(self, n, seed=None):
      dim0_seed, otherdims_seed = samplers.split_seed(
          seed, salt="von_mises_fisher"
      )
      # The sampling strategy relies on the fact that vMF variates are symmetric
      # about the mean direction. Accordingly, if we have a sampling strategy for
      # the away-from-mean angle, then we can uniformly sample the remaining
      # dimensions on the S^{dim-2} sphere for , and rotate these samples from a
      # (1, 0, 0, ..., 0)-mode distribution into the target orientation.
      #
      # This is easy to imagine on the 1-sphere (S^1; in 2-D space): sample a
      # von-Mises distributed `x` value in [-1, 1], then uniformly select what
      # amounts to a "up" or "down" additional degree of freedom after unit
      # normalizing, followed by a final rotation to the desired mean direction
      # from a basis of (1, 0).
      #
      # On S^2 (in 3-D), selecting a vMF `x` identifies a circle in `yz` on the
      # unit sphere over which the distribution is uniform, in particular the
      # circle where x = \hat{x} intersects the unit sphere. We pick a point on
      # that circle, then rotate to the desired mean direction from a basis of
      # (1, 0, 0).
      mean_direction = tf.convert_to_tensor(self.mean_direction)
      concentration = tf.convert_to_tensor(self.concentration)
      event_dim = (
          tf.compat.dimension_value(self.event_shape[0])
          or self._event_shape_tensor(mean_direction=mean_direction)[0]
      )
  
      sample_batch_shape = ps.concat(
          [
              [n],
              self._batch_shape_tensor(
                  mean_direction=mean_direction,
                  concentration=concentration,
              ),
          ],
          axis=0,
      )
      dim = tf.cast(event_dim - 1, self.dtype)
      if event_dim == 3:
          samples_dim0 = self._sample_3d(
              n,
              mean_direction=mean_direction,
              concentration=concentration,
              seed=dim0_seed,
          )
      else:
          # Wood'94 provides a rejection algorithm to sample the x coordinate.
          # Wood'94 definition of b:
          # b = (-2 * kappa + tf.sqrt(4 * kappa**2 + dim**2)) / dim
          # https://stats.stackexchange.com/questions/156729 suggests:
          b = dim / (
              2 * concentration
              + tf.sqrt(4 * concentration**2 + dim**2)
          )
          # TODO(bjp): Integrate any useful numerical tricks from hyperspherical VAE
          #     https://github.com/nicola-decao/s-vae-tf/
          x = (1 - b) / (1 + b)
          c = concentration * x + dim * tf.math.log1p(-(x**2))
          beta = beta_lib.Beta(dim / 2, dim / 2)
  
          def cond_fn(w, should_continue, seed):
              del w, seed
              return tf.reduce_any(should_continue)
  
          def body_fn(w, should_continue, seed):
              """While loop body for sampling the angle `w`."""
              beta_seed, unif_seed, next_seed = samplers.split_seed(
                  seed, n=3
              )
              z = beta.sample(
                  sample_shape=sample_batch_shape, seed=beta_seed
              )
              # set_shape needed here because of b/139013403
              tensorshape_util.set_shape(z, w.shape)
              w = tf.where(
                  should_continue,
                  (1.0 - (1.0 + b) * z) / (1.0 - (1.0 - b) * z),
                  w,
              )
              if not self.allow_nan_stats:
                  w = tf.debugging.check_numerics(w, "w")
              unif = samplers.uniform(
                  sample_batch_shape, seed=unif_seed, dtype=self.dtype
              )
              # set_shape needed here because of b/139013403
              tensorshape_util.set_shape(unif, w.shape)
              should_continue = should_continue & (
                  concentration * w + dim * tf.math.log1p(-x * w) - c
                  <
                  # Use log1p(-unif) to prevent log(0) and ensure that log(1) is
                  # possible.
                  tf.math.log1p(-unif)
              )
              return w, should_continue, next_seed
  
          w = tf.zeros(sample_batch_shape, dtype=self.dtype)
          should_continue = tf.ones(sample_batch_shape, dtype=tf.bool)
          samples_dim0, _, _ = tf.while_loop(
              cond=cond_fn,
              body=body_fn,
              loop_vars=(w, should_continue, dim0_seed),
          )
          samples_dim0 = samples_dim0[..., tf.newaxis]
      if not self._allow_nan_stats:
          # Verify samples are w/in -1, 1, with useful error output tensors (top
          # value rather than all values).
          with tf.control_dependencies(
              [
                  assert_util.assert_less_equal(
                      samples_dim0,
                      dtype_util.as_numpy_dtype(self.dtype)(1.01),
                  ),
                  assert_util.assert_greater_equal(
                      samples_dim0,
                      dtype_util.as_numpy_dtype(self.dtype)(-1.01),
                  ),
              ]
          ):
              samples_dim0 = tf.identity(samples_dim0)
      samples_otherdims_shape = ps.concat(
          [sample_batch_shape, [event_dim - 1]], axis=0
      )
      unit_otherdims = tf.math.l2_normalize(
          samplers.normal(
              samples_otherdims_shape,
              seed=otherdims_seed,
              dtype=self.dtype,
          ),
          axis=-1,
      )
      samples = tf.concat(
          [
              samples_dim0,  # we must avoid sqrt(1 - (>1)**2)
              tf.sqrt(tf.maximum(1 - samples_dim0**2, 0.0))
              * unit_otherdims,
          ],
          axis=-1,
      )
      samples = tf.math.l2_normalize(samples, axis=-1)
      if not self.allow_nan_stats:
          samples = tf.debugging.check_numerics(samples, "samples")
  
      # Runtime assert that samples are unit length.
      if not self.allow_nan_stats:
          worst, _ = tf.math.top_k(
              tf.reshape(
                  tf.abs(1 - tf.linalg.norm(samples, axis=-1)), [-1]
              )
          )
          with tf.control_dependencies(
              [
                  assert_util.assert_near(
                      dtype_util.as_numpy_dtype(self.dtype)(0),
                      worst,
                      atol=1e-4,
                      summarize=100,
                  )
              ]
          ):
              samples = tf.identity(samples)
      # The samples generated are symmetric around a mode at (1, 0, 0, ...., 0).
      # Now, we move the mode to `self.mean_direction` using a rotation matrix.
      if not self.allow_nan_stats:
          # Assert that the basis vector rotates to the mean direction, as expected.
          basis = tf.cast(
              tf.concat([[1.0], tf.zeros([event_dim - 1])], axis=0),
              self.dtype,
          )
          with tf.control_dependencies(
              [
                  assert_util.assert_less(
                      tf.linalg.norm(
                          self._rotate(
                              basis, mean_direction=mean_direction
                          )
                          - mean_direction,
                          axis=-1,
                      ),
                      dtype_util.as_numpy_dtype(self.dtype)(1e-5),
                  )
              ]
          ):
              return self._rotate(
                  samples, mean_direction=mean_direction
              )
      return self._rotate(samples, mean_direction=mean_direction)
  
imports: import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import chain
  from tensorflow_probability.python.bijectors import invert
  from tensorflow_probability.python.bijectors import softmax_centered
  from tensorflow_probability.python.bijectors import softplus
  from tensorflow_probability.python.bijectors import square
  from tensorflow_probability.python.distributions import beta
  from tensorflow_probability.python.distributions import distribution
  from tensorflow_probability.python.distributions import kullback_leibler
  from tensorflow_probability.python.distributions import spherical_uniform
  from tensorflow_probability.python.internal import assert_util
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import parameter_properties
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import reparameterization
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import tensor_util
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.math import bessel
label_desc: Use image node_id 17 for calling the VonMisesFisher obj's underlying member method code with example usage: obj._sample_n(n, seed) and returns: self, tf, w,  should_continue,  next_seed, self