remove hidden state rollback indicators

Author: jakegrigsby

amago/nets/traj_encoders.py

@@ -461,12 +461,6 @@ def make_cache():
             val_cache=make_cache(),
             seq_lens=torch.zeros((batch_size,), dtype=torch.int32, device=device),
         )
-        # FIXME: remove after kv-cache validation
-        print(
-            f"[TformerTrajEncoder] init_hidden_state max_seq_len={self.max_seq_len} "
-            f"batch_size={batch_size} n_layers={self.n_layers} "
-            f"n_heads={self.n_heads} head_dim={self.head_dim}"
-        )
         return hidden_state
 
     def reset_hidden_state(

amago/nets/transformer.py

@@ -604,13 +604,6 @@ def __len__(self):
 
     def roll_back(self, seq_lens):
         idxs = torch.where(seq_lens == self.max_seq_len)[0]
-        # FIXME: remove after kv-cache validation
-        if idxs.numel() > 0 and not getattr(self, "_logged_first_roll", False):
-            print(
-                f"[Cache] FIRST roll_back triggered: max_seq_len={self.max_seq_len} "
-                f"n_actors_overflowing={idxs.numel()}"
-            )
-            self._logged_first_roll = True
         roll = self.data[:, idxs, 1:].clone()
         self.data[:, idxs, :-1] = roll
         self.data[:, idxs, -1] = torch.nan  # no silent bugs
@@ -631,13 +624,6 @@ def __init__(self, key_cache: Cache, val_cache: Cache, seq_lens: torch.Tensor):
         self.device = key_cache.device
 
     def reset(self, idxs):
-        # FIXME: remove after kv-cache validation
-        n_reset = int(idxs.sum()) if hasattr(idxs, "sum") else len(idxs)
-        if n_reset > 0:
-            print(
-                f"[TformerHiddenState] reset n_actors_reset={n_reset} "
-                f"max_seq_len={self.key_cache.max_seq_len}"
-            )
         self.seq_lens[idxs] = 0
 
     def update(self):