Test that state resets correctly after num_inference_steps.
(self)
| 193 | ) |
| 194 | |
| 195 | def test_mag_cache_reset(self): |
| 196 | """Test that state resets correctly after num_inference_steps.""" |
| 197 | model = DummyTransformer() |
| 198 | config = MagCacheConfig( |
| 199 | threshold=100.0, num_inference_steps=2, retention_ratio=0.0, mag_ratios=np.array([1.0, 1.0]) |
| 200 | ) |
| 201 | apply_mag_cache(model, config) |
| 202 | self._set_context(model, "test_context") |
| 203 | |
| 204 | input_t = torch.ones(1, 1, 1) |
| 205 | |
| 206 | model(input_t) # Step 0 |
| 207 | model(input_t) # Step 1 (Skipped) |
| 208 | |
| 209 | # Step 2 (Reset -> Step 0) -> Should Compute |
| 210 | # Input 2.0 -> Output 8.0 |
| 211 | input_t2 = torch.tensor([[[2.0]]]) |
| 212 | output_t2 = model(input_t2) |
| 213 | |
| 214 | self.assertTrue(torch.allclose(output_t2, torch.tensor([[[8.0]]])), "State did not reset correctly") |
| 215 | |
| 216 | def test_mag_cache_calibration(self): |
| 217 | """Test that calibration mode records ratios.""" |
nothing calls this directly
no test coverage detected