class_docstr: Module that produces a multivariate normal distribution using tfd.Independent or tfd.MultivariateNormalDiag.
code: def __call__(self, inputs: tf.Tensor) -> tfd.Distribution:
      zero = tf.constant(0, dtype=inputs.dtype)
      mean = self._mean_layer(inputs)
  
      if self._fixed_scale:
          scale = tf.ones_like(mean) * self._init_scale
      else:
          scale = tf.nn.softplus(self._scale_layer(inputs))
          scale *= self._init_scale / tf.nn.softplus(zero)
          scale += self._min_scale
  
      # Maybe transform the mean.
      if self._tanh_mean:
          mean = tf.tanh(mean)
  
      if self._use_tfd_independent:
          dist = tfd.Independent(tfd.Normal(loc=mean, scale=scale))
      else:
          dist = tfd.MultivariateNormalDiag(loc=mean, scale_diag=scale)
  
      return dist
  
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 2 for calling the MultivariateNormalDiagHead obj's underlying member method code with example usage: obj.__call__(inputs) and returns: dist