# Path: flaxformer/types.py



from typing import Any, Callable
from flax import linen as nn
from flax.linen import partitioning
from jax import lax
from flaxformer.types import Array
import jax.numpy as jnp
import numpy as np
        (num_buckets - max_exact)).astype(jnp.int32)
    val_if_large = jnp.minimum(val_if_large, num_buckets - 1)
    ret += jnp.where(is_small, n, val_if_large)
    return ret

  def full_att_rp_bucket(self, qlen, klen, bidirectional=True):
    """Gets relative position buckets for full attention.

    Args:
      qlen: attention query length.
      klen: attention key length.
      bidirectional: a boolean - whether the attention is bidirectional

    Returns:
      int32 (qlen, klen) shaped array containing values in the range
      [0, num_buckets).
    """
    # TODO: should we be computing this w. numpy as a program
    # constant?
    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,
        bidirectional=bidirectional,
        num_buckets=self.num_buckets,
        max_distance=self.max_distance)
    return rp_bucket

  @nn.compact

# Path: flaxformer/types.py



# Path: flaxformer/architectures/longt5/relative_position_biases_general.py
from typing import Any, Callable
from flax import linen as nn
from flax.linen import partitioning
from jax import lax
from flaxformer.types import Array
import jax.numpy as jnp
import numpy as np

        (num_buckets - max_exact)).astype(jnp.int32)
    val_if_large = jnp.minimum(val_if_large, num_buckets - 1)
    ret += jnp.where(is_small, n, val_if_large)
    return ret

  def full_att_rp_bucket(self, qlen, klen, bidirectional=True):
    """Gets relative position buckets for full attention.

    Args:
      qlen: attention query length.
      klen: attention key length.
      bidirectional: a boolean - whether the attention is bidirectional

    Returns:
      int32 (qlen, klen) shaped array containing values in the range
      [0, num_buckets).
    """
    # TODO: should we be computing this w. numpy as a program
    # constant?
    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,
        bidirectional=bidirectional,
        num_buckets=self.num_buckets,
        max_distance=self.max_distance)
    return rp_bucket

  @nn.compact

  def __call__(self, rp_bucket: Array):