Module neuralops.utils.helpers

Module providing helpers implementations for NeuralOps.

Functions

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

Expand source code

def count_parameters(model: torch.nn.Module, trainable_only: bool = True) -> 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 ensure_dir(path: pathlib.Path | str) ‑> pathlib.Path

Expand source code

def ensure_dir(path: str | Path) -> Path:
    """Ensure dir."""
    path = Path(path)
    path.mkdir(parents=True, exist_ok=True)
    return path

Ensure dir.

def flatten_dict(d: dict, parent_key: str = '', sep: str = '.') ‑> dict

Expand source code

def flatten_dict(d: dict, parent_key: str = "", sep: str = ".") -> dict:
    """Flatten dict."""
    items = []
    for k, v in d.items():
        new_key = f"{parent_key}{sep}{k}" if parent_key else k
        if isinstance(v, dict):
            items.extend(flatten_dict(v, new_key, sep=sep).items())
        else:
            items.append((new_key, v))
    return dict(items)

Flatten dict.

def format_number(num: int) ‑> str

Expand source code

def format_number(num: int) -> str:
    """Format number."""
    if abs(num) >= 1_000_000_000:
        return f"{num / 1_000_000_000:.1f}B"
    if abs(num) >= 1_000_000:
        return f"{num / 1_000_000:.1f}M"
    if abs(num) >= 1_000:
        return f"{num / 1_000:.0f}K"
    return str(num)

Format number.

def format_time(seconds: float) ‑> str

Expand source code

def format_time(seconds: float) -> str:
    """Format time."""
    hours = int(seconds // 3600)
    minutes = int((seconds % 3600) // 60)
    secs = int(seconds % 60)
    parts = []
    if hours:
        parts.append(f"{hours}h")
    if minutes:
        parts.append(f"{minutes}m")
    parts.append(f"{secs}s")
    return " ".join(parts)

Format time.

def load_json(path: pathlib.Path | str) ‑> Any

Expand source code

def load_json(path: str | Path) -> Any:
    """Load json from disk or config."""
    with open(path) as f:
        return json.load(f)

Load json from disk or config.

def save_json(data: Any, path: pathlib.Path | str, indent: int = 2) ‑> None

Expand source code

def save_json(data: Any, path: str | Path, indent: int = 2) -> None:
    """Persist json to disk."""
    ensure_dir(os.path.dirname(path))
    with open(path, "w") as f:
        json.dump(data, f, indent=indent)

Persist json to disk.

def set_seed(seed: int, deterministic: bool = False) ‑> None

Expand source code

def set_seed(seed: int, deterministic: bool = False) -> None:
    """Set the seed."""
    set_all_seeds(seed)
    if deterministic and torch.cuda.is_available():
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False

Set the seed.

def unflatten_dict(d: dict, sep: str = '.') ‑> dict

Expand source code

def unflatten_dict(d: dict, sep: str = ".") -> dict:
    """Unflatten dict."""
    result = {}
    for key, value in d.items():
        parts = key.split(sep)
        current = result
        for part in parts[:-1]:
            if part not in current:
                current[part] = {}
            current = current[part]
        current[parts[-1]] = value
    return result

Unflatten dict.