No public description

PiperOrigin-RevId: 568424151
google-deepmind · Nov 13, 2023 · 71c33e2 · 71c33e2
1 parent 69c993c
commit 71c33e2
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/tracr/craft/transformers.py b/tracr/craft/transformers.py
@@ -75,6 +75,7 @@ def __post_init__(self):
     assert self.w_ov.output_space.issubspace(self.residual_space)
 
   def apply(self, x: bases.VectorInBasis) -> bases.VectorInBasis:
+    assert self.residual_space is not None
     assert x in self.residual_space
     # seq_len x query_space
     queries = x.project(self.w_qk.left_space)