(
fn: Callable,
module: nn.Module,
name: str = "",
depth_first: bool = True,
include_root: bool = False,
)
| 26 | |
| 27 | |
| 28 | def named_replace( |
| 29 | fn: Callable, |
| 30 | module: nn.Module, |
| 31 | name: str = "", |
| 32 | depth_first: bool = True, |
| 33 | include_root: bool = False, |
| 34 | ) -> nn.Module: |
| 35 | if not depth_first and include_root: |
| 36 | module = fn(module=module, name=name) |
| 37 | for child_name_o, child_module in list(module.named_children()): |
| 38 | child_name = ".".join((name, child_name_o)) if name else child_name_o |
| 39 | new_child = named_replace( |
| 40 | fn=fn, |
| 41 | module=child_module, |
| 42 | name=child_name, |
| 43 | depth_first=depth_first, |
| 44 | include_root=True, |
| 45 | ) |
| 46 | setattr(module, child_name_o, new_child) |
| 47 | |
| 48 | if depth_first and include_root: |
| 49 | module = fn(module=module, name=name) |
| 50 | return module |
| 51 | |
| 52 | |
| 53 | def named_apply( |
nothing calls this directly
no outgoing calls
no test coverage detected