code: def categorical_sample(key, probs):
      """Sample from a set of discrete probabilities."""
      warnings.warn(
          "Rlax categorical_sample will be deprecated. "
          "Please use distrax.Categorical.sample instead.",
          PendingDeprecationWarning,
          stacklevel=2,
      )
      return distrax.Categorical(probs=probs).sample(seed=key)
  
file_path: rlax/rlax/_src/distributions.py
imports: import collections
  import warnings
  import chex
  import distrax
  import jax.numpy
label_desc: Use image node_id 1 for calling a global function with example usage: categorical_sample(key, probs) and returns: distrax