Update pep8 (#78)

* simplify requirements file

* style

* apply linter

ch03/02_bonus_efficient-multihead-attention/ch03.py

@@ -8,33 +8,33 @@ class CausalAttention(nn.Module):
         super().__init__()
         self.d_out = d_out
         self.W_query = nn.Linear(d_in, d_out, bias=qkv_bias)
-        self.W_key   = nn.Linear(d_in, d_out, bias=qkv_bias)
+        self.W_key = nn.Linear(d_in, d_out, bias=qkv_bias)
         self.W_value = nn.Linear(d_in, d_out, bias=qkv_bias)
-        self.dropout = nn.Dropout(dropout) # New
-        self.register_buffer('mask', torch.triu(torch.ones(block_size, block_size), diagonal=1)) # New
+        self.dropout = nn.Dropout(dropout)  # New
+        self.register_buffer('mask', torch.triu(torch.ones(block_size, block_size), diagonal=1))  # New
 
     def forward(self, x):
-        b, num_tokens, d_in = x.shape # New batch dimension b
+        b, num_tokens, d_in = x.shape  # New batch dimension b
         keys = self.W_key(x)
         queries = self.W_query(x)
         values = self.W_value(x)
 
-        attn_scores = queries @ keys.transpose(1, 2) # Changed transpose
+        attn_scores = queries @ keys.transpose(1, 2)  # Changed transpose
         attn_scores.masked_fill_(  # New, _ ops are in-place
-            self.mask.bool()[:num_tokens, :num_tokens], -torch.inf) 
+            self.mask.bool()[:num_tokens, :num_tokens], -torch.inf)
         attn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1)
-        attn_weights = self.dropout(attn_weights) # New
+        attn_weights = self.dropout(attn_weights)  # New
 
         context_vec = attn_weights @ values
         return context_vec
-    
-    
+
+
 class MultiHeadAttentionWrapper(nn.Module):
 
     def __init__(self, d_in, d_out, block_size, dropout, num_heads, qkv_bias=False):
         super().__init__()
         self.heads = nn.ModuleList(
-            [CausalAttention(d_in, d_out, block_size, dropout, qkv_bias) 
+            [CausalAttention(d_in, d_out, block_size, dropout, qkv_bias)
              for _ in range(num_heads)]
         )
         self.out_proj = nn.Linear(d_out*num_heads, d_out*num_heads)
@@ -44,7 +44,6 @@ class MultiHeadAttentionWrapper(nn.Module):
         return self.out_proj(context_vec)
 
 
-
 class MultiHeadAttention(nn.Module):
     def __init__(self, d_in, d_out, block_size, dropout, num_heads, qkv_bias=False):
         super().__init__()
@@ -52,7 +51,7 @@ class MultiHeadAttention(nn.Module):
 
         self.d_out = d_out
         self.num_heads = num_heads
-        self.head_dim = d_out // num_heads # Reduce the projection dim to match desired output dim
+        self.head_dim = d_out // num_heads  # Reduce the projection dim to match desired output dim
 
         self.W_query = nn.Linear(d_in, d_out, bias=qkv_bias)
         self.W_key = nn.Linear(d_in, d_out, bias=qkv_bias)
@@ -70,7 +69,7 @@ class MultiHeadAttention(nn.Module):
 
         # We implicitly split the matrix by adding a `num_heads` dimension
         # Unroll last dim: (b, num_tokens, d_out) -> (b, num_tokens, num_heads, head_dim)
-        keys = keys.view(b, num_tokens, self.num_heads, self.head_dim) 
+        keys = keys.view(b, num_tokens, self.num_heads, self.head_dim)
         values = values.view(b, num_tokens, self.num_heads, self.head_dim)
         queries = queries.view(b, num_tokens, self.num_heads, self.head_dim)
 
@@ -92,10 +91,10 @@ class MultiHeadAttention(nn.Module):
         attn_weights = self.dropout(attn_weights)
 
         # Shape: (b, num_tokens, num_heads, head_dim)
-        context_vec = (attn_weights @ values).transpose(1, 2) 
+        context_vec = (attn_weights @ values).transpose(1, 2)
 
         # Combine heads, where self.d_out = self.num_heads * self.head_dim
         context_vec = context_vec.contiguous().view(b, num_tokens, self.d_out)
-        context_vec = self.out_proj(context_vec) # optional projection
+        context_vec = self.out_proj(context_vec)  # optional projection
 
-        return context_vec
+        return context_vec