code: def calc_rel_pos(n):
      pos = tf.stack(meshgrid(tf.range(n), tf.range(n), indexing="ij"))
      pos = Rearrange("n i j -> (i j) n")(
          pos
      )  # [n*n, 2] pos[n] = (i, j)
      rel_pos = (
          pos[None, :] - pos[:, None]
      )  # [n*n, n*n, 2] rel_pos[n, m] = (rel_i, rel_j)
      rel_pos += (
          n - 1
      )  # shift value range from [-n+1, n-1] to [0, 2n-2]
      return rel_pos
  
file_path: lambda-networks/lambda_networks/tfkeras.py
imports: import tensorflow
  from einops.layers.tensorflow import Rearrange
  from tensorflow.keras.layers import Conv2D, BatchNormalization, Conv3D, ZeroPadding3D, Softmax, Lambda, Add, Layer
  from tensorflow.keras import initializers
  from tensorflow import einsum, nn, meshgrid
label_desc: Use image node_id 3 for calling a global function with example usage: calc_rel_pos(n) and returns: rel_pos