pyzag.reparametrization

Helper methods for reparameterizing modules, for example to scale parameter values and gradients

class pyzag.reparametrization.RangeRescale(lb: Tensor | float, ub: Tensor | float, clamp: bool = True)

Bases: Module

Scale parameter within bounds

forward(X: Tensor) → Tensor

Go from scaled to natural parameters

Parameters:

X (torch.tensor) – scaled parameter values

reverse(X: Tensor) → Tensor

Go from natural to scaled parameter values

Parameters:

X (torch.tensor) – natural parameter values

forward_std_dev(X: Tensor) → Tensor

Go from the standard deviation of a scaled normal to the actual standard deviation

Parameters:

X (torch.tensor) – scaled standard deviation

reverse_std_dev(X: Tensor) → Tensor

Go from the standard deviation of the actual normal to the standard deviation of the scaled normal

Parameters:

X (torch.tensor) – natural standard deviation

class pyzag.reparametrization.CurvatureRescale(scale, offset=0.0, lb=None, ub=None)

Bases: Module

Scale a parameter by a curvature-derived step size, with optional bounds.

The data-driven counterpart of RangeRescale. Where that one takes the step scale from a hand-picked range width (ub - lb), this takes it from the Gauss-Newton curvature, scale = 1 / sqrt(diag H) – build one per parameter with pyzag.preconditioning.gauss_newton_rescalers().

This separates two jobs that a range rescale conflates. A range width both bounds the parameter and sets its step size, and those pull in opposite directions: bounds you can trust are wide, and wide bounds condition the problem badly. Here the scale comes from the data and lb / ub are bounds and nothing else, so they can be as generous as honesty requires without costing anything.

Being a reparametrization rather than a gradient preconditioner, it works with any optimizer – including Adam, which is invariant to pyzag.preconditioning.GaussNewtonPreconditioner and rejects it. The optimizer’s own state (momentum, second moments) lives in the scaled coordinates along with the metric, which is what makes the pair coherent.

It is static: the scale is fixed when it is built. If the curvature drifts materially over the fit, use GaussNewtonPreconditioner, which can refresh; re-scaling mid-run would invalidate a stateful optimizer’s moments.

Parameters:

scale (torch.tensor) – per-element step scale, 1 / sqrt(diag H).

Keyword Arguments:

• offset (torch.tensor or float) – natural value at scaled zero (default 0).

• ub (lb,) – bounds in natural units, clamped after scaling. Both must be given, or neither. Note that swapping a RangeRescale for this one drops its clamp unless you pass these.

forward(X)

Go from scaled to natural parameters

Parameters:

X (torch.tensor) – scaled parameter values

reverse(X)

Go from natural to scaled parameter values

Parameters:

X (torch.tensor) – natural parameter values

forward_std_dev(X)

Go from the standard deviation of a scaled normal to the actual standard deviation

Parameters:

X (torch.tensor) – scaled standard deviation

reverse_std_dev(X)

Go from the standard deviation of the actual normal to the standard deviation of the scaled normal

Parameters:

X (torch.tensor) – natural standard deviation

class pyzag.reparametrization.Reparameterizer(map_dict: Mapping[str, RangeRescale], error_not_provided: bool = False)

Bases: object

Reparameterize a torch Module by adding the appropriate rescale function to each parameter

Parameters:

map_dict (dict mapping str to rescaler) – dictionary mapping the parameter name to the appropriate rescaler

Keyword Arguments:

error_not_provided (bool) – if True, error out if a rescaler is missing