code: def split_arg_into_blocks(block_dims, block_dims_fn, arg, axis=-1):
      """Split `x` into blocks matching `operators`'s `domain_dimension`.
  
      Specifically, if we have a blockwise lower-triangular matrix, with block
      sizes along the diagonal `[M_j, M_j] j = 0,1,2..J`,  this method splits `arg`
      on `axis` into `J` tensors, whose shape at `axis` is `M_j`.
  
      Args:
        block_dims: Iterable of `TensorShapes`.
        block_dims_fn: Callable returning an iterable of `Tensor`s.
        arg: `Tensor`. `arg` is split into `J` tensors.
        axis: Python `Integer` representing the axis to split `arg` on.
  
      Returns:
        A list of `Tensor`s.
      """
      block_sizes = [dim.value for dim in block_dims]
      if any(d is None for d in block_sizes):
          block_sizes = block_dims_fn()
      return array_ops.split(arg, block_sizes, axis=axis)
  
imports: import numpy
  from tensorflow_probability.python.internal.backend.numpy import dtype
  from tensorflow_probability.python.internal.backend.numpy import ops
  from tensorflow_probability.python.internal.backend.numpy import ops
  from tensorflow_probability.python.internal.backend.numpy import numpy_array
  from tensorflow_probability.python.internal.backend.numpy import debugging
  from tensorflow_probability.python.internal.backend.numpy import control_flow
  from tensorflow_probability.python.internal.backend.numpy import linalg_impl
  from tensorflow_probability.python.internal.backend.numpy import numpy_math
  from tensorflow_probability.python.internal.backend.numpy import ops
  from tensorflow_probability.python.internal.backend.numpy import nest
  import numpy
  from tensorflow_probability.python.internal.backend.numpy import linalg_impl
  from tensorflow_probability.python.internal.backend.numpy import ops
  from tensorflow_probability.python.internal.backend.numpy.gen import tensor_shape
  from tensorflow_probability.python.internal.backend.numpy import private
label_desc: Use image node_id 19 for calling a global function with example usage: split_arg_into_blocks(block_dims, block_dims_fn, arg, axis) and returns: array_ops