code: def degree_matrix(A, return_sparse_batch=False):
      """
      Computes the degree matrix of A, deals with sparse A and batch mode
      automatically.
      :param A: Tensor or SparseTensor with rank k = {2, 3}.
      :param return_sparse_batch: if operating in batch mode, return a
      SparseTensor. Note that the sparse degree Tensor returned by this function
      cannot be used for sparse matrix multiplication afterwards.
      :return: SparseTensor of rank k.
      """
      D = degrees(A)
  
      batch_mode = K.ndim(D) == 2
      N = tf.shape(D)[-1]
      batch_size = tf.shape(D)[0] if batch_mode else 1
  
      inner_index = tf.tile(
          tf.stack([tf.range(N)] * 2, axis=1), (batch_size, 1)
      )
      if batch_mode:
          if return_sparse_batch:
              outer_index = ops.repeat(
                  tf.range(batch_size),
                  tf.ones(batch_size) * tf.cast(N, tf.float32),
              )
              indices = tf.concat(
                  [outer_index[:, None], inner_index], 1
              )
              dense_shape = (batch_size, N, N)
          else:
              return tf.linalg.diag(D)
      else:
          indices = inner_index
          dense_shape = (N, N)
  
      indices = tf.cast(indices, tf.int64)
      values = tf.reshape(D, (-1,))
      return tf.SparseTensor(indices, values, dense_shape)
  
imports: import tensorflow
  from tensorflow.keras import backend
  from .None import ops
label_desc: Use image node_id 3 for calling a global function with example usage: degree_matrix(A, return_sparse_batch) and returns: tf, tf