code: def _trainable_linear_operator_diag(
      shape,
      scale_initializer=1e-2,
      diag_bijector=None,
      dtype=None,
      name=None,
  ):
      """Build a trainable `LinearOperatorDiag` instance.
  
      Args:
        shape: Shape of the `LinearOperator`, equal to `[b0, ..., bn, d]`, where
          `b0...bn` are batch dimensions and `d` is the length of the diagonal.
        scale_initializer: Variables are initialized with samples from
          `Normal(0, scale_initializer)`.
        diag_bijector: Bijector to apply to the diagonal of the operator.
        dtype: `tf.dtype` of the `LinearOperator`.
        name: str, name for `tf.name_scope`.
      Yields:
        *parameters: sequence of `trainable_state_util.Parameter` namedtuples.
          These are intended to be consumed by
          `trainable_state_util.as_stateful_builder` and
          `trainable_state_util.as_stateless_builder` to define stateful and
          stateless variants respectively.
      """
      with tf.name_scope(name or "trainable_linear_operator_diag"):
          if dtype is None:
              dtype = dtype_util.common_dtype(
                  [scale_initializer], dtype_hint=tf.float32
              )
          scale_initializer = tf.convert_to_tensor(
              scale_initializer, dtype=dtype
          )
  
          diag_bijector = diag_bijector or _DefaultScaleDiagonal()
          scale_diag = yield trainable_state_util.Parameter(
              init_fn=lambda seed: diag_bijector.forward(  # pylint: disable=g-long-lambda
                  samplers.normal(
                      mean=0.0,
                      stddev=scale_initializer,
                      shape=shape,
                      dtype=dtype,
                      seed=seed,
                  )
              ),
              name="scale_diag",
              constraining_bijector=diag_bijector,
          )
          return tf.linalg.LinearOperatorDiag(
              scale_diag, is_non_singular=True
          )
  
imports: import functools
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import bijector
  from tensorflow_probability.python.bijectors import fill_scale_tril
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import nest_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import trainable_state_util
  from tensorflow.python.util import nest
label_desc: Use image node_id 3 for calling a global function with example usage: _trainable_linear_operator_diag(shape, scale_initializer, diag_bijector, dtype, name) and returns: tf