| 1 | from typing import Callable, NamedTuple |
| 2 | |
| 3 | import torch |
| 4 | |
| 5 | |
| 6 | class ModuleOps(NamedTuple): |
| 7 | """ |
| 8 | Defines a named operation for matching and mutating PyTorch modules. |
| 9 | Used to selectively transform modules in a model (e.g., replacing layers with quantized versions). |
| 10 | """ |
| 11 | |
| 12 | name: str |
| 13 | matcher: Callable[[torch.nn.Module], bool] |
| 14 | mutator: Callable[[torch.nn.Module], torch.nn.Module] |
| 15 |