class_docstr: Module that outputs a Gaussian Mixture Distribution.
class_docstr_tok: ["Module","that","outputs","a","Gaussian","Mixture","Distribution","."]
docstr: Run the networks through inputs.
  
  Args:
    inputs: hidden activations of the policy network body.
    low_noise_policy: whether to set vanishingly small scales for each
      component. If this flag is set to True, the policy is effectively run
      without Gaussian noise.
  
  Returns:
    Mixture Gaussian distribution.
docstr_tok: ["Run","the","networks","through","inputs",".","Args",":","inputs",":","hidden","activations","of","the","policy","network","body",".","low_noise_policy",":","whether","to","set","vanishingly","small","scales","for","each","component",".","If","this","flag","is","set","to","True",",","the","policy","is","effectively","run","without","Gaussian","noise",".","Returns",":","Mixture","Gaussian","distribution","."]
code: def __call__(
      self, inputs: jnp.ndarray, low_noise_policy: bool = False
  ) -> tfd.Distribution:
      """Run the networks through inputs.
  
      Args:
        inputs: hidden activations of the policy network body.
        low_noise_policy: whether to set vanishingly small scales for each
          component. If this flag is set to True, the policy is effectively run
          without Gaussian noise.
  
      Returns:
        Mixture Gaussian distribution.
      """
  
      # Define the weight initializer.
      w_init = hk.initializers.VarianceScaling(scale=1e-5)
  
      # Create a layer that outputs the unnormalized log-weights.
      if self._multivariate:
          logits_size = self._num_components
      else:
          logits_size = self._num_dimensions * self._num_components
      logit_layer = hk.Linear(logits_size, w_init=w_init)
  
      # Create two layers that outputs a location and a scale, respectively, for
      # each dimension and each component.
      loc_layer = hk.Linear(
          self._num_dimensions * self._num_components, w_init=w_init
      )
      scale_layer = hk.Linear(
          self._num_dimensions * self._num_components, w_init=w_init
      )
  
      # Compute logits, locs, and scales if necessary.
      logits = logit_layer(inputs)
      locs = loc_layer(inputs)
  
      # When a low_noise_policy is requested, set the scales to its minimum value.
      if low_noise_policy:
          scales = jnp.full(locs.shape, _MIN_SCALE)
      else:
          scales = scale_layer(inputs)
          scales = (
              self._scale_factor * jax.nn.softplus(scales) + _MIN_SCALE
          )
  
      if self._multivariate:
          components_class = tfd.MultivariateNormalDiag
          shape = [
              -1,
              self._num_components,
              self._num_dimensions,
          ]  # [B, C, D]
          # In this case, no need to reshape logits as they are in the correct shape
          # already, namely [batch_size, num_components].
      else:
          components_class = tfd.Normal
          shape = [
              -1,
              self._num_dimensions,
              self._num_components,
          ]  # [B, D, C]
          if self._append_singleton_event_dim:
              shape.insert(2, 1)  # [B, D, 1, C]
          logits = logits.reshape(shape)
  
      # Reshape the mixture's location and scale parameters appropriately.
      locs = locs.reshape(shape)
      scales = scales.reshape(shape)
  
      if self._multivariate:
          components_distribution = components_class(
              loc=locs, scale_diag=scales
          )
      else:
          components_distribution = components_class(
              loc=locs, scale=scales
          )
  
      # Transformed the component distributions in the mixture.
      if self._transformation_fn:
          components_distribution = self._transformation_fn(
              components_distribution
          )
  
      # Create the mixture distribution.
      distribution = tfd.MixtureSameFamily(
          mixture_distribution=tfd.Categorical(logits=logits),
          components_distribution=components_distribution,
      )
  
      if not self._multivariate:
          distribution = tfd.Independent(
              distribution,
              reinterpreted_batch_ndims=self._reinterpreted_batch_ndims,
          )
  
      return distribution
  
imports: from typing import Any, List, Optional, Sequence, Union, Callable
  import chex
  import haiku
  import jax
  import jax.numpy
  import numpy
  import tensorflow_probability
  import tensorflow_probability.substrates.jax
label_desc: Use image node_id 2 for calling the GaussianMixture obj's underlying member method code with example usage: obj.__call__(inputs, low_noise_policy) and returns: distribution