code: def compute_relative_positions(
      query_spatial_shape, key_spatial_shape, spatial_axis=None
  ):
      """Generate relative positions of queries and keys.
  
  
      For relative attention, the pairwise positional distance between each query
      and key point is used in the attention weight computation. This function
      generates the positional distances between each query-key pair, given the
      offset of first position in the query with respect to first position in the
      key.
  
      For example, if the query and key are 1d and query has 2 entries and the key
      has 3 entries, the relative distance matrix is:
        [[0, 1, 2],
         [-1, 0, 1]]
      where each [i, j] entry = j - i (j = key index, i = query index). Note that
      the values in this matrix are being used by an embedding lookup, so we shift
      them such that the smallest index is zero:
        [[1, 2, 3],
         [0, 1, 2]]
  
      This function produces the multi-dimensional distance for a query and key.
      It factorizes the distance computation such that there is a positional
      distance per dimension. An input with 3 dimensions will have a total of
      3 distances, 1 per dimension.
  
      Args:
        query_spatial_shape: tuple; Indicating the spatial shape of the query.
        key_spatial_shape: tuple; Indicating the spatial shape of the key.
        spatial_axis: tuple; The axis over which the distance is calculated. Default
          is None, which means distances over all axis is calculated.
  
      Returns:
        a numpy (np) int array of shape [len(spatial_axis),
          query_spatial_shape(spatial_axis), key_spatial_shape(spatial_axis)]
          holding the distance between each query  and key pair across dimensions
          that are determined by `spatial_axis`,  where the query and key are
          indexed by their position. The smallest value in the array is zero.
      """
      assert len(query_spatial_shape) == len(key_spatial_shape)
      if spatial_axis is None:
          spatial_axis = range(len(query_spatial_shape))
      for sa in spatial_axis:
          if not 0 <= sa < len(query_spatial_shape):
              raise ValueError(
                  "Element of `spatial_axis` should be between 0 and "
                  "length of `query_spatial_shape`."
              )
  
      num_dims = len(spatial_axis)
      # Keep only dimensions we are iterested in.
      query_spatial_shape = tuple(
          [query_spatial_shape[a] for a in spatial_axis]
      )
      key_spatial_shape = tuple(
          [key_spatial_shape[a] for a in spatial_axis]
      )
  
      total_queries = np.prod(query_spatial_shape)
  
      total_keys = np.prod(key_spatial_shape)
      # A distance per dimension in the flattened query-key arrays.
  
      relative_positions = np.empty(
          (num_dims, total_queries, total_keys), dtype=np.int32
      )
  
      # Convert flattened indices to multi-dimension coordinate indices.
      coordinates_query = np.unravel_index(
          range(total_queries), query_spatial_shape
      )
      coordinates_key = np.unravel_index(
          range(total_keys), key_spatial_shape
      )
  
      # Compute distances between each query-key point.
      for dim in range(num_dims):
          for flat_index_query in range(total_queries):
              for flat_index_key in range(total_keys):
                  relative_positions[
                      dim, flat_index_query, flat_index_key
                  ] = (
                      coordinates_key[dim][flat_index_key]
                      - coordinates_query[dim][flat_index_query]
                  )
          relative_positions[dim] = relative_positions[dim]
  
      # These indices are being used by an embedding lookup, so shift the indices
      # such that the smallest index is zero.
      relative_positions -= np.amin(
          relative_positions, axis=(1, 2), keepdims=True
      )
      # Reshape to original dim.
      relative_positions = relative_positions.reshape(
          (num_dims,) + query_spatial_shape + key_spatial_shape
      )
      return relative_positions
  
file_path: scenic/scenic/model_lib/layers/nn_ops.py
imports: from typing import Callable, Sequence
  import flax.linen
  import jax
  from jax import lax
  import jax.numpy
  import numpy
label_desc: Use image node_id 3 for calling a global function with example usage: compute_relative_positions(query_spatial_shape, key_spatial_shape, spatial_axis) and returns: relative_positions