class_docstr: Tests for token selector module.
class_docstr_tok: ["Tests","for","token","selector","module","."]
docstr: Test categorial sampler.
docstr_tok: ["Test","categorial","sampler","."]
code: def test_sample_categorial(
      self,
      logit_shape,
      num_samples,
      replacement,
      expected_shape,
      expected_raise,
      expected_unique,
      expeted_nonunique,
  ):
      """Test categorial sampler."""
      rng, sample_rng = random.split(random.PRNGKey(0))
      logits = random.normal(rng, logit_shape)
  
      kwargs = {}
      if replacement is not None:
          kwargs["replacement"] = replacement
  
      if expected_raise:
          with self.assertRaises(ValueError):
              samples = model_utils.sample_categorical(
                  sample_rng, logits, num_samples, **kwargs
              )
      else:
          samples = model_utils.sample_categorical(
              sample_rng, logits, num_samples, **kwargs
          )
          self.assertEqual(samples.shape, expected_shape)
  
          if expected_unique or expeted_nonunique:
              samples = jnp.reshape(samples, (-1, expected_shape[-1]))
              samples = jax.device_get(samples)
              for sample in samples:
                  if expected_unique:
                      self.assertEqual(
                          len(set(sample.tolist())), len(sample)
                      )
                  if expeted_nonunique:
                      self.assertLess(
                          len(set(sample.tolist())), len(sample)
                      )
  
file_path: scenic/scenic/projects/fast_vit/tests/test_model_utils.py
imports: import functools
  from absl.testing import absltest
  from absl.testing import parameterized
  import jax
  from jax import random
  import jax.numpy
  import numpy
  from scenic.projects.fast_vit import model_utils
label_desc: Use image node_id 2 for calling the TopKTokenSelectorTest obj's underlying member method code with example usage: obj.test_sample_categorial(logit_shape, num_samples, replacement, expected_shape, expected_raise, expected_unique, expeted_nonunique) without return types