code: def clipped_entropy_softmax(temperature=1.0, entropy_clip=1.0):
      """A softmax distribution with clipped entropy (1 is eq to not clipping)."""
  
      warnings.warn(
          "Rlax clipped_entropy_softmax will be deprecated. "
          "Please use distrax.Softmax instead.",
          PendingDeprecationWarning,
          stacklevel=2,
      )
  
      def sample_fn(key: Array, logits: Array, action_spec=None):
          del action_spec
          return distrax.Softmax(logits, temperature).sample(seed=key)
  
      def probs_fn(logits: Array, action_spec=None):
          del action_spec
          return distrax.Softmax(logits, temperature).probs
  
      def logprob_fn(sample: Array, logits: Array, action_spec=None):
          del action_spec
          return distrax.Softmax(logits, temperature).log_prob(sample)
  
      def entropy_fn(logits: Array):
          return jnp.minimum(
              distrax.Softmax(logits, temperature).entropy(),
              entropy_clip * jnp.log(logits.shape[-1]),
          )
  
      def kl_fn(p_logits: Array, q_logits: Array):
          return categorical_kl_divergence(
              p_logits, q_logits, temperature
          )
  
      return DiscreteDistribution(
          sample_fn, probs_fn, logprob_fn, entropy_fn, kl_fn
      )
  
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 3 for calling a global function with example usage: clipped_entropy_softmax(temperature, entropy_clip) and returns: DiscreteDistribution, distrax, distrax, distrax, jnp, categorical_kl_divergence