Generic function generator to replace base_t module with dest_t wrapper. Args: base_t : module type to replace dest_t : destination module type Returns: swap function to replace base_t module with dest_t
(base_t: type[nn.Module], dest_t: type[nn.Module])
| 1312 | |
| 1313 | |
| 1314 | def wrap_module(base_t: type[nn.Module], dest_t: type[nn.Module]) -> Callable[[nn.Module], nn.Module | None]: |
| 1315 | """ |
| 1316 | Generic function generator to replace base_t module with dest_t wrapper. |
| 1317 | Args: |
| 1318 | base_t : module type to replace |
| 1319 | dest_t : destination module type |
| 1320 | Returns: |
| 1321 | swap function to replace base_t module with dest_t |
| 1322 | """ |
| 1323 | |
| 1324 | def expansion_fn(mod: nn.Module) -> nn.Module | None: |
| 1325 | out = dest_t(mod) |
| 1326 | return out |
| 1327 | |
| 1328 | return expansion_fn |
| 1329 | |
| 1330 | |
| 1331 | def simple_replace(base_t: type[nn.Module], dest_t: type[nn.Module]) -> Callable[[nn.Module], nn.Module | None]: |
no outgoing calls
no test coverage detected
searching dependent graphs…