code: def _random_categorical_sample(
      num_samples: int,
      weights: type_alias.TensorLike,
      seed: Optional[type_alias.TensorLike] = None,
      stateless: bool = False,
      name: str = "random_categorical_sample",
      sample_dtype: tf.DType = tf.int32,
  ) -> type_alias.TensorLike:
      """Samples from a categorical distribution with arbitrary batch dimensions.
  
      Note:
        In the following, A1 to An are optional batch dimensions.
  
      Args:
        num_samples: An `int32` scalar denoting the number of samples to generate
          per mesh.
        weights: A `float` tensor of shape `[A1, ..., An, F]` where F is number of
          faces.
          All weights must be > 0.
        seed: Optional random seed, value depends on `stateless`.
        stateless: Optional flag to use stateless random sampler. If stateless=True,
          then `seed` must be provided as shape `[2]` int tensor. Stateless random
          sampling is useful for testing to generate the same reproducible sequence
          across calls. If stateless=False, then a stateful random number generator
          is used (default behavior).
        name: Name for op. Defaults to "random_categorical_sample".
        sample_dtype: Type of output samples.
  
      Returns:
        A `sample_dtype` tensor of shape `[A1, ..., An, num_samples]`.
      """
      with tf.name_scope(name):
          asserts.assert_all_above(weights, 0)
          logits = tf.math.log(weights)
          num_faces = tf.shape(input=logits)[-1]
          batch_shape = tf.shape(input=logits)[:-1]
          logits_2d = tf.reshape(logits, [-1, num_faces])
          if stateless:
              seed = tf.convert_to_tensor(value=seed)
              shape.check_static(
                  tensor=seed,
                  tensor_name="seed",
                  has_dim_equals=(-1, 2),
              )
              sample_fn = tf.random.stateless_categorical
          else:
              sample_fn = tf.random.categorical
          draws = sample_fn(
              logits=logits_2d,
              num_samples=num_samples,
              dtype=sample_dtype,
              seed=seed,
          )
          samples = tf.reshape(
              draws,
              shape=tf.concat((batch_shape, (num_samples,)), axis=0),
          )
          return samples
  
file_path: graphics/tensorflow_graphics/geometry/representation/mesh/sampler.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from typing import Optional, Tuple
  import tensorflow
  from tensorflow_graphics.geometry.representation import triangle
  from tensorflow_graphics.geometry.representation.mesh import normals
  from tensorflow_graphics.util import asserts
  from tensorflow_graphics.util import export_api
  from tensorflow_graphics.util import shape
  from tensorflow_graphics.util import type_alias
label_desc: Use image node_id 2 for calling a global function with example usage: _random_categorical_sample(num_samples, weights, seed, stateless, name, sample_dtype) and returns: samples