code: def testSampleMeanDir2d(self):
      mean_dirs = tf.math.l2_normalize(
          [[1.0, 1], [-2, 1], [0, -1]], axis=-1
      )
      concentration = [[0], [0.1], [2], [40], [1000]]
      vmf = von_mises_fisher.VonMisesFisher(
          mean_direction=mean_dirs,
          concentration=concentration,
          validate_args=True,
          allow_nan_stats=False,
      )
      self.assertEqual(
          [5, 3], tensorshape_util.as_list(vmf.batch_shape)
      )
      self.assertEqual([2], tensorshape_util.as_list(vmf.event_shape))
      nsamples = int(2e4)
      samples = vmf.sample(nsamples, seed=test_util.test_seed())
      self.assertEqual(
          [nsamples, 5, 3, 2], tensorshape_util.as_list(samples.shape)
      )
      sample_mean = self.evaluate(samples).mean(axis=0)
      # Assert that positive-concentration distributions have samples with
      # the expected mean direction.
      sample_dir = sample_mean / np.linalg.norm(
          sample_mean, axis=-1, keepdims=True
      )
      inner_product = self.evaluate(
          tf.reduce_sum(sample_dir * vmf.mean_direction, axis=-1)
      )
      # All except the 0-concentration distribution should have >0 inner product
      # with the mean direction of the distribution.
      self.assertAllGreater(inner_product[1:], 0.1)
      # Pick out >1 concentration distributions to assert ~1 inner product with
      # mean direction.
      self.assertAllClose(
          np.ones_like(inner_product)[2:], inner_product[2:], atol=1e-3
      )
      # Inner products should be roughly ascending by concentration.
      self.assertAllClose(
          np.round(np.sort(inner_product, axis=0), decimals=3),
          np.round(inner_product, decimals=3),
          atol=0.005,
      )
      means = self.evaluate(vmf.mean())
      # Mean vector for 0-concentration is precisely (0, 0).
      self.assertAllEqual(np.zeros_like(means[0]), means[0])
      mean_lengths = np.linalg.norm(means, axis=-1)
      # Length of the mean vector is strictly ascending with concentration.
      self.assertAllEqual(mean_lengths, np.sort(mean_lengths, axis=0))
      self.assertAllClose(
          np.linalg.norm(sample_mean, axis=-1), mean_lengths, atol=0.03
      )
  
imports: from absl.testing import parameterized
  import numpy
  from scipy import special
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import kullback_leibler
  from tensorflow_probability.python.distributions import spherical_uniform
  from tensorflow_probability.python.distributions import von_mises_fisher
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.internal import test_util
label_desc: Use image node_id 2 for calling the VonMisesFisherTest obj's underlying member method code with example usage: obj.testSampleMeanDir2d() without return types