self_attention_fillin.py

class SelfAttention(torch.nn.Module):
  def __init__(self, input_dimension, output_dimension):
    super().__init__()

    self.input_dimension = input_dimension
    self.output_dimension = output_dimension

    self.Qw = torch.nn.Linear(self.input_dimension, self.output_dimension)
    self.Kw = torch.nn.Linear(self.input_dimension, self.output_dimension)
    self.Vw = torch.nn.Linear(self.input_dimension, self.output_dimension)

  def forward(self, x):
    Q = # Fill in
    K = # Fill in
    V = # Fill in

    raw_attention_scores = # Fill in
    normalized_attention_scores = # Fill in
    attended_inputs = # Fill in

    return attended_inputs