torch_vn_algebra.cost¶

Source: torch_vn_algebra/cost.py. Signatures and docstrings as in the source.

Module documentation¶

Safeguards against unexpectedly long or memory-hungry computations.

* Time: batch-separable kernels (eigh, SVD, QR, matrix exponentials) are routed through
  `batched_call`. For large problems it first runs the kernel on a small slice of the batch,
  extrapolates the time for the whole batch and issues a CostWarning with the estimate when
  it exceeds `warn_seconds`. The result is identical to a single call.
* Memory: `check_memory` compares the estimated size of a large allocation with the memory
  available on the device; above `warn_memory_fraction` it warns, above `max_memory_fraction`
  it raises InsufficientMemoryError (instead of an out-of-memory crash halfway through).
* Integrators (dynamics.rk4 and everything built on it) warn with an ETA after the first step
  and can show a progress bar.

    import torch_vn_algebra as tv
    tv.cost.set_limits(warn_seconds=30)          # thresholds
    with tv.cost.disabled():                     # switch the checks off locally
        ...
    import warnings; warnings.simplefilter('ignore', tv.cost.CostWarning)   # silence warnings only

CostWarning¶

class CostWarning()

A computation is expected to take long or to use a large share of the memory.

InsufficientMemoryError¶

class InsufficientMemoryError()

The estimated memory of an operation exceeds what is available on the device.

set_limits¶

def set_limits(**kwargs)

Change thresholds: warn_seconds, warn_memory_fraction, max_memory_fraction, probe_items,
probe_flops, enabled.

get_limits¶

def get_limits() -> dict

disabled¶

def disabled()

Temporarily switch all cost checks off.

available_memory¶

def available_memory(device) -> Optional[int]

Free memory on `device` in bytes (CUDA: mem_get_info; CPU: MemAvailable on Linux, psutil otherwise).

check_memory¶

def check_memory(nbytes: float, device, what: str)

Warn or raise if an allocation of about `nbytes` is large compared with the free memory.

itemsize¶

def itemsize(dtype: torch.dtype) -> int

Bytes per entry of `dtype` (torch.dtype.itemsize exists only from PyTorch 2.1).

abs_max¶

def abs_max(x: torch.Tensor, dim=None) -> torch.Tensor

max |x_i| over `dim` (all axes by default). On PyTorch versions whose complex |z| overflows for large
entries, max(|Re|, |Im|) instead (within a factor sqrt(2) of max |x_i|, never overflows).

tensor_bytes¶

def tensor_bytes(shape, dtype) -> int

kernel_flops¶

def kernel_flops(kind: str, k: int, complex_: bool) -> float

Rough flop count of one k x k kernel (only used to skip probing of small problems).

batched_call¶

def batched_call(fn: Callable, x: torch.Tensor, what: str, kind: str='eigh')

fn(x) for a function acting independently on the items x[0], x[1], ... (leading dimension). The first large
problem of a kind (same fn, block shape, dtype and device) is timed on slices of the batch (from 8 items on a
CPU, 1/16 of the batch on a GPU, where small slices do not fill the device) and the cost per item is remembered;
a CostWarning reports the expected time if it exceeds warn_seconds. A call has a fixed cost (dispatch, workspace;
milliseconds for batched LAPACK kernels) that is not proportional to the number of items: it is measured on one
item and subtracted, and the slice is enlarged (x4) until the per-item part dominates or the time it implies
for the whole batch is below warn_seconds. Later calls cost nothing extra. The result is fn(x), computed in one
call.

StepTimer¶

class StepTimer(total_steps: int, what: str, progress: bool=False)

ETA warning (and optional progress bar) for loops with a known number of equal steps.

StepTimer.step¶

def step(self, device=None)

StepTimer.close¶

def close(self)