code: def add_self_loops_indices(indices, n_nodes=None):
      """
      Given the indices of a square SparseTensor, adds the diagonal entries (i, i)
      and returns the reordered indices.
      :param indices: Tensor of rank 2, the indices to a SparseTensor.
      :param n_nodes: the size of the n_nodes x n_nodes SparseTensor indexed by
      the indices. If `None`, n_nodes is calculated as the maximum entry in the
      indices plus 1.
      :return: Tensor of rank 2, the indices to a SparseTensor.
      """
      n_nodes = (
          tf.reduce_max(indices) + 1 if n_nodes is None else n_nodes
      )
      row, col = indices[..., 0], indices[..., 1]
      mask = tf.ensure_shape(row != col, row.shape)
      sl_indices = tf.range(n_nodes, dtype=row.dtype)[:, None]
      sl_indices = tf.repeat(sl_indices, 2, -1)
      indices = tf.concat((indices[mask], sl_indices), 0)
      dummy_values = tf.ones_like(indices[:, 0])
      indices, _ = gen_sparse_ops.sparse_reorder(
          indices, dummy_values, (n_nodes, n_nodes)
      )
      return indices
  
file_path: spektral/spektral/layers/ops/sparse.py
imports: import tensorflow
  from tensorflow.python.ops import gen_sparse_ops
  from .None import ops
label_desc: Use image node_id 2 for calling a global function with example usage: add_self_loops_indices(indices, n_nodes) and returns: indices