Relax attend_to_offset value width

Author: physicsrob

tests/compile/forward/test_forward_compile.py

@@ -145,6 +145,23 @@ def test_compile_attend_to_offset():
     )
 
 
+def test_compile_attend_to_offset_wide_value():
+    """attend_to_offset supports values wider than the position encoding."""
+    pos = create_pos_encoding()
+    v = create_input("v", 20)
+    out = pos.attend_to_offset(v, delta_pos=-1)
+
+    n_pos = 5
+    _verify(
+        out,
+        n_pos=n_pos,
+        input_values={
+            "v": torch.randn(n_pos, 20),
+        },
+        pos_encoding=pos,
+    )
+
+
 def test_compile_map_to_table():
     """map_to_table — table lookup via MLP (large d_hidden)."""
     x = create_input("x", 2)

tests/graph/test_pos_encoding.py

@@ -32,7 +32,18 @@ def test_attend_to_offset():
     pos_encoding = PosEncoding(16)
     last_input = pos_encoding.attend_to_offset(value_input, delta_pos=-1)
     output = last_input.compute(n_pos=5, input_values={"value": input_values})
-    assert torch.allclose(output[:, 1:], input_values[:, :-1])
+    assert torch.allclose(output[0], input_values[0])
+    assert torch.allclose(output[1:], input_values[:-1])
+
+
+def test_attend_to_offset_value_can_be_wider_than_position_encoding():
+    input_values = torch.arange(5 * 20, dtype=torch.float32).reshape(5, 20)
+    value_input = InputNode("value", 20, value_range=(-100.0, 100.0))
+    pos_encoding = PosEncoding(16)
+    last_input = pos_encoding.attend_to_offset(value_input, delta_pos=-1)
+    output = last_input.compute(n_pos=5, input_values={"value": input_values})
+    assert torch.allclose(output[0], input_values[0])
+    assert torch.allclose(output[1:], input_values[:-1])
 
 
 def test_delta_matrix():

torchwright/graph/pos_encoding.py

@@ -61,30 +61,30 @@ def attend_to_offset(self, value: Node, delta_pos=-1) -> Node:
             # NOOP -- supporting this simplifies some use cases.
             return value
 
-        # We're applying the query/key to the position
-        d_head = self.d_pos
-        assert self.d_pos <= d_head
-        assert len(value) <= d_head
+        # Q/K match positions in positional-encoding space; V/O transports
+        # the payload independently and may be wider than the position vector.
+        d_qk = self.d_pos
+        d_v = len(value)
 
-        # key_matrix shape (d_key_in, d_head)
+        # key_matrix shape (d_key_in, d_qk)
         delta = get_pos_delta_matrix(delta_pos, self.d_pos)
         # The last sin/cos pair (columns d_pos-2, d_pos-1) contains the
         # raw position counter instead of a sinusoid.  Zero it out so the
         # counter doesn't participate in the trig-identity attention.
         delta[self.d_pos - 2 :, :] = 0.0
         delta[:, self.d_pos - 2 :] = 0.0
-        key_matrix = torch.zeros((len(self), d_head))
+        key_matrix = torch.zeros((len(self), d_qk))
         key_matrix[:, : self.d_pos] = delta
 
-        # query_matrix shape (d_query_in, d_head)
-        query_matrix = attention_hardness * torch.eye(len(self), d_head)
+        # query_matrix shape (d_query_in, d_qk)
+        query_matrix = attention_hardness * torch.eye(len(self), d_qk)
         query_matrix[self.d_pos - 2 :, :] = 0.0
 
-        # value_matrix shape (d_value_in, d_head)
-        value_matrix = torch.eye(len(value), d_head)
+        # value_matrix shape (d_value_in, d_v)
+        value_matrix = torch.eye(d_v)
 
-        # output_matrix shape (d_head, d_output)
-        output_matrix = torch.eye(d_head, len(value))
+        # output_matrix shape (d_v, d_output)
+        output_matrix = torch.eye(d_v)
 
         return Attn(
             query_in=self,