code: def _sub_diag(nonmatrix):
      """Get the first sub-diagonal of a shape [N, N, ...] 'non matrix'."""
      with tf.name_scope("sub_matrix"):
          # TODO(b/143702351) Once array_ops.matrix_diag_part_v3 is ready and exposed,
          # replace the call to matrix_diag_part_v2 below with tf.linalg.matrix_diag.
          # We can also stop special casing for matrix_dim < 2 at that point.
          # Until then, OpError raised for 1x1 matricies without static shape.
          # In fact, non-static shape breaks matrix_diag_part_v2, so we must raise
          # this message now.
          # See http://b/138403336 for the TF issue tracker.
          if not tensorshape_util.is_fully_defined(nonmatrix.shape[:2]):
              raise ValueError(
                  "`inverse_temperatures did not have statically defined shape, "
                  "which breaks tracking of is_swap_{proposed,accepted}.  "
                  "Please provide an inverse_temperatures with statically known shape."
              )
  
          # The sub-matrix of a 1x1 matrix is not defined (throws exception), so in
          # this special case return an empty matrix.
          # TODO(b/143702351) Remove this special case handling once
          # matrix_diag_part_v3 is ready.
          matrix_dim = ps.size0(nonmatrix)
          if matrix_dim is not None and matrix_dim < 2:
              # Shape is [..., 0], so returned tensor is empty, thus contains no
              # values...and therefore the fact that we use 'ones' doesn't matter.
              shape = ps.pad(
                  ps.shape(nonmatrix)[2:],
                  paddings=[[0, 1]],
                  constant_values=0,
              )
              matrix_sub_diag = tf.cast(tf.ones(shape), nonmatrix.dtype)
  
          else:
              # Get first sub-diagonal.  `padding_value` is not used (since matrix is
              # square), but is required for the API since this is raw gen_array_ops.
              matrix_sub_diag = tf.raw_ops.MatrixDiagPartV2(
                  input=distribution_util.rotate_transpose(
                      nonmatrix, shift=-2
                  ),
                  k=ps.convert_to_shape_tensor(-1, dtype=tf.int32),
                  padding_value=tf.cast(0.0, dtype=nonmatrix.dtype),
              )
  
          return distribution_util.rotate_transpose(
              matrix_sub_diag, shift=1
          )
  
imports: import collections
  import warnings
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import assert_util
  from tensorflow_probability.python.internal import broadcast_util
  from tensorflow_probability.python.internal import distribution_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.internal import unnest
  from tensorflow_probability.python.mcmc import kernel
  from tensorflow_probability.python.mcmc.internal import util
label_desc: Use image node_id 9 for calling a global function with example usage: _sub_diag(nonmatrix) and returns: distribution_util