tensorflow_probability/python/distributions/von_mises_fisher.py
def _sample_3d(self, n, seed=None):
    """Specialized inversion sampler for 3D."""
    seed = seed_stream.SeedStream(seed, salt='von_mises_fisher_3d')
    u_shape = tf.concat([[n], self._batch_shape_tensor()], axis=0)
    z = tf.random.uniform(u_shape, seed=seed(), dtype=self.dtype)
    # TODO(bjp): Higher-order odd dim analytic CDFs are available in [1], could
    # be bisected for bounded sampling runtime (i.e. not rejection sampling).
    # [1]: Inversion sampler via: https://ieeexplore.ieee.org/document/7347705/
    # The inversion is: u = 1 + log(z + (1-z)*exp(-2*kappa)) / kappa
    # We must protect against both kappa and z being zero.
    safe_conc = tf.where(self.concentration > 0,
                         self.concentration,
                         tf.ones_like(self.concentration))
    safe_z = tf.where(z > 0, z, tf.ones_like(z))
    safe_u = 1 + tf.reduce_logsumexp(
        input_tensor=[
            tf.math.log(safe_z),
            tf.math.log1p(-safe_z) - 2 * safe_conc
        ],
        axis=0) / safe_conc
    # Limit of the above expression as kappa->0 is 2*z-1
    u = tf.where(self.concentration > tf.zeros_like(safe_u), safe_u,
                 2 * z - 1)
    # Limit of the expression as z->0 is -1.
    u = tf.where(tf.equal(z, 0), -tf.ones_like(u), u)
    if not self._allow_nan_stats:
      u = tf.debugging.check_numerics(u, 'u in _sample_3d')
    return u[..., tf.newaxis]
def _sample_3d(self, n, seed=None):
    """Specialized inversion sampler for 3D."""
    seed = seed_stream.SeedStream(seed, salt='von_mises_fisher_3d')
    u_shape = tf.concat([[n], self._batch_shape_tensor()], axis=0)
    z = tf.random.uniform(u_shape, seed=seed(), dtype=self.dtype)
    # TODO(bjp): Higher-order odd dim analytic CDFs are available in [1], could
    # be bisected for bounded sampling runtime (i.e. not rejection sampling).
    # [1]: Inversion sampler via: https://ieeexplore.ieee.org/document/7347705/
    # The inversion is: u = 1 + log(z + (1-z)*exp(-2*kappa)) / kappa
    # We must protect against both kappa and z being zero.
    safe_conc = tf.where(self.concentration > 0,
                         self.concentration,
                         tf.ones_like(self.concentration))
    safe_z = tf.where(z > 0, z, tf.ones_like(z))
    safe_u = 1 + tf.reduce_logsumexp(
        input_tensor=[
            tf.math.log(safe_z),
            tf.math.log1p(-safe_z) - 2 * safe_conc
        ],
        axis=0) / safe_conc
    # Limit of the above expression as kappa->0 is 2*z-1
    u = tf.where(self.concentration > tf.zeros_like(safe_u), safe_u,
                 2 * z - 1)
    # Limit of the expression as z->0 is -1.
    u = tf.where(tf.equal(z, 0), -tf.ones_like(u), u)
    if not self._allow_nan_stats:
      u = tf.debugging.check_numerics(u, 'u in _sample_3d')
    return u[..., tf.newaxis]
Specialized inversion sampler for 3D.