class_docstr: Module that outputs a Gaussian Mixture Distribution.
docstr: Initialization.
  
  Args:
    num_dimensions: dimensionality of the output distribution
    num_components: number of mixture components.
    multivariate: whether the resulting distribution is multivariate or not.
    init_scale: the initial scale for the Gaussian mixture components.
    name: name of the module passed to snt.Module parent class.
docstr_tok: ["Initialization",".","Args",":","num_dimensions",":","dimensionality","of","the","output","distribution","num_components",":","number","of","mixture","components",".","multivariate",":","whether","the","resulting","distribution","is","multivariate","or","not",".","init_scale",":","the","initial","scale","for","the","Gaussian","mixture","components",".","name",":","name","of","the","module","passed","to","snt.Module","parent","class","."]
code: def __init__(
      self,
      num_dimensions: int,
      num_components: int,
      multivariate: bool,
      init_scale: Optional[float] = None,
      name: str = "GaussianMixture",
  ):
      """Initialization.
  
      Args:
        num_dimensions: dimensionality of the output distribution
        num_components: number of mixture components.
        multivariate: whether the resulting distribution is multivariate or not.
        init_scale: the initial scale for the Gaussian mixture components.
        name: name of the module passed to snt.Module parent class.
      """
      super().__init__(name=name)
  
      self._num_dimensions = num_dimensions
      self._num_components = num_components
      self._multivariate = multivariate
  
      if init_scale is not None:
          self._scale_factor = init_scale / tf.nn.softplus(0.0)
      else:
          self._scale_factor = (
              1.0  # Corresponds to init_scale = softplus(0).
          )
  
      # Define the weight initializer.
      w_init = tf.initializers.VarianceScaling(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
      self._logit_layer = snt.Linear(logits_size, w_init=w_init)
  
      # Create two layers that outputs a location and a scale, respectively, for
      # each dimension and each component.
      self._loc_layer = snt.Linear(
          self._num_dimensions * self._num_components, w_init=w_init
      )
      self._scale_layer = snt.Linear(
          self._num_dimensions * self._num_components, w_init=w_init
      )
  
imports: import types
  from typing import Optional, Union
  from absl import logging
  from acme.tf.networks import distributions
  import numpy
  import sonnet
  import tensorflow
  import tensorflow_probability
label_desc: Use image node_id 1 to create a new GaussianMixture object from inherited base classes: snt with example: obj = GaussianMixture(num_dimensions, num_components, multivariate, init_scale, name)