code: def make_kernel_bias_posterior_mvn_diag(
      kernel_shape,
      bias_shape,
      kernel_initializer=None,
      bias_initializer=None,
      kernel_batch_ndims=0,
      bias_batch_ndims=0,  # pylint: disable=unused-argument
      dtype=tf.float32,
      kernel_name="posterior_kernel",
      bias_name="posterior_bias",
  ):
      """Create learnable posterior for Variational layers with kernel and bias.
  
      Args:
        kernel_shape: ...
        bias_shape: ...
        kernel_initializer: ...
          Default value: `None` (i.e., `tf_keras.initializers.glorot_uniform()`).
        bias_initializer: ...
          Default value: `None` (i.e., `tf_keras.initializers.zeros()`).
        kernel_batch_ndims: ...
          Default value: `0`.
        bias_batch_ndims: ...
          Default value: `0`.
        dtype: ...
          Default value: `tf.float32`.
        kernel_name: ...
          Default value: `"posterior_kernel"`.
        bias_name: ...
          Default value: `"posterior_bias"`.
  
      Returns:
        kernel_and_bias_distribution: ...
      """
      if kernel_initializer is None:
          kernel_initializer = nn_init_lib.glorot_uniform()
      if bias_initializer is None:
          bias_initializer = tf_keras.initializers.zeros()
      make_loc = lambda init_fn, shape, batch_ndims, name: tf.Variable(  # pylint: disable=g-long-lambda
          _try_call_init_fn(init_fn, shape, dtype, batch_ndims),
          name=name + "_loc",
      )
      # Setting the initial scale to a relatively small value causes the `loc` to
      # quickly move toward a lower loss value.
      make_scale = lambda shape, name: TransformedVariable(  # pylint: disable=g-long-lambda
          tf.fill(shape, value=tf.constant(1e-3, dtype=dtype)),
          Chain([Shift(1e-5), Softplus()]),
          name=name + "_scale",
      )
      return JointDistributionSequential(
          [
              Independent(
                  Normal(
                      loc=make_loc(
                          kernel_initializer,
                          kernel_shape,
                          kernel_batch_ndims,
                          kernel_name,
                      ),
                      scale=make_scale(kernel_shape, kernel_name),
                  ),
                  reinterpreted_batch_ndims=ps.size(kernel_shape),
                  name=kernel_name,
              ),
              Independent(
                  Normal(
                      loc=make_loc(
                          bias_initializer,
                          bias_shape,
                          kernel_batch_ndims,
                          bias_name,
                      ),
                      scale=make_scale(bias_shape, bias_name),
                  ),
                  reinterpreted_batch_ndims=ps.size(bias_shape),
                  name=bias_name,
              ),
          ]
      )
  
file_path: probability/tensorflow_probability/python/experimental/nn/util/kernel_bias.py
imports: import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors.chain import Chain
  from tensorflow_probability.python.bijectors.shift import Shift
  from tensorflow_probability.python.bijectors.softplus import Softplus
  from tensorflow_probability.python.distributions.categorical import Categorical
  from tensorflow_probability.python.distributions.independent import Independent
  from tensorflow_probability.python.distributions.joint_distribution_sequential import JointDistributionSequential
  from tensorflow_probability.python.distributions.mixture_same_family import MixtureSameFamily
  from tensorflow_probability.python.distributions.normal import Normal
  from tensorflow_probability.python.distributions.sample import Sample
  from tensorflow_probability.python.experimental.nn import initializers
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import tf_keras
  from tensorflow_probability.python.util.deferred_tensor import TransformedVariable
label_desc: Use image node_id 3 for calling a global function with example usage: make_kernel_bias_posterior_mvn_diag(kernel_shape, bias_shape, kernel_initializer, bias_initializer, kernel_batch_ndims, bias_batch_ndims, dtype, kernel_name, bias_name) and returns: JointDistributionSequential