class_docstr: Provides learnable NxN relative attention bias.
  
  Attributes:
    num_heads: Number of heads for which to provide relative attention.
    nd_shape: Shape for which to provided relative attention bias. For instance,
      for images we we would provide a 2D shape. Note that batch and feature
      dimensions should be excluded here.
    initializer: Initializer for the bias.
class_docstr_tok: ["Provides","learnable","NxN","relative","attention","bias",".","Attributes",":","num_heads",":","Number","of","heads","for","which","to","provide","relative","attention",".","nd_shape",":","Shape","for","which","to","provided","relative","attention","bias",".","For","instance",",","for","images","we","we","would","provide","a","2D","shape",".","Note","that","batch","and","feature","dimensions","should","be","excluded","here",".","initializer",":","Initializer","for","the","bias","."]
docstr: Creates relative attention bias that factorizes over dimensions.
  
  length = prod(nd_shape)
  
  Returns:
    Bias of shape `[num_heads, length, length]`.
docstr_tok: ["Creates","relative","attention","bias","that","factorizes","over","dimensions",".","length","=","prod","(","nd_shape",")","Returns",":","Bias","of","shape","`","[","num_heads",",","length",",","length","]","`","."]
code: def __call__(self) -> jnp.ndarray:
      """Creates relative attention bias that factorizes over dimensions.
  
      length = prod(nd_shape)
  
      Returns:
        Bias of shape `[num_heads, length, length]`.
      """
      length = np.prod(self.nd_shape)
      tile = 1
      biases = []
      for i, l in enumerate(self.nd_shape):
          # Relative attention in every dimension separately.
          if l > 1:
              new_bias = self.relative_attn_bias(
                  l, self.num_heads, f"bias_{i}"
              )
              repeat = length // (tile * l)
              if repeat > 1:
                  new_bias = new_bias[:, :, jnp.newaxis, :, jnp.newaxis]
                  new_bias = jnp.tile(
                      new_bias, [1, tile, repeat, tile, repeat]
                  )
                  new_bias = jnp.reshape(
                      new_bias, [self.num_heads, length, length]
                  )
              elif tile > 1:
                  new_bias = jnp.tile(new_bias, [1, tile, tile])
              tile *= l
              biases.append(new_bias)
  
      return sum(biases)
  
imports: import functools
  from typing import Callable, Optional, Sequence, Tuple
  import flax.linen
  import jax
  import jax.numpy
  import numpy
  from scenic.model_lib.layers import nn_layers
label_desc: Use image node_id 1 for calling the RelativeAttentionBias obj's underlying member method code with example usage: obj.__call__() and returns: sum