docstr: Produce relative position embedding attention biases.
  Returns:
    output: `(heads, q_len, k_len)` attention bias
docstr_tok: ["Produce","relative","position","embedding","attention","biases",".","Returns",":","output",":","`","(","heads",",","q_len",",","k_len",")","`","attention","bias"]
code: def __call__(self, qlen, klen, heads, num_buckets):
      """Produce relative position embedding attention biases.
      Returns:
        output: `(heads, q_len, k_len)` attention bias
      """
      context_position = np.arange(qlen, dtype=jnp.int32)[:, None]
      memory_position = np.arange(klen, dtype=jnp.int32)[None, :]
      relative_position = (
          memory_position - context_position
      )  # shape (qlen, klen)
      rp_bucket = self._relative_position_bucket(relative_position)
      relative_attention_bias = hk.get_parameter(
          "rel_embedding",
          [heads, num_buckets],
          init=hk.initializers.TruncatedNormal(stddev=0.02),
      )
      # Instead of using a slow gather, we create a leading-dimension one-hot
      # array from rp_bucket and use it to perform the gather-equivalent via a
      # contraction, i.e.:
      # (num_head, num_buckets) x (num_buckets one-hot, qlen, klen).
      # This is equivalent to relative_attention_bias[:, rp_bucket]
      bcast_iota = jax.lax.broadcasted_iota(
          jnp.int32, (num_buckets, 1, 1), 0
      )
      rp_bucket_one_hot = jnp.array(
          rp_bucket[jnp.newaxis, Ellipsis] == bcast_iota
      ).astype(relative_attention_bias.dtype)
      # --> shape (qlen, klen, num_heads)
      values = jax.lax.dot_general(
          relative_attention_bias,
          rp_bucket_one_hot,
          (((1,), (0,)), ((), ())),  # rhs, lhs contracting dims
      )  # no batched dims
      return values
  
imports: import haiku
  import jax
  import jax.numpy
  import numpy
  from einops import rearrange, repeat
  from mesh_transformer.util import f_psum, g_psum, maybe_shard, head_print
  from jax.experimental import PartitionSpec
  from jax.experimental.maps import thread_resources
label_desc: Use image node_id 2 for calling the RelativePositionEmbs obj's underlying member method code with example usage: obj.__call__(qlen, klen, heads, num_buckets) and returns: values