code: def squashed_gaussian(sigma_min=-4, sigma_max=0.0):
      """A squashed gaussian distribution with diagonal covariance matrix."""
  
      warnings.warn(
          "Rlax squashed_gaussian will be deprecated. "
          "Please use distrax Transformed MultivariateNormalDiag distribution "
          "with chained Tanh/ScalarAffine bijector instead.",
          PendingDeprecationWarning,
          stacklevel=2,
      )
  
      def sigma_activation(
          sigma, sigma_min=sigma_min, sigma_max=sigma_max
      ):
          return jnp.exp(
              sigma_min
              + 0.5 * (sigma_max - sigma_min) * (jnp.tanh(sigma) + 1.0)
          )
  
      def mu_activation(mu):
          return jnp.tanh(mu)
  
      def get_squashed_gaussian_dist(mu, sigma, action_spec=None):
          if action_spec is not None:
              scale = 0.5 * (action_spec.maximum - action_spec.minimum)
              shift = action_spec.minimum
              bijector = distrax.Chain(
                  [
                      distrax.ScalarAffine(shift=shift, scale=scale),
                      distrax.ScalarAffine(shift=1.0),
                      distrax.Tanh(),
                  ]
              )
          else:
              bijector = distrax.Tanh()
          return distrax.Transformed(
              distribution=distrax.MultivariateNormalDiag(
                  loc=mu_activation(mu),
                  scale_diag=sigma_activation(sigma),
              ),
              bijector=distrax.Block(bijector, ndims=1),
          )
  
      def sample_fn(key: Array, mu: Array, sigma: Array, action_spec):
          return get_squashed_gaussian_dist(
              mu, sigma, action_spec
          ).sample(seed=key)
  
      def prob_fn(sample: Array, mu: Array, sigma: Array, action_spec):
          return get_squashed_gaussian_dist(
              mu, sigma, action_spec
          ).prob(sample)
  
      def logprob_fn(
          sample: Array, mu: Array, sigma: Array, action_spec
      ):
          return get_squashed_gaussian_dist(
              mu, sigma, action_spec
          ).log_prob(sample)
  
      def entropy_fn(mu: Array, sigma: Array):
          return get_squashed_gaussian_dist(
              mu, sigma
          ).distribution.entropy()
  
      def kl_to_standard_normal_fn(mu: Array, sigma: Array):
          return get_squashed_gaussian_dist(
              mu, sigma
          ).distribution.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 get_squashed_gaussian_dist(
              mu_0, sigma_0
          ).distribution.kl_divergence(
              get_squashed_gaussian_dist(mu_1, sigma_1).distribution
          )
  
      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 7 for calling a global function with example usage: squashed_gaussian(sigma_min, sigma_max) and returns: ContinuousDistribution, jnp, jnp, distrax, get_squashed_gaussian_dist, get_squashed_gaussian_dist, get_squashed_gaussian_dist, get_squashed_gaussian_dist, get_squashed_gaussian_dist, get_squashed_gaussian_dist