Module neuralops.core.model_utils

Generic helpers: calibration, weight‑copying, safe module replacement, activation collection.

Functions

def activation_hooks(model: torch.nn.modules.module.Module,
target_module_type: type = torch.nn.modules.linear.Linear,
collect_input: bool = True,
collect_output: bool = False) ‑> Any

Expand source code

@contextmanager
def activation_hooks(
    model: nn.Module,
    target_module_type: type = nn.Linear,
    collect_input: bool = True,
    collect_output: bool = False,
) -> Any:
    """Context manager that collects activations for all modules of a given type."""
    activations: dict[str, Any] = {}
    hooks = []

    def make_hook(name) -> Any:
        """Make hook."""

        def hook(module, input, output) -> None:
            """Hook."""
            data = {}
            if collect_input:
                data["input"] = input[0].detach()
            if collect_output:
                data["output"] = (
                    output.detach() if isinstance(output, torch.Tensor) else output[0].detach()
                )
            activations[name] = data

        return hook

    for name, module in model.named_modules():
        if isinstance(module, target_module_type):
            hooks.append(module.register_forward_hook(make_hook(name)))

    try:
        yield activations
    finally:
        for hook in hooks:
            hook.remove()

Context manager that collects activations for all modules of a given type.

def copy_attention_weights(src_module: torch.nn.modules.module.Module,
dst_module: torch.nn.modules.module.Module,
source_map: dict[str, list[str]] | None = None) ‑> None

Expand source code

def copy_attention_weights(
    src_module: nn.Module,
    dst_module: nn.Module,
    source_map: dict[str, list[str]] | None = None,
) -> None:
    """
    Copy pretrained attention weights from src to dst.
    Tries multiple possible attribute names for each projection.
    """
    if source_map is None:
        source_map = ATTENTION_WEIGHT_MAP
    for target_attr, src_attrs in source_map.items():
        src = None
        for attr in src_attrs:
            if hasattr(src_module, attr):
                src = getattr(src_module, attr)
                break
        if src is None:
            raise ValueError(
                f"Cannot find source for '{target_attr}' in "
                f"{src_module.__class__.__name__}. Tried: {src_attrs}"
            )
        dst = getattr(dst_module, target_attr)
        dst.weight.data.copy_(src.weight.data)
        if src.bias is not None and dst.bias is not None:
            dst.bias.data.copy_(src.bias.data)
        elif src.bias is None and dst.bias is not None:
            nn.init.zeros_(dst.bias)

Copy pretrained attention weights from src to dst. Tries multiple possible attribute names for each projection.

def count_parameters(model: torch.nn.modules.module.Module, trainable_only: bool = False) ‑> int

Expand source code

def count_parameters(model: nn.Module, trainable_only: bool = False) -> int:
    """Count parameters."""
    if trainable_only:
        return sum(p.numel() for p in model.parameters() if p.requires_grad)
    return sum(p.numel() for p in model.parameters())

Count parameters.

def count_params(model: torch.nn.modules.module.Module, trainable_only: bool = False) ‑> int

Expand source code

def count_parameters(model: nn.Module, trainable_only: bool = False) -> int:
    """Count parameters."""
    if trainable_only:
        return sum(p.numel() for p in model.parameters() if p.requires_grad)
    return sum(p.numel() for p in model.parameters())

Count parameters.

def get_parent_and_child_name(model: torch.nn.modules.module.Module, full_name: str) ‑> Any

Expand source code

def get_parent_and_child_name(model: nn.Module, full_name: str) -> Any:
    """Return (parent_module, child_attribute_name), supporting ModuleList and ModuleDict."""
    parts = [p for p in full_name.split(".") if p]
    parent = model
    for part in parts[:-1]:
        if isinstance(parent, (nn.ModuleList, list)):
            parent = parent[int(part)]
        elif isinstance(parent, (nn.ModuleDict, dict)):
            parent = parent[part]
        else:
            parent = getattr(parent, part)
    return parent, parts[-1]

Return (parent_module, child_attribute_name), supporting ModuleList and ModuleDict.

def replace_module(model: torch.nn.modules.module.Module,
full_name: str,
new_module: torch.nn.modules.module.Module) ‑> None

Expand source code

def replace_module(model: nn.Module, full_name: str, new_module: nn.Module) -> None:
    """Replace a sub-module identified by its dotted name, supporting ModuleList and ModuleDict."""
    parts = [p for p in full_name.split(".") if p]
    parent = model
    for part in parts[:-1]:
        if isinstance(parent, (nn.ModuleList, list)):
            parent = parent[int(part)]
        elif isinstance(parent, (nn.ModuleDict, dict)):
            parent = parent[part]
        else:
            parent = getattr(parent, part)
    last = parts[-1]
    if isinstance(parent, (nn.ModuleList, list)):
        parent[int(last)] = new_module
    elif isinstance(parent, (nn.ModuleDict, dict)):
        parent[last] = new_module
    else:
        setattr(parent, last, new_module)

Replace a sub-module identified by its dotted name, supporting ModuleList and ModuleDict.

def run_calibration_forward(model: torch.nn.modules.module.Module,
tokenizer,
texts: list[str],
device: torch.device | None = None,
max_length: int = 128) ‑> torch.nn.modules.module.Module

Expand source code

def run_calibration_forward(
    model: nn.Module,
    tokenizer,
    texts: list[str],
    device: torch.device | None = None,
    max_length: int = 128,
) -> nn.Module:
    """Run a single forward pass on calibration texts; return model (in eval)."""
    model.eval()
    inputs = tokenizer(
        texts, return_tensors="pt", padding=True, truncation=True, max_length=max_length
    )
    if device is None:
        device = next(model.parameters()).device
    inputs = {k: v.to(device) for k, v in inputs.items()}
    with torch.no_grad():
        model(**inputs)
    return model

Run a single forward pass on calibration texts; return model (in eval).