code: def _rolling_window_indices(size, rw_size, num_valid_entries):
      """Returns the rolling windows indices and mask for valid ones.
  
      When size = 3, rw_size = 2, returns [[0, 1], [1, 2], [2, 0]]. When size = 2,
      rw_size = 3, returns [[0, 1, 0], [1, 0, 1]].
  
      When num_valid_entries = 2, the first returns [[0, 1], [1, 0], [0, 1]] and the
      first 2 are valid with mask as [True, True, False].
  
      Args:
        size: A scalar int `Tensor` for the size.
        rw_size: A scalar int `Tensor` for the rw_size.
        num_valid_entries: A 1-D `Tensor` with shape [batch_size] representing the
          number of valid entries for each instance in a batch.
  
      Returns:
        A tuple of Tensors (batch_rw_indices, batch_indices_mask). The first has
        shape [batch_size, size, rw_size] and the second has shape [batch_size,
        size].
      """
      with tf.compat.v1.name_scope(name="rolling_window_indices"):
          # shape = [size, rw_size] with value [[0, 1, ...], [1, 2, ...], ...].
          rw_indices = tf.expand_dims(
              tf.range(rw_size), 0
          ) + tf.expand_dims(tf.range(size), 1)
          # shape = [batch_size, size, rw_size]. Make batch_size copies.
          batch_size = tf.shape(input=num_valid_entries)[0]
          batch_rw_indices = tf.tile(
              tf.expand_dims(rw_indices, 0),
              multiples=[batch_size, 1, 1],
          )
          # Mark the first n indices as valid where n = num_valid_entries.
          batch_indices_mask = tf.less(
              tf.reduce_min(input_tensor=batch_rw_indices, axis=2),
              tf.reshape(num_valid_entries, [-1, 1]),
          )
          # Mod the indices to the range of num_valid_entries.
          num_valid_entries = tf.where(
              tf.less(num_valid_entries, 1),
              tf.ones_like(num_valid_entries),
              num_valid_entries,
          )
          batch_rw_indices = tf.math.mod(
              batch_rw_indices,
              tf.reshape(num_valid_entries, [-1, 1, 1]),
          )
          return batch_rw_indices, batch_indices_mask
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import abc
  import six
  import tensorflow
  from tensorflow import estimator
  from tensorflow.python.util import function_utils
  from tensorflow_ranking.python import feature
  from tensorflow_ranking.python import utils
label_desc: Use image node_id 2 for calling a global function with example usage: _rolling_window_indices(size, rw_size, num_valid_entries) and returns: batch_rw_indices,  batch_indices_mask