class_docstr: Module that produces a multivariate normal distribution using tfd.Independent or tfd.MultivariateNormalDiag.
docstr: Initialization.
  
  Args:
    num_dimensions: Number of dimensions of MVN distribution.
    init_scale: Initial standard deviation.
    min_scale: Minimum standard deviation.
    tanh_mean: Whether to transform the mean (via tanh) before passing it to
      the distribution.
    fixed_scale: Whether to use a fixed variance.
    use_tfd_independent: Whether to use tfd.Independent or
      tfd.MultivariateNormalDiag class
    w_init: Initialization for linear layer weights.
    b_init: Initialization for linear layer biases.
docstr_tok: ["Initialization",".","Args",":","num_dimensions",":","Number","of","dimensions","of","MVN","distribution",".","init_scale",":","Initial","standard","deviation",".","min_scale",":","Minimum","standard","deviation",".","tanh_mean",":","Whether","to","transform","the","mean","(","via","tanh",")","before","passing","it","to","the","distribution",".","fixed_scale",":","Whether","to","use","a","fixed","variance",".","use_tfd_independent",":","Whether","to","use","tfd.Independent","or","tfd.MultivariateNormalDiag","class","w_init",":","Initialization","for","linear","layer","weights",".","b_init",":","Initialization","for","linear","layer","biases","."]
code: def __init__(
      self,
      num_dimensions: int,
      init_scale: float = 0.3,
      min_scale: float = 1e-6,
      tanh_mean: bool = False,
      fixed_scale: bool = False,
      use_tfd_independent: bool = False,
      w_init: snt_init.Initializer = tf.initializers.VarianceScaling(
          1e-4
      ),
      b_init: snt_init.Initializer = tf.initializers.Zeros(),
  ):
      """Initialization.
  
      Args:
        num_dimensions: Number of dimensions of MVN distribution.
        init_scale: Initial standard deviation.
        min_scale: Minimum standard deviation.
        tanh_mean: Whether to transform the mean (via tanh) before passing it to
          the distribution.
        fixed_scale: Whether to use a fixed variance.
        use_tfd_independent: Whether to use tfd.Independent or
          tfd.MultivariateNormalDiag class
        w_init: Initialization for linear layer weights.
        b_init: Initialization for linear layer biases.
      """
      super().__init__(name="MultivariateNormalDiagHead")
      self._init_scale = init_scale
      self._min_scale = min_scale
      self._tanh_mean = tanh_mean
      self._mean_layer = snt.Linear(
          num_dimensions, w_init=w_init, b_init=b_init
      )
      self._fixed_scale = fixed_scale
  
      if not fixed_scale:
          self._scale_layer = snt.Linear(
              num_dimensions, w_init=w_init, b_init=b_init
          )
      self._use_tfd_independent = use_tfd_independent
  
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 MultivariateNormalDiagHead object from inherited base classes: snt with example: obj = MultivariateNormalDiagHead(num_dimensions, init_scale, min_scale, tanh_mean, fixed_scale, use_tfd_independent, w_init, b_init)