src/transformers/models/longt5/modeling_flax_longt5.py/FlaxLongT5LocalAttention/compute_bias
class FlaxLongT5LocalAttention:    def compute_bias(self, block_length: int):
        """Compute binned relative position bias"""
        memory_position = jnp.arange(3 * block_length, dtype="i4")
        context_position = memory_position[block_length:-block_length]

        relative_position = memory_position[None, :] - context_position[:, None]
        relative_position_bucket = self._relative_position_bucket(
            relative_position,
            bidirectional=True,
            num_buckets=self.relative_attention_num_buckets,
            max_distance=self.relative_attention_max_distance,
        )

        values = self.relative_attention_bias(relative_position_bucket)
        values = values.transpose((2, 0, 1))[None, None, :, :, :]
        return values