code: def greedy():
      """A greedy distribution."""
      warnings.warn(
          "Rlax greedy will be deprecated. "
          "Please use distrax.Greedy instead.",
          PendingDeprecationWarning,
          stacklevel=2,
      )
  
      def sample_fn(key: Array, preferences: Array):
          return distrax.Greedy(preferences).sample(seed=key)
  
      def probs_fn(preferences: Array):
          return distrax.Greedy(preferences).probs
  
      def log_prob_fn(sample: Array, preferences: Array):
          return distrax.Greedy(preferences).log_prob(sample)
  
      def entropy_fn(preferences: Array):
          return distrax.Greedy(preferences).entropy()
  
      return DiscreteDistribution(
          sample_fn, probs_fn, log_prob_fn, entropy_fn, None
      )
  
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 4 for calling a global function with example usage: greedy() and returns: DiscreteDistribution, distrax, distrax, distrax, distrax