class_docstr: Module that outputs a Gaussian Mixture Distribution.
class_docstr_tok: ["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.
    append_singleton_event_dim: (univariate only) Whether to add an extra
      singleton dimension to the event shape.
    reinterpreted_batch_ndims: (univariate only) Number of batch dimensions to
      reinterpret as event dimensions.
    transformation_fn: Distribution transform such as TanhTransformation
      applied to individual 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",".","append_singleton_event_dim",":","(","univariate","only",")","Whether","to","add","an","extra","singleton","dimension","to","the","event","shape",".","reinterpreted_batch_ndims",":","(","univariate","only",")","Number","of","batch","dimensions","to","reinterpret","as","event","dimensions",".","transformation_fn",":","Distribution","transform","such","as","TanhTransformation","applied","to","individual","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,
      append_singleton_event_dim: bool = False,
      reinterpreted_batch_ndims: Optional[int] = None,
      transformation_fn: Optional[
          Callable[[tfd.Distribution], tfd.Distribution]
      ] = 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.
        append_singleton_event_dim: (univariate only) Whether to add an extra
          singleton dimension to the event shape.
        reinterpreted_batch_ndims: (univariate only) Number of batch dimensions to
          reinterpret as event dimensions.
        transformation_fn: Distribution transform such as TanhTransformation
          applied to individual 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
      self._append_singleton_event_dim = append_singleton_event_dim
      self._reinterpreted_batch_ndims = reinterpreted_batch_ndims
  
      if init_scale is not None:
          self._scale_factor = init_scale / jax.nn.softplus(0.0)
      else:
          self._scale_factor = (
              1.0  # Corresponds to init_scale = softplus(0).
          )
  
      self._transformation_fn = transformation_fn
  
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 1 to create a new GaussianMixture object from inherited base classes: hk with example: obj = GaussianMixture(num_dimensions, num_components, multivariate, init_scale, append_singleton_event_dim, reinterpreted_batch_ndims, transformation_fn, name)