code: def gaussian_diagonal(sigma=None):
      """A gaussian distribution with diagonal covariance matrix."""
  
      warnings.warn(
          "Rlax gaussian_diagonal will be deprecated. "
          "Please use distrax MultivariateNormalDiag instead.",
          PendingDeprecationWarning,
          stacklevel=2,
      )
  
      def sample_fn(key: Array, mu: Array, sigma: Array = sigma):
          return distrax.MultivariateNormalDiag(
              mu, jnp.ones_like(mu) * sigma
          ).sample(seed=key)
  
      def prob_fn(sample: Array, mu: Array, sigma: Array = sigma):
          return distrax.MultivariateNormalDiag(
              mu, jnp.ones_like(mu) * sigma
          ).prob(sample)
  
      def logprob_fn(sample: Array, mu: Array, sigma: Array = sigma):
          return distrax.MultivariateNormalDiag(
              mu, jnp.ones_like(mu) * sigma
          ).log_prob(sample)
  
      def entropy_fn(mu: Array, sigma: Array = sigma):
          return distrax.MultivariateNormalDiag(
              mu, jnp.ones_like(mu) * sigma
          ).entropy()
  
      def kl_to_standard_normal_fn(mu: Array, sigma: Array = sigma):
          return distrax.MultivariateNormalDiag(
              mu, jnp.ones_like(mu) * sigma
          ).kl_divergence(
              distrax.MultivariateNormalDiag(
                  jnp.zeros_like(mu), jnp.ones_like(mu)
              )
          )
  
      def kl_fn(
          mu_0: Array, sigma_0: Numeric, mu_1: Array, sigma_1: Numeric
      ):
          return distrax.MultivariateNormalDiag(
              mu_0, jnp.ones_like(mu_0) * sigma_0
          ).kl_divergence(
              distrax.MultivariateNormalDiag(
                  mu_1, jnp.ones_like(mu_1) * sigma_1
              )
          )
  
      return ContinuousDistribution(
          sample_fn,
          prob_fn,
          logprob_fn,
          entropy_fn,
          kl_to_standard_normal_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 6 for calling a global function with example usage: gaussian_diagonal(sigma) and returns: ContinuousDistribution, distrax, distrax, distrax, distrax, distrax, distrax