src/transformers/models/big_bird/modeling_flax_big_bird.py/FlaxBigBirdAttention/__call__
class FlaxBigBirdAttention:    def __call__(
        self,
        hidden_states,
        attention_mask,
        layer_head_mask,
        key_value_states=None,
        init_cache=False,
        deterministic=True,
        output_attentions: bool = False,
    ):
        # Attention mask comes in as attention_mask.shape == (*batch_sizes, kv_length)
        # FLAX expects: attention_mask.shape == (*batch_sizes, 1, 1, kv_length) such that it is broadcastable
        # with attn_weights.shape == (*batch_sizes, num_heads, q_length, kv_length)
        if self.config.attention_type == "original_full":
            attn_outputs = self.self(
                hidden_states,
                attention_mask,
                layer_head_mask=layer_head_mask,
                key_value_states=key_value_states,
                init_cache=init_cache,
                deterministic=deterministic,
                output_attentions=output_attentions,
            )
        else:
            attn_outputs = self.self(
                hidden_states,
                attention_mask,
                deterministic=deterministic,
                output_attentions=output_attentions,
            )
        attn_output = attn_outputs[0]
        hidden_states = self.output(attn_output, hidden_states, deterministic=deterministic)

        outputs = (hidden_states,)

        if output_attentions:
            outputs += (attn_outputs[1],)

        return outputs