API — `nefi.utils`¶

Generated by scripts/gen_api_docs.py — do not edit by hand. Index

`nefi.utils`¶

Small utilities shared across nefi.

`nefi.utils.compat`¶

Version shims for the PyTorch APIs used by the performance switches (torch >= 2.1).

`def autocast_enabled(device_type: str) -> bool`¶

Whether torch.autocast is active for device_type.

`def cudagraph_mark_step_begin()`¶

Mark a new iteration for CUDA-graph trees (no-op when unavailable).

`def grad_scaler(device_type: str)`¶

Dynamic loss scaler for fp16 autocast on device_type.

`def is_functorch_wrapped(t: torch.Tensor) -> bool`¶

Whether t is a vmap / grad wrapper (batched) tensor.

`nefi.utils.device`¶

Device resolution. CUDA-first; MPS only when explicitly requested (no float64, partial FFT).

`def peak_memory_mb(device: torch.device) -> float | None`¶

Peak allocated memory in MB for CUDA devices, None otherwise.

`def reset_peak_memory(device: torch.device)`¶

`def resolve_device(device: str | torch.device | None = 'auto') -> torch.device`¶

Resolve a device spec.

`def synchronize(device: torch.device)`¶

Block until all queued kernels on device finished (no-op on CPU).

`nefi.utils.seed`¶

Reproducibility helpers.

`def seed_everything(seed: int | None, deterministic: bool = False)`¶

Seed python, numpy and torch (all devices).

`nefi.utils.tensor`¶

Tensor helpers: resampling across resolutions, cropping, FFT sizes.

`def as_tensor(x, device=None, dtype=None) -> torch.Tensor`¶

`def center_crop(x: torch.Tensor, shape: Sequence[int], offsets: Sequence[int] | None = None)`¶

Crop the trailing dims of x to shape starting at offsets (default: centered).

`def next_fast_len(n: int) -> int`¶

Smallest FFT-friendly length >= n (real FFT).

`def pad_trailing(x: torch.Tensor, shape: Sequence[int], value: float = 0.0) -> torch.Tensor`¶

Zero-pad the trailing dims of x (at the end of each axis) up to shape.

`def resample(x: torch.Tensor, shape: Sequence[int], mode: str = 'auto') -> torch.Tensor`¶

Resample the trailing len(shape) dims of x to shape (differentiable).

`def shape_tuple(shape: Sequence[int] | int) -> tuple[int, ...]`¶

`def to_numpy(x: torch.Tensor) -> np.ndarray`¶

`nefi.utils.timing`¶

Wall-clock timing helper.

`class Timer()`¶

Context manager measuring wall-clock seconds; Timer().elapsed after the block.

API — nefi.utils¶

nefi.utils¶

nefi.utils.compat¶

def autocast_enabled(device_type: str) -> bool¶

def cudagraph_mark_step_begin()¶

def grad_scaler(device_type: str)¶

def is_functorch_wrapped(t: torch.Tensor) -> bool¶

nefi.utils.device¶

def peak_memory_mb(device: torch.device) -> float | None¶

def reset_peak_memory(device: torch.device)¶

def resolve_device(device: str | torch.device | None = 'auto') -> torch.device¶

def synchronize(device: torch.device)¶

nefi.utils.seed¶

def seed_everything(seed: int | None, deterministic: bool = False)¶

nefi.utils.tensor¶

def as_tensor(x, device=None, dtype=None) -> torch.Tensor¶

def center_crop(x: torch.Tensor, shape: Sequence[int], offsets: Sequence[int] | None = None)¶

def next_fast_len(n: int) -> int¶

def pad_trailing(x: torch.Tensor, shape: Sequence[int], value: float = 0.0) -> torch.Tensor¶

def resample(x: torch.Tensor, shape: Sequence[int], mode: str = 'auto') -> torch.Tensor¶

def shape_tuple(shape: Sequence[int] | int) -> tuple[int, ...]¶

def to_numpy(x: torch.Tensor) -> np.ndarray¶

nefi.utils.timing¶

class Timer()¶