code: def sparse_cosine_similarity(x, n_nodes, mask, i):
      mask = tf.cast(mask, tf.int32)
      leader_idx = tf.where(mask)
  
      # Number of nodes in each graph
      ns = tf.math.segment_sum(tf.ones_like(i), i)
      ks = tf.math.segment_sum(mask, i)
  
      # s will be a block-diagonal matrix where entry i,j is the cosine
      # similarity between node i and leader j.
      # The code below creates the indices of the sparse block-diagonal matrix
      # Row indices of the block-diagonal S
      starts = tf.cumsum(ns) - ns
      starts = tf.repeat(starts, ks)
      stops = tf.cumsum(ns)
      stops = tf.repeat(stops, ks)
      index_n = tf.ragged.range(starts, stops).flat_values
  
      # Column indices of the block-diagonal S
      index_k = tf.repeat(leader_idx, tf.repeat(ns, ks))
      index_k_for_s = tf.repeat(
          tf.range(tf.reduce_sum(ks)), tf.repeat(ns, ks)
      )
  
      # Make index int64
      index_n = tf.cast(index_n, tf.int64)
      index_k = tf.cast(index_k, tf.int64)
      index_k_for_s = tf.cast(index_k_for_s, tf.int64)
  
      # Compute similarity between nodes and leaders
      x_n = tf.gather(x, index_n)
      x_n_norm = tf.norm(x_n, axis=-1)
      x_k = tf.gather(x, index_k)
      x_k_norm = tf.norm(x_k, axis=-1)
      values = tf.reduce_sum(x_n * x_k, -1) / (x_n_norm * x_k_norm)
  
      # Create a sparse tensor for S
      indices = tf.stack((index_n, index_k_for_s), 1)
      s = tf.SparseTensor(
          values=values,
          indices=indices,
          dense_shape=(n_nodes, tf.reduce_sum(ks)),
      )
      s = tf.sparse.reorder(s)
  
      return s
  
file_path: spektral/spektral/layers/pooling/la_pool.py
imports: import tensorflow
  from scipy import sparse
  from tensorflow.keras import backend
  from spektral.layers import ops
  from spektral.layers.pooling.src import SRCPool
label_desc: Use image node_id 3 for calling a global function with example usage: sparse_cosine_similarity(x, n_nodes, mask, i) and returns: s