class_docstr: Module that produces a tfd.MultivariateNormalDiag distribution.
class_docstr_tok: ["Module","that","produces","a","tfd.MultivariateNormalDiag","distribution","."]
docstr: Initialization.
  
  Args:
    num_dimensions: Number of dimensions of MVN distribution.
    init_scale: Initial standard deviation.
    min_scale: Minimum standard deviation.
    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",".","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,
      w_init: hk_init.Initializer = hk_init.VarianceScaling(1e-4),
      b_init: hk_init.Initializer = hk_init.Constant(0.0),
  ):
      """Initialization.
  
      Args:
        num_dimensions: Number of dimensions of MVN distribution.
        init_scale: Initial standard deviation.
        min_scale: Minimum standard deviation.
        w_init: Initialization for linear layer weights.
        b_init: Initialization for linear layer biases.
      """
      super().__init__(name="MultivariateNormalDiagHead")
      self._min_scale = min_scale
      self._init_scale = init_scale
      self._loc_layer = hk.Linear(
          num_dimensions, w_init=w_init, b_init=b_init
      )
      self._scale_layer = hk.Linear(
          num_dimensions, w_init=w_init, b_init=b_init
      )
  
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 MultivariateNormalDiagHead object from inherited base classes: hk with example: obj = MultivariateNormalDiagHead(num_dimensions, init_scale, min_scale, w_init, b_init)