better naming

src/transformers/configuration_reformer.py

@@ -100,8 +100,8 @@ def __init__(
         vocab_size=10,
         attention_head_size=32,
         hidden_size=64,
-        num_attention_heads=2,
-        num_buckets=2,
+        num_attention_heads=1,
+        num_buckets=[2, 4],
         num_hashes=2,
         lsh_attn_chunk_length=64,
         local_attn_chunk_length=64,
@@ -122,7 +122,7 @@ def __init__(
         layer_norm_eps=1e-12,
         sinusoidal_pos_embds=False,
         axial_pos_embds=False,
-        axial_pos_shape=[8, 8],
+        axial_pos_shape=[32, 16],
         axial_pos_embds_dim=[32, 32],
         attn_layers=["lsh", "lsh", "lsh", "lsh"],
 #        attn_layers=["local", "local", "local", "local"],

src/transformers/modeling_reformer.py

@@ -471,6 +471,7 @@ def _hash_vectors(self, vectors, num_hashes):
         # See https://arxiv.org/pdf/1509.02897.pdf
         # We sample a different random rotation for each round of hashing to
         # decrease the probability of hash misses.
+
         if isinstance(self.num_buckets, int):
             assert (
                 self.num_buckets % 2 == 0
@@ -480,12 +481,13 @@ def _hash_vectors(self, vectors, num_hashes):
         else:
             # Factorize the hash if self.num_buckets is a list or tuple
             rotation_size, num_buckets = 0, 1
-            for num_bucket in self.num_buckets:
-                assert num_bucket % 2 == 0, "The number of buckets should be even, but `num_bucket`: {}".format(
-                    num_bucket
-                )
-                rotation_size += num_bucket
-                num_buckets *= num_bucket
+            for bucket_factor in self.num_buckets:
+                assert bucket_factor % 2 == 0, "The number of buckets should be even, but `num_bucket`: {}".format(bucket_factor)
+                rotation_size = rotation_size + bucket_factor
+                num_buckets = num_buckets * bucket_factor
+
+        # remove gradient
+        vectors = vectors.detach()
 
         # TODO: delete later when integration tests are ok
         if self.hash_seed is not None:
@@ -497,7 +499,7 @@ def _hash_vectors(self, vectors, num_hashes):
             rotated_vectors = torch.einsum("bmtd,dhr->bmhtr", vectors, random_rotations)
         else:
             rotations_shape = (self.num_attention_heads, vectors.shape[-1], num_hashes, rotation_size // 2)
-            # create a random self.attention_head_size x num_hashes x self.num_buckets/2
+            # create a random self.attention_head_size x num_hashes x num_buckets/2
             random_rotations = torch.randn(rotations_shape, device=vectors.device).to(vectors.dtype)
 
             # rotated_vectors has dim:
@@ -513,17 +515,17 @@ def _hash_vectors(self, vectors, num_hashes):
         else:
             # Get the buckets for them and combine.
             buckets, cur_sum, cur_product = None, 0, 1
-            for num_bucket in self.num_buckets:
-                rotated_vectors = rotated_vectors[..., cur_sum : cur_sum + (num_bucket // 2)]
-                cur_sum += num_bucket // 2
-                rotated_vectors = torch.cat([rotated_vectors, -rotated_vectors], dim=-1)
+            for bucket_factor in self.num_buckets:
+                rotated_vectors_factor = rotated_vectors[..., cur_sum : cur_sum + (bucket_factor // 2)]
+                cur_sum = cur_sum + bucket_factor // 2
+                rotated_vectors_factor = torch.cat([rotated_vectors_factor, -rotated_vectors_factor], dim=-1)
 
                 if buckets is None:
-                    buckets = torch.argmax(rotated_vectors, dim=-1)
+                    buckets = torch.argmax(rotated_vectors_factor, dim=-1)
                 else:
-                    buckets += cur_product * torch.argmax(rotated_vectors, dim=-1)
+                    buckets = buckets + (cur_product * torch.argmax(rotated_vectors_factor, dim=-1))
 
-                cur_product *= num_bucket
+                cur_product = cur_product * bucket_factor
 
         # buckets is now (Batch_size x Num_Attn_Heads x Num_Hashes x Seq_Len).
         # Next we add offsets so that bucket numbers from different hashing rounds don't overlap.
@@ -1511,8 +1513,8 @@ def forward(
         if labels is not None:
             # Shift so that tokens < n predict n
             # Uncomment this line for integration test with Trax
-#            shift_logits = logits.contiguous()
-            shift_logits = logits[..., :-1, :].contiguous()
+            shift_logits = logits.contiguous()
+#            shift_logits = logits[..., :-1, :].contiguous()
 
             shift_labels = labels[..., 1:].contiguous()
             # Flatten the tokens

tests/test_modeling_reformer.py

@@ -631,7 +631,7 @@ def test_local_layer(self):
     def test_reformer_lm_model(self):
         config = ReformerConfig(axial_pos_embds=True, hash_seed=0, is_decoder=True)
 
-        shape = (1, 64)  # Batch x SeqLen x ModelDimPerHead
+        shape = (1, 512)  # Batch x SeqLen x ModelDimPerHead
 
         np_input = np.random.randint(0, config.vocab_size, size=shape)
         np_input_2 = np.asarray(np_input, np.float32)