def __call__(self, hidden_states:Tensor, position_bias:Optional[Tensor]=None) -> Tuple[Tensor,Tensor]:
    """
    Self-attention (if key_value_states is None) or attention over source sentence (provided by key_value_states).
    """
    # Input is (batch_size, seq_length, dim)
    batch_size, key_length = hidden_states.shape[:2]

    def shape(states):
      """projection"""
      return states.view(batch_size, -1, self.n_heads, self.key_value_proj_dim).transpose(1, 2)

    def unshape(states):
      """reshape"""
      return states.transpose(1, 2).contiguous().view(batch_size, -1, self.inner_dim)

    def project(hidden_states, proj_layer):
      """projects hidden states correctly to key/query states"""
      # self-attn
      # (batch_size, n_heads, seq_length, dim_per_head)
      return shape(proj_layer(hidden_states))

    # get query states
    query_states = shape(self.q(hidden_states))  # (batch_size, n_heads, seq_length, dim_per_head)

    # get key/value states
    key_states = project(hidden_states, self.k)
    value_states = project(hidden_states, self.v)

    # compute scores
    scores = Tensor.matmul(query_states, key_states.transpose(3, 2))

    if position_bias is None:
      position_bias = self.compute_bias(key_length, key_length, device=scores.device)

    scores += position_bias
    attn_weights = Tensor.softmax(scores.float(), axis=-1).cast(scores.dtype)  # (batch_size, n_heads, seq_length, key_length)

    attn_output = unshape(Tensor.matmul(attn_weights, value_states))  # (batch_size, seq_length, dim)
    attn_output = self.o(attn_output)

    return attn_output, position_bias

    Self-attention (if key_value_states is None) or attention over source sentence (provided by key_value_states).
    
extra/models/t5.py