code: def categorical_kl_divergence(
      p_logits: Array, q_logits: Array, temperature: float = 1.0
  ) -> Array:
      """Compute the KL between two categorical distributions from their logits.
  
      Args:
        p_logits: unnormalized logits for the first distribution.
        q_logits: unnormalized logits for the second distribution.
        temperature: the temperature for the softmax distribution, defaults at 1.
  
      Returns:
        the kl divergence between the distributions.
      """
      warnings.warn(
          "Rlax categorical_kl_divergence will be deprecated. "
          "Please use distrax.Softmax.kl_divergence instead.",
          PendingDeprecationWarning,
          stacklevel=2,
      )
      return distrax.Softmax(p_logits, temperature).kl_divergence(
          distrax.Softmax(q_logits, temperature)
      )
  
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 10 for calling a global function with example usage: categorical_kl_divergence(p_logits, q_logits, temperature) and returns: distrax