code: def _trainable_linear_operator_block(
      operators, block_dims=None, batch_shape=(), dtype=None, name=None
  ):
      """Builds a trainable blockwise `tf.linalg.LinearOperator`.
  
      This function returns a trainable blockwise `LinearOperator`. If `operators`
      is a flat list, it is interpreted as blocks along the diagonal of the
      structure and an instance of `tf.linalg.LinearOperatorBlockDiag` is returned.
      If `operators` is a doubly nested list, then a
      `tf.linalg.LinearOperatorBlockLowerTriangular` instance is returned, with
      the block in row `i` column `j` (`i >= j`) given by `operators[i][j]`.
      The `operators` list may contain `LinearOperator` instances, `LinearOperator`
      subclasses, or callables defining custom constructors (see example below).
      The dimensions of the blocks are given by `block_dims`; this argument may be
      omitted if `operators` contains only `LinearOperator` instances.
  
      Args:
        operators: A list or tuple containing `LinearOperator` subclasses,
          `LinearOperator` instances, and/or callables returning
          `(init_fn, apply_fn)` pairs. If the list is flat, a
          `tf.linalg.LinearOperatorBlockDiag` instance is returned. Otherwise, the
          list must be singly nested, with the
          first element of length 1, second element of length 2, etc.; the
          elements of the outer list are interpreted as rows of a lower-triangular
          block structure, and a `tf.linalg.LinearOperatorBlockLowerTriangular`
          instance is returned. Callables contained in the lists must take two
          arguments -- `shape`, the shape of the parameter instantiating the
          `LinearOperator`, and `dtype`, the `tf.dtype` of the `LinearOperator` --
          and return a further pair of callables representing a stateless trainable
          operator (see example below).
        block_dims: List or tuple of integers, representing the sizes of the blocks
          along one dimension of the (square) blockwise `LinearOperator`. If
          `operators` contains only `LinearOperator` instances, `block_dims` may be
          `None` and the dimensions are inferred.
        batch_shape: Batch shape of the `LinearOperator`.
        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.
  
      ### Examples
  
      To build a 5x5 trainable `LinearOperatorBlockDiag` given `LinearOperator`
      subclasses and `block_dims`:
  
      ```python
      op = build_trainable_linear_operator_block(
        operators=(tf.linalg.LinearOperatorDiag,
                   tf.linalg.LinearOperatorLowerTriangular),
        block_dims=[3, 2],
        dtype=tf.float32)
      ```
  
      If `operators` contains only `LinearOperator` instances, the `block_dims`
      argument is not necessary:
  
      ```python
      # Builds a 6x6 `LinearOperatorBlockDiag` with batch shape `(4,).
      op = build_trainable_linear_operator_block(
        operators=(tf.linalg.LinearOperatorDiag(tf.Variable(tf.ones((4, 3)))),
                   tf.linalg.LinearOperatorFullMatrix([4.]),
                   tf.linalg.LinearOperatorIdentity(2)))
  
      ```
  
      A custom operator constructor may be specified as a callable taking
      arguments `shape` and `dtype`, and returning a pair of callables
      `(init_fn, apply_fn)` describing a parameterized operator, with the following
      signatures:
  
      ```python
      raw_parameters = init_fn(seed)
      linear_operator = apply_fn(raw_parameters)
      ```
  
      For example, to define a custom initialization for a diagonal operator:
  
      ```python
      import functools
  
      def diag_operator_with_uniform_initialization(shape, dtype):
        init_fn = functools.partial(
            samplers.uniform, shape, maxval=2., dtype=dtype)
        apply_fn = lambda scale_diag: tf.linalg.LinearOperatorDiag(
            scale_diag, is_non_singular=True)
        return init_fn, apply_fn
  
      # Build an 8x8 `LinearOperatorBlockLowerTriangular`, with our custom diagonal
      # operator in the upper left block, and `LinearOperator` subclasses in the
      # lower two blocks.
      op = build_trainable_linear_operator_block(
        operators=(diag_operator_with_uniform_initialization,
                   (tf.linalg.LinearOperatorFullMatrix,
                    tf.linalg.LinearOperatorLowerTriangular)),
        block_dims=[4, 4],
        dtype=tf.float64)
      ```
  
      """
      with tf.name_scope(name or "trainable_linear_operator_block"):
          operator_instances = [
              op
              for op in nest.flatten(operators)
              if isinstance(op, tf.linalg.LinearOperator)
          ]
          if block_dims is None and len(operator_instances) < len(
              nest.flatten(operators)
          ):
              # If `operator_instances` contains fewer elements than `operators`,
              # then some elements of `operators` are not instances of `LinearOperator`.
              raise ValueError(
                  "Argument `block_dims` must be defined unless "
                  "`operators` contains only `tf.linalg.LinearOperator` "
                  "instances."
              )
  
          batch_shape = ps.cast(batch_shape, tf.int32)
          if dtype is None:
              dtype = dtype_util.common_dtype(operator_instances)
  
          def convert_operator(path, op):
              if isinstance(op, tf.linalg.LinearOperator):
                  return op
              if len(set(path)) == 1:  # for operators on the diagonal
                  shape = ps.concat(
                      [batch_shape, [block_dims[path[0]]]], axis=0
                  )
              else:
                  shape = ps.concat(
                      [
                          batch_shape,
                          [block_dims[path[0]], block_dims[path[1]]],
                      ],
                      axis=0,
                  )
              if op in _OPERATOR_COROUTINES:
                  operator = yield from _OPERATOR_COROUTINES[op](
                      shape=shape, dtype=dtype
                  )
              else:  # Custom stateless constructor.
                  init_fn, apply_fn = op(shape=shape, dtype=dtype)
                  raw_params = yield trainable_state_util.Parameter(
                      init_fn
                  )
                  operator = apply_fn(raw_params)
              return operator
  
          # Build a structure of component trainable LinearOperators.
          operator_blocks = (
              yield from nest_util.map_structure_coroutine(
                  convert_operator, operators, _with_tuple_paths=True
              )
          )
          paths = nest.yield_flat_paths(operators)
          if all(len(p) == 1 for p in paths):
              return tf.linalg.LinearOperatorBlockDiag(
                  operator_blocks, is_non_singular=True
              )
          elif all(len(p) == 2 for p in paths):
              return tf.linalg.LinearOperatorBlockLowerTriangular(
                  operator_blocks, is_non_singular=True
              )
          else:
              raise ValueError(
                  "Argument `operators` must be a flat or singly-nested sequence."
              )
  
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 1 for calling a global function with example usage: _trainable_linear_operator_block(operators, block_dims, batch_shape, dtype, name) and returns: operator, tf, op, tf