omnibias.jax.optim

0.5.0rc1Release candidatePython >=3.10Apache-2.0

JAX activation derivatives, functional network jets, direct contractions, and parameter-differentiable numerical models.

import omnibias.jax.optim

Module contract & mathematical conventions

Optimization primitives for neural fields and residual least squares.

Includes safeguarded Newton, Gauss–Newton, Krylov, and trust-region methods. Model parameter derivatives may use autodiff; coordinate derivative jets use the activation tower.

Shared formulas do not imply universal floating-point identity. Verify dtype, shape, derivative-order support, and the stated numerical assumptions for the path you use.

Module source ↗Source revision 9a2d73df0573Running reference examples →

Imports & aliases

These names share a definition with another public module. Follow a name for its full contract, parameters, and examples.

Import from this module	Canonical reference
`BlockSpec`	`omnibias.core.BlockSpec`
`JetLineSearchConfig`	`omnibias.core.JetLineSearchConfig`
`LineSearchResult`	`omnibias.core.LineSearchResult`
`arrangement_w_block`	`omnibias.core.arrangement_w_block`
`block_direction`	`omnibias.jax.block_direction`
`block_exact_search`	`omnibias.jax.block_exact_search`
`block_exact_sweep`	`omnibias.jax.block_exact_sweep`
`jet_line_search`	`omnibias.jax.jet_line_search`
`jet_line_search_on_ray`	`omnibias.jax.jet_line_search_on_ray`
`last_linear_block`	`omnibias.core.last_linear_block`
`ombu_bias_block`	`omnibias.core.ombu_bias_block`

CLASS

CubicRegularizedGNConfig

omnibias.jax.optim.CubicRegularizedGNConfigSource ↗

class CubicRegularizedGNConfig

Dataclass constructor fields; inherited fields, if any, are declared by the base class. A field(...) expression describes a default factory or field options, not a literal argument default.

Hyper-parameters for cubic_regularized_gauss_newton_minimize.

Parameters

Name	Type & default	Contract
`steps`dataclass field	`int`Default: 50	See the operation contract and implementation below.
`sigma`dataclass field	`float`Default: 1.0	See the operation contract and implementation below.
`eta_accept`dataclass field	`float`Default: 0.1	See the operation contract and implementation below.
`eta_success`dataclass field	`float`Default: 0.9	See the operation contract and implementation below.
`sigma_increase`dataclass field	`float`Default: 2.0	See the operation contract and implementation below.
`sigma_decrease`dataclass field	`float`Default: 2.0	See the operation contract and implementation below.
`min_sigma`dataclass field	`float`Default: 1e-08	See the operation contract and implementation below.
`max_sigma`dataclass field	`float`Default: 1e+16	See the operation contract and implementation below.
`krylov_dim`dataclass field	`int \| None`Default: None	See the operation contract and implementation below.
`max_line_search`dataclass field	`int`Default: 12	See the operation contract and implementation below.

Executed regression · cubic gn config validation

Recorded coverage confirms this test executed the API body. This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗

def test_cubic_gn_config_validation() -> None:
    with pytest.raises(ValueError):
        CubicRegularizedGNConfig(sigma=0.0)
    with pytest.raises(ValueError):
        CubicRegularizedGNConfig(krylov_dim=0)

python -m pytest "packages/omnibias-jax/tests/test_optim.py::test_cubic_gn_config_validation" -q

Regression example · cubic gn rejects at minimum

This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗

def test_cubic_gn_rejects_at_minimum() -> None:
    t = jnp.linspace(0.0, 1.0, 16)
    y = 1.5 * jnp.exp(-0.3 * t)

    def residual_fn(p: Array) -> Array:
        return p[0] * jnp.exp(p[1] * t) - y

    p_star = jnp.array([1.5, -0.3])
    p1, hist = cubic_regularized_gauss_newton_minimize(
        residual_fn,
        p_star,
        config=CubicRegularizedGNConfig(steps=1, sigma=1e-4, krylov_dim=2),
    )
    assert jnp.allclose(p1, p_star, atol=1e-12)
    assert float(hist[-1]) <= float(hist[0]) + 1e-18

python -m pytest packages/omnibias-jax/tests/test_optim.py -k test_cubic_gn_rejects_at_minimum -q

Fields & methods

stepsfield

steps: int = 50

Declared int field on CubicRegularizedGNConfig. Default expression: 50.

View source ↗

sigmafield

sigma: float = 1.0

Declared float field on CubicRegularizedGNConfig. Default expression: 1.0.

View source ↗

eta_acceptfield

eta_accept: float = 0.1

Declared float field on CubicRegularizedGNConfig. Default expression: 0.1.

View source ↗

eta_successfield

eta_success: float = 0.9

Declared float field on CubicRegularizedGNConfig. Default expression: 0.9.

View source ↗

sigma_increasefield

sigma_increase: float = 2.0

Declared float field on CubicRegularizedGNConfig. Default expression: 2.0.

View source ↗

sigma_decreasefield

sigma_decrease: float = 2.0

Declared float field on CubicRegularizedGNConfig. Default expression: 2.0.

View source ↗

min_sigmafield

min_sigma: float = 1e-08

Declared float field on CubicRegularizedGNConfig. Default expression: 1e-08.

View source ↗

max_sigmafield

max_sigma: float = 1e+16

Declared float field on CubicRegularizedGNConfig. Default expression: 1e+16.

View source ↗

krylov_dimfield

krylov_dim: int | None = None

None = full parameter dimension. A small default (e.g. 20) silently truncates ARC on anything wider than a toy residual.

View source ↗

max_line_searchfield

max_line_search: int = 12

Declared int field on CubicRegularizedGNConfig. Default expression: 12.

View source ↗

__post_init__method

__post_init__(self) -> None

See the class contract above and the exact method implementation below.

Returns: None

Validation & failure cases

When self.sigma <= 0.0, raises ValueError(f'sigma must be > 0, got {self.sigma}').
When not 0.0 < self.eta_accept <= self.eta_success < 1.0, raises ValueError(f'need 0 < eta_accept <= eta_success < 1, got {self.eta_accept}, {self.eta_success}').
When self.sigma_increase <= 1.0, raises ValueError(f'sigma_increase must be > 1, got {self.sigma_increase}').
When self.sigma_decrease <= 1.0, raises ValueError(f'sigma_decrease must be > 1, got {self.sigma_decrease}').
When self.krylov_dim is not None and self.krylov_dim < 1, raises ValueError(f'krylov_dim must be >= 1 or None, got {self.krylov_dim}').
When self.max_line_search < 1, raises ValueError(f'max_line_search must be >= 1, got {self.max_line_search}').

Executed regression · cubic gn config validation

Full file, imports & fixtures ↗

def test_cubic_gn_config_validation() -> None:
    with pytest.raises(ValueError):
        CubicRegularizedGNConfig(sigma=0.0)
    with pytest.raises(ValueError):
        CubicRegularizedGNConfig(krylov_dim=0)

python -m pytest "packages/omnibias-jax/tests/test_optim.py::test_cubic_gn_config_validation" -q

Method implementation

    def __post_init__(self) -> None:
        if self.sigma <= 0.0:
            raise ValueError(f"sigma must be > 0, got {self.sigma}")
        if not 0.0 < self.eta_accept <= self.eta_success < 1.0:
            raise ValueError(
                f"need 0 < eta_accept <= eta_success < 1, got "
                f"{self.eta_accept}, {self.eta_success}"
            )
        if self.sigma_increase <= 1.0:
            raise ValueError(f"sigma_increase must be > 1, got {self.sigma_increase}")
        if self.sigma_decrease <= 1.0:
            raise ValueError(f"sigma_decrease must be > 1, got {self.sigma_decrease}")
        if self.krylov_dim is not None and self.krylov_dim < 1:
            raise ValueError(f"krylov_dim must be >= 1 or None, got {self.krylov_dim}")
        if self.max_line_search < 1:
            raise ValueError(f"max_line_search must be >= 1, got {self.max_line_search}")

View source ↗

Read the implementation

class CubicRegularizedGNConfig:
    """Hyper-parameters for :func:`cubic_regularized_gauss_newton_minimize`."""

    steps: int = 50
    sigma: float = 1.0
    eta_accept: float = 0.1
    eta_success: float = 0.9
    sigma_increase: float = 2.0
    sigma_decrease: float = 2.0
    min_sigma: float = 1e-8
    max_sigma: float = 1e16
    # None = full parameter dimension. A small default (e.g. 20) silently
    # truncates ARC on anything wider than a toy residual.
    krylov_dim: int | None = None
    max_line_search: int = 12

    def __post_init__(self) -> None:
        if self.sigma <= 0.0:
            raise ValueError(f"sigma must be > 0, got {self.sigma}")
        if not 0.0 < self.eta_accept <= self.eta_success < 1.0:
            raise ValueError(
                f"need 0 < eta_accept <= eta_success < 1, got "
                f"{self.eta_accept}, {self.eta_success}"
            )
        if self.sigma_increase <= 1.0:
            raise ValueError(f"sigma_increase must be > 1, got {self.sigma_increase}")
        if self.sigma_decrease <= 1.0:
            raise ValueError(f"sigma_decrease must be > 1, got {self.sigma_decrease}")
        if self.krylov_dim is not None and self.krylov_dim < 1:
            raise ValueError(f"krylov_dim must be >= 1 or None, got {self.krylov_dim}")
        if self.max_line_search < 1:
            raise ValueError(f"max_line_search must be >= 1, got {self.max_line_search}")

CONSTANT

GNSolver

omnibias.jax.optim.GNSolverSource ↗

GNSolver = Literal['dense', 'qr', 'cgls']

This exported constant is defined by the expression above. Its exact value and annotation come from the pinned source; use it with the module contract.

Read the implementation

GNSolver = Literal["dense", "qr", "cgls"]

CLASS

GaussNewtonState

omnibias.jax.optim.GaussNewtonStateSource ↗

class GaussNewtonState

Dataclass constructor fields; inherited fields, if any, are declared by the base class. A field(...) expression describes a default factory or field options, not a literal argument default.

Immutable LM optimiser state (functional update via gauss_newton_step).

Parameters

Name	Type & default	Contract
`params`dataclass field	`Array`Required	See the operation contract and implementation below.
`damping`dataclass field	`float`Required	See the operation contract and implementation below.
`loss`dataclass field	`float`Required	See the operation contract and implementation below.
`accepted`dataclass field	`bool`Required	See the operation contract and implementation below.
`n_iter`dataclass field	`int`Required	See the operation contract and implementation below.

Fields & methods

paramsfield

params: Array

Declared Array field on GaussNewtonState. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.

View source ↗

dampingfield

damping: float

Declared float field on GaussNewtonState. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.

View source ↗

lossfield

loss: float

Declared float field on GaussNewtonState. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.

View source ↗

acceptedfield

accepted: bool

Declared bool field on GaussNewtonState. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.

View source ↗

n_iterfield

n_iter: int

Declared int field on GaussNewtonState. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.

View source ↗

Read the implementation

class GaussNewtonState:
    """Immutable LM optimiser state (functional update via :func:`gauss_newton_step`)."""

    params: Array
    damping: float
    loss: float
    accepted: bool
    n_iter: int

CLASS

HomotopyGNConfig

omnibias.jax.optim.HomotopyGNConfigSource ↗

class HomotopyGNConfig

Dataclass constructor fields; inherited fields, if any, are declared by the base class. A field(...) expression describes a default factory or field options, not a literal argument default.

Stage-wise GN on r(theta, t) for a coupling homotopy t in [0, 1].

Each stage freezes t and runs Martens–Grosse or cubic GN. Designed for residuals that split as L + t Quad (nonlocal quadratic coupling). Does not weaken any downstream residual gate.

Parameters

Name	Type & default	Contract
`stages`dataclass field	`tuple[float, ...]`Default: (0.0, 0.25, 0.5, 0.75, 1.0)	See the operation contract and implementation below.
`steps_per_stage`dataclass field	`int`Default: 20	See the operation contract and implementation below.
`method`dataclass field	`Literal['martens_grosse', 'cubic']`Default: 'cubic'	See the operation contract and implementation below.
`damping`dataclass field	`float`Default: 0.001	See the operation contract and implementation below.
`cubic_sigma`dataclass field	`float`Default: 1.0	See the operation contract and implementation below.
`solver`dataclass field	`GNSolver`Default: 'qr'	See the operation contract and implementation below.
`krylov_dim`dataclass field	`int \| None`Default: None	See the operation contract and implementation below.

Executed regression · homotopy gn tracks linear coupling

Full file, imports & fixtures ↗

def test_homotopy_gn_tracks_linear_coupling() -> None:
    from omnibias.jax.optim import HomotopyGNConfig, homotopy_gauss_newton_minimize

    def residual_fn(p: Array, t: float) -> Array:
        return p - (1.0 + t)

    p1, info = homotopy_gauss_newton_minimize(
        residual_fn,
        jnp.asarray([0.0]),
        config=HomotopyGNConfig(
            stages=(0.0, 0.5, 1.0),
            steps_per_stage=4,
            method="martens_grosse",
            damping=1e-8,
        ),
    )
    assert float(jnp.abs(p1[0] - 2.0)) < 1e-8
    assert info["stages"][-1]["t"] == 1.0

python -m pytest "packages/omnibias-jax/tests/test_optim.py::test_homotopy_gn_tracks_linear_coupling" -q

Fields & methods

stagesfield

stages: tuple[float, ...] = (0.0, 0.25, 0.5, 0.75, 1.0)

Declared tuple[float, ...] field on HomotopyGNConfig. Default expression: (0.0, 0.25, 0.5, 0.75, 1.0).

View source ↗

steps_per_stagefield

steps_per_stage: int = 20

Declared int field on HomotopyGNConfig. Default expression: 20.

View source ↗

methodfield

method: Literal['martens_grosse', 'cubic'] = 'cubic'

Declared Literal['martens_grosse', 'cubic'] field on HomotopyGNConfig. Default expression: 'cubic'.

View source ↗

dampingfield

damping: float = 0.001

Declared float field on HomotopyGNConfig. Default expression: 0.001.

View source ↗

cubic_sigmafield

cubic_sigma: float = 1.0

Declared float field on HomotopyGNConfig. Default expression: 1.0.

View source ↗

solverfield

solver: GNSolver = 'qr'

Declared GNSolver field on HomotopyGNConfig. Default expression: 'qr'.

View source ↗

krylov_dimfield

krylov_dim: int | None = None

Declared int | None field on HomotopyGNConfig. Default expression: None.

View source ↗

__post_init__method

__post_init__(self) -> None

See the class contract above and the exact method implementation below.

Returns: None

Validation & failure cases

When not self.stages, raises ValueError('stages must be non-empty').
When any((not 0.0 <= float(t) <= 1.0 for t in self.stages)), raises ValueError(f'every stage t must lie in [0, 1], got {self.stages}').
When self.steps_per_stage < 1, raises ValueError(f'steps_per_stage must be >= 1, got {self.steps_per_stage}').
When self.method not in ('martens_grosse', 'cubic'), raises ValueError(f'method must be martens_grosse or cubic, got {self.method!r}').

Executed regression · homotopy gn tracks linear coupling

Full file, imports & fixtures ↗

def test_homotopy_gn_tracks_linear_coupling() -> None:
    from omnibias.jax.optim import HomotopyGNConfig, homotopy_gauss_newton_minimize

    def residual_fn(p: Array, t: float) -> Array:
        return p - (1.0 + t)

    p1, info = homotopy_gauss_newton_minimize(
        residual_fn,
        jnp.asarray([0.0]),
        config=HomotopyGNConfig(
            stages=(0.0, 0.5, 1.0),
            steps_per_stage=4,
            method="martens_grosse",
            damping=1e-8,
        ),
    )
    assert float(jnp.abs(p1[0] - 2.0)) < 1e-8
    assert info["stages"][-1]["t"] == 1.0

python -m pytest "packages/omnibias-jax/tests/test_optim.py::test_homotopy_gn_tracks_linear_coupling" -q

Method implementation

    def __post_init__(self) -> None:
        if not self.stages:
            raise ValueError("stages must be non-empty")
        if any(not 0.0 <= float(t) <= 1.0 for t in self.stages):
            raise ValueError(f"every stage t must lie in [0, 1], got {self.stages}")
        if self.steps_per_stage < 1:
            raise ValueError(f"steps_per_stage must be >= 1, got {self.steps_per_stage}")
        if self.method not in ("martens_grosse", "cubic"):
            raise ValueError(f"method must be martens_grosse or cubic, got {self.method!r}")

View source ↗

Read the implementation

class HomotopyGNConfig:
    """Stage-wise GN on ``r(theta, t)`` for a coupling homotopy ``t`` in ``[0, 1]``.

    Each stage freezes ``t`` and runs Martens–Grosse or cubic GN. Designed for
    residuals that split as ``L + t Quad`` (nonlocal quadratic coupling).
    Does not weaken any downstream residual gate.
    """

    stages: tuple[float, ...] = (0.0, 0.25, 0.5, 0.75, 1.0)
    steps_per_stage: int = 20
    method: Literal["martens_grosse", "cubic"] = "cubic"
    damping: float = 1e-3
    cubic_sigma: float = 1.0
    solver: GNSolver = "qr"
    krylov_dim: int | None = None

    def __post_init__(self) -> None:
        if not self.stages:
            raise ValueError("stages must be non-empty")
        if any(not 0.0 <= float(t) <= 1.0 for t in self.stages):
            raise ValueError(f"every stage t must lie in [0, 1], got {self.stages}")
        if self.steps_per_stage < 1:
            raise ValueError(f"steps_per_stage must be >= 1, got {self.steps_per_stage}")
        if self.method not in ("martens_grosse", "cubic"):
            raise ValueError(f"method must be martens_grosse or cubic, got {self.method!r}")

CONSTANT

HomotopyResidualFn

omnibias.jax.optim.HomotopyResidualFnSource ↗

HomotopyResidualFn = Callable[[Array, float], Array]

This exported constant is defined by the expression above. Its exact value and annotation come from the pinned source; use it with the module contract.

Read the implementation

HomotopyResidualFn = Callable[[Array, float], Array]

CLASS

LinfMinimaxConfig

omnibias.jax.optim.LinfMinimaxConfigSource ↗

class LinfMinimaxConfig

Dataclass constructor fields; inherited fields, if any, are declared by the base class. A field(...) expression describes a default factory or field options, not a literal argument default.

Hyper-parameters for linf_minimax_minimize / linf_minimax_step.

box is the trust-region radius passed straight through to linearized_linf_direction's box clip -- it is not a Levenberg-Marquardt damping (there is no matrix to damp; the epigraph LP / Lawson-IRLS direction is already a full minimax solve of the *linearized* problem).

Parameters

Name	Type & default	Contract
`steps`dataclass field	`int`Default: 50	See the operation contract and implementation below.
`box0`dataclass field	`float`Default: 1.0	See the operation contract and implementation below.
`box_decrease`dataclass field	`float`Default: 0.5	See the operation contract and implementation below.
`box_increase`dataclass field	`float`Default: 1.5	See the operation contract and implementation below.
`min_box`dataclass field	`float`Default: 1e-10	See the operation contract and implementation below.
`max_box`dataclass field	`float`Default: 1000000.0	See the operation contract and implementation below.
`irls_iters`dataclass field	`int`Default: 16	See the operation contract and implementation below.
`max_line_search`dataclass field	`int`Default: 20	See the operation contract and implementation below.
`accept_tol`dataclass field	`float`Default: 0.0	See the operation contract and implementation below.

Executed regression · linf minimax minimize reaches exact minimax on toy problem

Full file, imports & fixtures ↗

def test_linf_minimax_minimize_reaches_exact_minimax_on_toy_problem() -> None:
    from omnibias.jax.optim import LinfMinimaxConfig, linf_minimax_minimize

    c = jnp.asarray([1.0, -0.2, 0.4])

    def residual_fn(theta: Array) -> Array:
        return theta[0] + c

    theta0 = jnp.asarray([0.0])
    theta1, history = linf_minimax_minimize(
        residual_fn, theta0, config=LinfMinimaxConfig(steps=5, box0=1.0)
    )
    assert float(theta1[0]) == pytest.approx(-0.4, abs=1e-8)
    assert float(history[-1]) == pytest.approx(0.6, abs=1e-8)
    # The safeguard makes the recorded max|r| history non-increasing.
    assert all(
        float(history[i + 1]) <= float(history[i]) + 1e-12 for i in range(len(history) - 1)
    )

python -m pytest "packages/omnibias-jax/tests/test_optim.py::test_linf_minimax_minimize_reaches_exact_minimax_on_toy_problem" -q

Regression example · linf minimax config rejects invalid hyperparameters

This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗

def test_linf_minimax_config_rejects_invalid_hyperparameters() -> None:
    from omnibias.jax.optim import LinfMinimaxConfig

    with pytest.raises(ValueError):
        LinfMinimaxConfig(steps=0)
    with pytest.raises(ValueError):
        LinfMinimaxConfig(box0=0.0)
    with pytest.raises(ValueError):
        LinfMinimaxConfig(box_decrease=1.0)
    with pytest.raises(ValueError):
        LinfMinimaxConfig(box_increase=0.5)
    with pytest.raises(ValueError):
        LinfMinimaxConfig(min_box=0.0)
    with pytest.raises(ValueError):
        LinfMinimaxConfig(max_box=1e-3, min_box=1e-2)
    with pytest.raises(ValueError):
        LinfMinimaxConfig(irls_iters=0)
    with pytest.raises(ValueError):
        LinfMinimaxConfig(max_line_search=0)
    with pytest.raises(ValueError):
        LinfMinimaxConfig(accept_tol=-1.0)

python -m pytest packages/omnibias-jax/tests/test_optim.py -k test_linf_minimax_config_rejects_invalid_hyperparameters -q

Fields & methods

stepsfield

steps: int = 50

Declared int field on LinfMinimaxConfig. Default expression: 50.

View source ↗

box0field

box0: float = 1.0

Declared float field on LinfMinimaxConfig. Default expression: 1.0.

View source ↗

box_decreasefield

box_decrease: float = 0.5

Declared float field on LinfMinimaxConfig. Default expression: 0.5.

View source ↗

box_increasefield

box_increase: float = 1.5

Declared float field on LinfMinimaxConfig. Default expression: 1.5.

View source ↗

min_boxfield

min_box: float = 1e-10

Declared float field on LinfMinimaxConfig. Default expression: 1e-10.

View source ↗

max_boxfield

max_box: float = 1000000.0

Declared float field on LinfMinimaxConfig. Default expression: 1000000.0.

View source ↗

irls_itersfield

irls_iters: int = 16

Declared int field on LinfMinimaxConfig. Default expression: 16.

View source ↗

max_line_searchfield

max_line_search: int = 20

Declared int field on LinfMinimaxConfig. Default expression: 20.

View source ↗

accept_tolfield

accept_tol: float = 0.0

Declared float field on LinfMinimaxConfig. Default expression: 0.0.

View source ↗

__post_init__method

__post_init__(self) -> None

See the class contract above and the exact method implementation below.

Returns: None

Validation & failure cases

When self.steps < 1, raises ValueError(f'steps must be >= 1, got {self.steps}').
When self.box0 <= 0.0, raises ValueError(f'box0 must be > 0, got {self.box0}').
When not 0.0 < self.box_decrease < 1.0, raises ValueError(f'box_decrease must be in (0, 1), got {self.box_decrease}').
When self.box_increase < 1.0, raises ValueError(f'box_increase must be >= 1, got {self.box_increase}').
When self.min_box <= 0.0, raises ValueError(f'min_box must be > 0, got {self.min_box}').
When self.max_box < self.min_box, raises ValueError(f'max_box must be >= min_box, got {self.max_box} < {self.min_box}').
When self.irls_iters < 1, raises ValueError(f'irls_iters must be >= 1, got {self.irls_iters}').
When self.max_line_search < 1, raises ValueError(f'max_line_search must be >= 1, got {self.max_line_search}').
When self.accept_tol < 0.0, raises ValueError(f'accept_tol must be >= 0, got {self.accept_tol}').

Executed regression · linf minimax minimize reaches exact minimax on toy problem

Full file, imports & fixtures ↗

def test_linf_minimax_minimize_reaches_exact_minimax_on_toy_problem() -> None:
    from omnibias.jax.optim import LinfMinimaxConfig, linf_minimax_minimize

    c = jnp.asarray([1.0, -0.2, 0.4])

    def residual_fn(theta: Array) -> Array:
        return theta[0] + c

    theta0 = jnp.asarray([0.0])
    theta1, history = linf_minimax_minimize(
        residual_fn, theta0, config=LinfMinimaxConfig(steps=5, box0=1.0)
    )
    assert float(theta1[0]) == pytest.approx(-0.4, abs=1e-8)
    assert float(history[-1]) == pytest.approx(0.6, abs=1e-8)
    # The safeguard makes the recorded max|r| history non-increasing.
    assert all(
        float(history[i + 1]) <= float(history[i]) + 1e-12 for i in range(len(history) - 1)
    )

python -m pytest "packages/omnibias-jax/tests/test_optim.py::test_linf_minimax_minimize_reaches_exact_minimax_on_toy_problem" -q

Method implementation

    def __post_init__(self) -> None:
        if self.steps < 1:
            raise ValueError(f"steps must be >= 1, got {self.steps}")
        if self.box0 <= 0.0:
            raise ValueError(f"box0 must be > 0, got {self.box0}")
        if not 0.0 < self.box_decrease < 1.0:
            raise ValueError(f"box_decrease must be in (0, 1), got {self.box_decrease}")
        if self.box_increase < 1.0:
            raise ValueError(f"box_increase must be >= 1, got {self.box_increase}")
        if self.min_box <= 0.0:
            raise ValueError(f"min_box must be > 0, got {self.min_box}")
        if self.max_box < self.min_box:
            raise ValueError(
                f"max_box must be >= min_box, got {self.max_box} < {self.min_box}"
            )
        if self.irls_iters < 1:
            raise ValueError(f"irls_iters must be >= 1, got {self.irls_iters}")
        if self.max_line_search < 1:
            raise ValueError(f"max_line_search must be >= 1, got {self.max_line_search}")
        if self.accept_tol < 0.0:
            raise ValueError(f"accept_tol must be >= 0, got {self.accept_tol}")

View source ↗

Read the implementation

class LinfMinimaxConfig:
    """Hyper-parameters for :func:`linf_minimax_minimize` / :func:`linf_minimax_step`.

    ``box`` is the trust-region radius passed straight through to
    :func:`linearized_linf_direction`'s ``box`` clip -- it is not a Levenberg-Marquardt
    damping (there is no matrix to damp; the epigraph LP / Lawson-IRLS direction is
    already a full minimax solve of the *linearized* problem).
    """

    steps: int = 50
    box0: float = 1.0
    box_decrease: float = 0.5
    box_increase: float = 1.5
    min_box: float = 1e-10
    max_box: float = 1e6
    irls_iters: int = 16
    max_line_search: int = 20
    accept_tol: float = 0.0

    def __post_init__(self) -> None:
        if self.steps < 1:
            raise ValueError(f"steps must be >= 1, got {self.steps}")
        if self.box0 <= 0.0:
            raise ValueError(f"box0 must be > 0, got {self.box0}")
        if not 0.0 < self.box_decrease < 1.0:
            raise ValueError(f"box_decrease must be in (0, 1), got {self.box_decrease}")
        if self.box_increase < 1.0:
            raise ValueError(f"box_increase must be >= 1, got {self.box_increase}")
        if self.min_box <= 0.0:
            raise ValueError(f"min_box must be > 0, got {self.min_box}")
        if self.max_box < self.min_box:
            raise ValueError(
                f"max_box must be >= min_box, got {self.max_box} < {self.min_box}"
            )
        if self.irls_iters < 1:
            raise ValueError(f"irls_iters must be >= 1, got {self.irls_iters}")
        if self.max_line_search < 1:
            raise ValueError(f"max_line_search must be >= 1, got {self.max_line_search}")
        if self.accept_tol < 0.0:
            raise ValueError(f"accept_tol must be >= 0, got {self.accept_tol}")

CLASS

MartensGrosseGNConfig

omnibias.jax.optim.MartensGrosseGNConfigSource ↗

class MartensGrosseGNConfig

Dataclass constructor fields; inherited fields, if any, are declared by the base class. A field(...) expression describes a default factory or field options, not a literal argument default.

Hyper-parameters for martens_grosse_gauss_newton_minimize.

Parameters

Name	Type & default	Contract
`steps`dataclass field	`int`Default: 50	See the operation contract and implementation below.
`damping`dataclass field	`float`Default: 0.001	See the operation contract and implementation below.
`damping_decrease`dataclass field	`float`Default: 0.7	See the operation contract and implementation below.
`damping_increase`dataclass field	`float`Default: 2.0	See the operation contract and implementation below.
`min_damping`dataclass field	`float`Default: 1e-08	See the operation contract and implementation below.
`max_damping`dataclass field	`float`Default: 1000.0	See the operation contract and implementation below.
`accept_tol`dataclass field	`float`Default: 0.0	See the operation contract and implementation below.
`use_martens_grosse`dataclass field	`bool`Default: True	See the operation contract and implementation below.
`solver`dataclass field	`GNSolver`Default: 'qr'	See the operation contract and implementation below.
`cgls_max_iter`dataclass field	`int`Default: 100	See the operation contract and implementation below.
`cgls_tol`dataclass field	`float`Default: 1e-08	See the operation contract and implementation below.

Regression example · martens grosse gauss newton minimize recovers nonlinear ls

This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗

def test_martens_grosse_gauss_newton_minimize_recovers_nonlinear_ls() -> None:
    from omnibias.jax.optim import (
        MartensGrosseGNConfig,
        martens_grosse_gauss_newton_minimize,
    )

    t = jnp.linspace(0.0, 1.0, 24)
    true = jnp.array([2.0, -0.7])
    y = true[0] * jnp.exp(true[1] * t)

    def residual_fn(p: Array) -> Array:
        return p[0] * jnp.exp(p[1] * t) - y

    p0 = jnp.array([1.0, 0.0])
    p1, hist = martens_grosse_gauss_newton_minimize(
        residual_fn,
        p0,
        config=MartensGrosseGNConfig(steps=40, damping=1e-3, solver="qr"),
    )
    assert float(hist[-1]) < 1e-16
    assert jnp.allclose(p1, true, atol=1e-6)

python -m pytest packages/omnibias-jax/tests/test_optim.py -k test_martens_grosse_gauss_newton_minimize_recovers_nonlinear_ls -q

Regression example · martens grosse gauss newton minimize unaffected by linf addition

This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗

def test_martens_grosse_gauss_newton_minimize_unaffected_by_linf_addition() -> None:
    """The pre-existing GN trainer's behavior on its own regression fixture is
    unchanged by adding the L-infinity sibling (same numbers as
    ``test_martens_grosse_gauss_newton_minimize_recovers_nonlinear_ls``)."""
    from omnibias.jax.optim import (
        MartensGrosseGNConfig,
        martens_grosse_gauss_newton_minimize,
    )

    t = jnp.linspace(0.0, 1.0, 24)
    true = jnp.array([2.0, -0.7])
    y = true[0] * jnp.exp(true[1] * t)

    def residual_fn(p: Array) -> Array:
        return p[0] * jnp.exp(p[1] * t) - y

    p0 = jnp.array([1.0, 0.0])
    p1, hist = martens_grosse_gauss_newton_minimize(
        residual_fn,
        p0,
        config=MartensGrosseGNConfig(steps=40, damping=1e-3, solver="qr"),
    )
    assert float(hist[-1]) < 1e-16
    assert jnp.allclose(p1, true, atol=1e-6)

python -m pytest packages/omnibias-jax/tests/test_optim.py -k test_martens_grosse_gauss_newton_minimize_unaffected_by_linf_addition -q

Fields & methods

stepsfield

steps: int = 50

Declared int field on MartensGrosseGNConfig. Default expression: 50.

View source ↗

dampingfield

damping: float = 0.001

Declared float field on MartensGrosseGNConfig. Default expression: 0.001.

View source ↗

damping_decreasefield

damping_decrease: float = 0.7

Declared float field on MartensGrosseGNConfig. Default expression: 0.7.

View source ↗

damping_increasefield

damping_increase: float = 2.0

Declared float field on MartensGrosseGNConfig. Default expression: 2.0.

View source ↗

min_dampingfield

min_damping: float = 1e-08

Declared float field on MartensGrosseGNConfig. Default expression: 1e-08.

View source ↗

max_dampingfield

max_damping: float = 1000.0

Declared float field on MartensGrosseGNConfig. Default expression: 1000.0.

View source ↗

accept_tolfield

accept_tol: float = 0.0

Declared float field on MartensGrosseGNConfig. Default expression: 0.0.

View source ↗

use_martens_grossefield

use_martens_grosse: bool = True

Declared bool field on MartensGrosseGNConfig. Default expression: True.

View source ↗

solverfield

solver: GNSolver = 'qr'

Declared GNSolver field on MartensGrosseGNConfig. Default expression: 'qr'.

View source ↗

cgls_max_iterfield

cgls_max_iter: int = 100

Declared int field on MartensGrosseGNConfig. Default expression: 100.

View source ↗

cgls_tolfield

cgls_tol: float = 1e-08

Declared float field on MartensGrosseGNConfig. Default expression: 1e-08.

View source ↗

Read the implementation

class MartensGrosseGNConfig:
    """Hyper-parameters for :func:`martens_grosse_gauss_newton_minimize`."""

    steps: int = 50
    damping: float = 1e-3
    damping_decrease: float = 0.7
    damping_increase: float = 2.0
    min_damping: float = 1e-8
    max_damping: float = 1e3
    accept_tol: float = 0.0
    use_martens_grosse: bool = True
    solver: GNSolver = "qr"
    cgls_max_iter: int = 100
    cgls_tol: float = 1e-8

CONSTANT

MatVec

omnibias.jax.optim.MatVecSource ↗

MatVec = Callable[[Array], Array]

This exported constant is defined by the expression above. Its exact value and annotation come from the pinned source; use it with the module contract.

Read the implementation

MatVec = Callable[[Array], Array]

CONSTANT

ResidualFn

omnibias.jax.optim.ResidualFnSource ↗

ResidualFn = Callable[[Array], Array]

This exported constant is defined by the expression above. Its exact value and annotation come from the pinned source; use it with the module contract.

Read the implementation

ResidualFn = Callable[[Array], Array]

FUNCTION

cgls

omnibias.jax.optim.cglsSource ↗

cgls(a_matvec: MatVec, at_matvec: MatVec, b: Array, *, damp: float=0.0, max_iter: int=100, tol: float=1e-08) -> Array

Solve min_x ||A x - b||^2 + damp^2 ||x||^2 by CGLS, matrix-free on A.

Twin of omnibias.torch.optim.cgls. Accuracy scales with kappa(A), not kappa(A)^2.

Parameters

Name	Type & default	Contract
`a_matvec`positional-or-keyword	`MatVec`Required	See the operation contract and implementation below.
`at_matvec`positional-or-keyword	`MatVec`Required	See the operation contract and implementation below.
`b`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`damp`keyword-only	`float`Default: 0.0	See the operation contract and implementation below.
`max_iter`keyword-only	`int`Default: 100	See the operation contract and implementation below.
`tol`keyword-only	`float`Default: 1e-08	See the operation contract and implementation below.

ReturnsArray

Executed regression · cgls matches qr on small problem

Full file, imports & fixtures ↗

def test_cgls_matches_qr_on_small_problem() -> None:
    from omnibias.jax.optim import (
        gauss_newton_direction_cgls,
        lstsq_gauss_newton_direction,
    )

    rng = np.random.RandomState(11)
    a = jnp.asarray(rng.randn(20, 5))
    true = jnp.asarray(rng.randn(5))
    b = a @ true

    def residual_fn(p: Array) -> Array:
        return a @ p - b

    p0 = jnp.zeros(5)
    jac = a
    r0 = residual_fn(p0)
    d_qr = lstsq_gauss_newton_direction(jac, r0, 1e-8)
    d_cgls = gauss_newton_direction_cgls(residual_fn, p0, 1e-8, cgls_tol=1e-14)
    assert jnp.allclose(d_qr, d_cgls, atol=1e-6)

python -m pytest "packages/omnibias-jax/tests/test_optim.py::test_cgls_matches_qr_on_small_problem" -q

Read the implementation

def cgls(
    a_matvec: MatVec,
    at_matvec: MatVec,
    b: Array,
    *,
    damp: float = 0.0,
    max_iter: int = 100,
    tol: float = 1e-8,
) -> Array:
    r"""Solve ``min_x ||A x - b||^2 + damp^2 ||x||^2`` by CGLS, matrix-free on ``A``.

    Twin of :func:`omnibias.torch.optim.cgls`. Accuracy scales with ``kappa(A)``, not
    ``kappa(A)^2``.
    """
    lam = float(damp) ** 2
    r = b
    s = at_matvec(r)
    x = jnp.zeros_like(s)
    p = s
    gamma = jnp.vdot(s, s)
    norm_s0 = float(jnp.sqrt(gamma))
    if norm_s0 == 0.0:
        return x
    thresh = tol * norm_s0
    # Eager Python loop (data-dependent stop) -- matches torch cgls control flow.
    for _ in range(max_iter):
        if float(jnp.sqrt(gamma)) <= thresh:
            break
        q = a_matvec(p)
        denom = float(jnp.vdot(q, q) + lam * jnp.vdot(p, p))
        if denom <= 0.0:
            break
        alpha = gamma / denom
        x = x + alpha * p
        r = r - alpha * q
        s = at_matvec(r) - lam * x
        gamma_new = jnp.vdot(s, s)
        if float(jnp.sqrt(gamma_new)) <= thresh:
            break
        p = s + (gamma_new / gamma) * p
        gamma = gamma_new
    return x

FUNCTION

champ_barrier_residual

omnibias.jax.optim.champ_barrier_residualSource ↗

champ_barrier_residual(res: Array, tau: float, *, barrier_weight: float=20.0, peak_power: float=4.0, peak_frac: float=0.85) -> Array

Active-set L^∞ residual that cannot raise already-small entries above tau.

Rows with |r| >= peak_frac * max|r| are peak-weighted so Gauss–Newton works the current spike. One-sided excess relu(|r| - tau) is a barrier against walking past a known champ. The peak block is zero wherever res is; the barrier block is zero whenever |r| <= tau.

Parameters

Name	Type & default	Contract
`res`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`tau`positional-or-keyword	`float`Required	See the operation contract and implementation below.
`barrier_weight`keyword-only	`float`Default: 20.0	See the operation contract and implementation below.
`peak_power`keyword-only	`float`Default: 4.0	See the operation contract and implementation below.
`peak_frac`keyword-only	`float`Default: 0.85	See the operation contract and implementation below.

ReturnsArray

Executed regression · champ barrier residual holds zeros and tau

Full file, imports & fixtures ↗

def test_champ_barrier_residual_holds_zeros_and_tau() -> None:
    from omnibias.jax.optim import champ_barrier_residual

    r = jnp.asarray([0.0, 0.01, -0.02, 0.018])
    out = champ_barrier_residual(r, 0.021, barrier_weight=10.0, peak_power=2.0)
    n = int(r.size)
    peak, barrier = out[:n], out[n:]
    assert float(peak[0]) == 0.0
    assert float(jnp.max(jnp.abs(barrier))) == 0.0
    walked = champ_barrier_residual(r, 0.015, barrier_weight=10.0, peak_power=2.0)
    assert float(jnp.max(jnp.abs(walked[n:]))) > 0.0

python -m pytest "packages/omnibias-jax/tests/test_optim.py::test_champ_barrier_residual_holds_zeros_and_tau" -q

Read the implementation

def champ_barrier_residual(
    res: Array,
    tau: float,
    *,
    barrier_weight: float = 20.0,
    peak_power: float = 4.0,
    peak_frac: float = 0.85,
) -> Array:
    """Active-set L^∞ residual that cannot raise already-small entries above ``tau``.

    Rows with ``|r| >= peak_frac * max|r|`` are peak-weighted so Gauss–Newton
    works the current spike. One-sided excess ``relu(|r| - tau)`` is a barrier
    against walking past a known champ. The peak block is zero wherever
    ``res`` is; the barrier block is zero whenever ``|r| <= tau``.
    """
    ar = jnp.abs(res)
    rmax = jnp.max(ar)
    peak_mask = ar >= (float(peak_frac) * rmax)
    peak = jnp.where(
        peak_mask,
        peak_weighted_residual(res, peak_power),
        jnp.zeros_like(res),
    )
    excess = jnp.maximum(ar - float(tau), 0.0) * jnp.sign(res)
    return jnp.concatenate([peak, float(barrier_weight) * excess])

FUNCTION

cubic_regularized_gauss_newton_minimize

omnibias.jax.optim.cubic_regularized_gauss_newton_minimizeSource ↗

cubic_regularized_gauss_newton_minimize(residual_fn: ResidualFn, params0: Array, *, config: CubicRegularizedGNConfig | None=None) -> tuple[Array, Array]

Minimise 0.5 ||r(params)||^2 by cubic-regularised Gauss-Newton (ARC).

Matrix-free twin of omnibias.torch.optim.CubicRegularizedGaussNewton: each step minimises g^T s + 0.5 s^T (J^T J) s + (sigma/3)||s||^3 with g = J^T r in a Lanczos subspace, then accepts by the ARC ratio test. Loss history matches martens_grosse_gauss_newton_minimize (value at the start of each step, then the final value).

Parameters

Name	Type & default	Contract
`residual_fn`positional-or-keyword	`ResidualFn`Required	See the operation contract and implementation below.
`params0`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`config`keyword-only	`CubicRegularizedGNConfig \| None`Default: None	See the operation contract and implementation below.

Returnstuple[Array, Array]

Executed regression · cubic gn rejects at minimum

Full file, imports & fixtures ↗

def test_cubic_gn_rejects_at_minimum() -> None:
    t = jnp.linspace(0.0, 1.0, 16)
    y = 1.5 * jnp.exp(-0.3 * t)

    def residual_fn(p: Array) -> Array:
        return p[0] * jnp.exp(p[1] * t) - y

    p_star = jnp.array([1.5, -0.3])
    p1, hist = cubic_regularized_gauss_newton_minimize(
        residual_fn,
        p_star,
        config=CubicRegularizedGNConfig(steps=1, sigma=1e-4, krylov_dim=2),
    )
    assert jnp.allclose(p1, p_star, atol=1e-12)
    assert float(hist[-1]) <= float(hist[0]) + 1e-18

python -m pytest "packages/omnibias-jax/tests/test_optim.py::test_cubic_gn_rejects_at_minimum" -q

Regression example · cubic gn default krylov is full dimension

This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗

def test_cubic_gn_default_krylov_is_full_dimension() -> None:
    rng = np.random.RandomState(2)
    mat = jnp.asarray(rng.randn(40, 18))
    b = jnp.asarray(rng.randn(40))

    def residual_fn(p: Array) -> Array:
        return mat @ p - b

    p_star, _, _, _ = jnp.linalg.lstsq(mat, b, rcond=None)
    p, _hist = cubic_regularized_gauss_newton_minimize(
        residual_fn,
        jnp.zeros(18),
        config=CubicRegularizedGNConfig(steps=20, sigma=1.0),
    )
    assert jnp.allclose(p, p_star, atol=1e-6)

python -m pytest packages/omnibias-jax/tests/test_optim.py -k test_cubic_gn_default_krylov_is_full_dimension -q

Read the implementation

def cubic_regularized_gauss_newton_minimize(
    residual_fn: ResidualFn,
    params0: Array,
    *,
    config: CubicRegularizedGNConfig | None = None,
) -> tuple[Array, Array]:
    r"""Minimise ``0.5 ||r(params)||^2`` by cubic-regularised Gauss-Newton (ARC).

    Matrix-free twin of :class:`omnibias.torch.optim.CubicRegularizedGaussNewton`:
    each step minimises ``g^T s + 0.5 s^T (J^T J) s + (sigma/3)||s||^3`` with
    ``g = J^T r`` in a Lanczos subspace, then accepts by the ARC ratio test.
    Loss history matches :func:`martens_grosse_gauss_newton_minimize` (value at
    the start of each step, then the final value).
    """
    cfg = CubicRegularizedGNConfig() if config is None else config
    params = params0
    sigma = float(cfg.sigma)
    n_params = int(params0.size)
    krylov = n_params if cfg.krylov_dim is None else min(int(cfg.krylov_dim), n_params)
    losses: list[float] = []

    for _ in range(int(cfg.steps)):
        res, jt, jvec = _linearize_gn(residual_fn, params)
        g = jt(res)
        f0 = _half_sum_sq(res)
        losses.append(f0)

        def matvec(v: Array, jt_fn: MatVec = jt, jvec_fn: MatVec = jvec) -> Array:
            return jt_fn(jvec_fn(v))

        def eval_f(trial: Array) -> float:
            return _half_sum_sq(residual_fn(trial))

        params, _accepted, sigma, _rho, _f1 = _cubic_arc_step(
            params,
            g,
            matvec,
            f0,
            eval_f,
            sigma,
            krylov_dim=krylov,
            max_line_search=int(cfg.max_line_search),
            eta_accept=float(cfg.eta_accept),
            eta_success=float(cfg.eta_success),
            sigma_increase=float(cfg.sigma_increase),
            sigma_decrease=float(cfg.sigma_decrease),
            min_sigma=float(cfg.min_sigma),
            max_sigma=float(cfg.max_sigma),
        )

    losses.append(_half_sum_sq(residual_fn(params)))
    return params, jnp.asarray(losses, dtype=jnp.float64)

FUNCTION

cubic_regularized_newton_step

omnibias.jax.optim.cubic_regularized_newton_stepSource ↗

cubic_regularized_newton_step(loss_fn: Callable[[Array], Array], params: Array, sigma: float, *, krylov_dim: int=20) -> Array

One cubic-regularised Newton step (ARC), matrix-free via Lanczos.

Twin of omnibias.torch.optim.cubic_regularized_newton_step. Minimises g^T s + 0.5 s^T H s + (sigma/3) ||s||^3 in a Krylov subspace of the exact Hessian.

Parameters

Name	Type & default	Contract
`loss_fn`positional-or-keyword	`Callable[[Array], Array]`Required	See the operation contract and implementation below.
`params`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`sigma`positional-or-keyword	`float`Required	See the operation contract and implementation below.
`krylov_dim`keyword-only	`int`Default: 20	See the operation contract and implementation below.

ReturnsArray

Read the implementation

def cubic_regularized_newton_step(
    loss_fn: Callable[[Array], Array],
    params: Array,
    sigma: float,
    *,
    krylov_dim: int = 20,
) -> Array:
    r"""One cubic-regularised Newton step (ARC), matrix-free via Lanczos.

    Twin of :func:`omnibias.torch.optim.cubic_regularized_newton_step`.
    Minimises ``g^T s + 0.5 s^T H s + (sigma/3) ||s||^3`` in a Krylov
    subspace of the exact Hessian.
    """
    grad_fn = jax.grad(loss_fn)
    g = grad_fn(params)
    gnorm = float(jnp.linalg.norm(g))
    if gnorm == 0.0:
        return jnp.zeros_like(params)

    def matvec(v: Array) -> Array:
        return cast(Array, jax.jvp(grad_fn, (params,), (v,))[1])

    q_basis, tri = lanczos_tridiag(matvec, g, krylov_dim)
    c = jnp.zeros((tri.shape[0],), dtype=params.dtype).at[0].set(gnorm)
    return q_basis @ _solve_cubic_subproblem(tri, c, sigma)

FUNCTION

gauss_newton_direction

omnibias.jax.optim.gauss_newton_directionSource ↗

gauss_newton_direction(jac: Array, res: Array, damping: float) -> Array

Levenberg-Marquardt Gauss-Newton step delta for min 1/2 ||r||^2.

Solves (J^T J + mu I) delta = -J^T r (primal, P x P) when P <= N and the equivalent dual form delta = -J^T (J J^T + mu I)^{-1} r (N x N) otherwise.

Parameters

jac: Residual Jacobian J = d r / d theta of shape (N, P). res: Residual vector r of shape (N,). damping: Levenberg-Marquardt damping mu >= 0.

Parameters

Name	Type & default	Contract
`jac`positional-or-keyword	`Array`Required	Residual Jacobian `J = d r / d theta` of shape `(N, P)`.
`res`positional-or-keyword	`Array`Required	Residual vector `r` of shape `(N,)`.
`damping`positional-or-keyword	`float`Required	Levenberg-Marquardt damping `mu >= 0`.

ReturnsArray

Explicit validation & failure cases

When jac.ndim != 2, raises ValueError(f'jac must be 2-D (N, P), got shape {jac.shape}').
When res.shape != (n,), raises ValueError(f'res must have shape ({n},) matching jac rows, got {res.shape}').

These are explicit source guards. Backend operations may also raise for invalid shapes, dtypes, or unsupported devices.

Executed regression · gn direction shape validation

Full file, imports & fixtures ↗

def test_gn_direction_shape_validation() -> None:
    with pytest.raises(ValueError):
        gauss_newton_direction(jnp.zeros(3), jnp.zeros(3), 1e-3)
    with pytest.raises(ValueError):
        gauss_newton_direction(jnp.zeros((4, 2)), jnp.zeros(3), 1e-3)

python -m pytest "packages/omnibias-jax/tests/test_optim.py::test_gn_direction_shape_validation" -q

Regression example · gn direction solves linear least squares in one step

This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗

def test_gn_direction_solves_linear_least_squares_in_one_step() -> None:
    rng = np.random.RandomState(0)
    a = jnp.asarray(rng.randn(20, 5))
    b = jnp.asarray(rng.randn(20))
    p0 = jnp.zeros(5)
    delta = gauss_newton_direction(a, a @ p0 - b, damping=1e-10)
    p_star, _, _, _ = jnp.linalg.lstsq(a, b, rcond=None)
    assert jnp.allclose(p0 + delta, p_star, atol=1e-6)

python -m pytest packages/omnibias-jax/tests/test_optim.py -k test_gn_direction_solves_linear_least_squares_in_one_step -q

Read the implementation

def gauss_newton_direction(jac: Array, res: Array, damping: float) -> Array:
    r"""Levenberg-Marquardt Gauss-Newton step ``delta`` for ``min 1/2 ||r||^2``.

    Solves ``(J^T J + mu I) delta = -J^T r`` (primal, ``P x P``) when ``P <= N`` and the
    equivalent dual form ``delta = -J^T (J J^T + mu I)^{-1} r`` (``N x N``) otherwise.

    Parameters
    ----------
    jac:
        Residual Jacobian ``J = d r / d theta`` of shape ``(N, P)``.
    res:
        Residual vector ``r`` of shape ``(N,)``.
    damping:
        Levenberg-Marquardt damping ``mu >= 0``.
    """
    if jac.ndim != 2:
        raise ValueError(f"jac must be 2-D (N, P), got shape {jac.shape}")
    n, p = jac.shape
    if res.shape != (n,):
        raise ValueError(f"res must have shape ({n},) matching jac rows, got {res.shape}")
    mu = jnp.asarray(damping, dtype=jac.dtype)
    if p <= n:
        a = jac.T @ jac + mu * jnp.eye(p, dtype=jac.dtype)
        delta: Array = jnp.linalg.solve(a, -(jac.T @ res))
    else:
        g = jac @ jac.T + mu * jnp.eye(n, dtype=jac.dtype)
        delta = -(jac.T @ jnp.linalg.solve(g, res))
    return delta

FUNCTION

gauss_newton_direction_cgls

omnibias.jax.optim.gauss_newton_direction_cglsSource ↗

gauss_newton_direction_cgls(residual_fn: ResidualFn, params: Array, damping: float, *, cgls_max_iter: int=100, cgls_tol: float=1e-08) -> Array

Matrix-free LM Gauss-Newton direction via CGLS.

Twin of omnibias.torch.optim.gauss_newton_direction_cgls.

Parameters

Name	Type & default	Contract
`residual_fn`positional-or-keyword	`ResidualFn`Required	See the operation contract and implementation below.
`params`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`damping`positional-or-keyword	`float`Required	See the operation contract and implementation below.
`cgls_max_iter`keyword-only	`int`Default: 100	See the operation contract and implementation below.
`cgls_tol`keyword-only	`float`Default: 1e-08	See the operation contract and implementation below.

ReturnsArray

Executed regression · cgls matches qr on small problem

Full file, imports & fixtures ↗

def test_cgls_matches_qr_on_small_problem() -> None:
    from omnibias.jax.optim import (
        gauss_newton_direction_cgls,
        lstsq_gauss_newton_direction,
    )

    rng = np.random.RandomState(11)
    a = jnp.asarray(rng.randn(20, 5))
    true = jnp.asarray(rng.randn(5))
    b = a @ true

    def residual_fn(p: Array) -> Array:
        return a @ p - b

    p0 = jnp.zeros(5)
    jac = a
    r0 = residual_fn(p0)
    d_qr = lstsq_gauss_newton_direction(jac, r0, 1e-8)
    d_cgls = gauss_newton_direction_cgls(residual_fn, p0, 1e-8, cgls_tol=1e-14)
    assert jnp.allclose(d_qr, d_cgls, atol=1e-6)

python -m pytest "packages/omnibias-jax/tests/test_optim.py::test_cgls_matches_qr_on_small_problem" -q

Read the implementation

def gauss_newton_direction_cgls(
    residual_fn: ResidualFn,
    params: Array,
    damping: float,
    *,
    cgls_max_iter: int = 100,
    cgls_tol: float = 1e-8,
) -> Array:
    r"""Matrix-free LM Gauss-Newton direction via CGLS.

    Twin of :func:`omnibias.torch.optim.gauss_newton_direction_cgls`.
    """
    res, jt, jvec = _linearize_gn(residual_fn, params)
    return cgls(
        jvec,
        jt,
        -res,
        damp=float(jnp.sqrt(jnp.asarray(damping, dtype=res.dtype))),
        max_iter=cgls_max_iter,
        tol=cgls_tol,
    )

FUNCTION

gauss_newton_fisher

omnibias.jax.optim.gauss_newton_fisherSource ↗

gauss_newton_fisher(residual_fn: ResidualFn, params: Array) -> tuple[Array, Array]

Dense Gauss-Newton Fisher F = (1/N) J^T J and gradient g = (1/N) J^T r.

The closed-form Fisher metric of the least-squares objective 0.5 mean(r^2): with J = d r / d theta (one jax.jacrev), F is the Gauss-Newton (and, at a zero residual, exact) Hessian and g its gradient. Feed the pair to natural_gradient_direction for a Fisher-scoring / natural-gradient step (which on a residual linear in theta equals Newton and recovers the least-squares minimiser in one step). mathematically equivalent twin of omnibias.torch.optim.gauss_newton_fisher.

Returns (F, g) of shape ((P, P), (P,)).

Backend scope: this identifies the corresponding mathematical construction. It does not guarantee bitwise-equal floating-point results across devices, compilers, dtypes, or backend versions. Validate the selected path with an appropriate absolute and relative tolerance.

Parameters

Name	Type & default	Contract
`residual_fn`positional-or-keyword	`ResidualFn`Required	See the operation contract and implementation below.
`params`positional-or-keyword	`Array`Required	See the operation contract and implementation below.

Returnstuple[Array, Array]

Explicit validation & failure cases

When res.ndim != 1, raises ValueError(f'residual_fn must return a 1-D vector, got shape {tuple(res.shape)}').

These are explicit source guards. Backend operations may also raise for invalid shapes, dtypes, or unsupported devices.

Executed regression · natural gradient step recovers regression in one step

Full file, imports & fixtures ↗

def test_natural_gradient_step_recovers_regression_in_one_step() -> None:
    """The Fisher-scoring parameter update lands on the least-squares solution in one step."""
    rng = np.random.RandomState(31)
    x = jnp.asarray(rng.randn(24, 5))
    y = jnp.asarray(rng.randn(24))

    def residual_fn(theta: Array) -> Array:
        return x @ theta - y

    theta0 = jnp.zeros(5)
    fisher, g = gauss_newton_fisher(residual_fn, theta0)
    theta1 = natural_gradient_step(theta0, g, fisher, learning_rate=1.0, damping=0.0)
    p_star, _, _, _ = jnp.linalg.lstsq(x, y, rcond=None)
    assert jnp.allclose(theta1, p_star, atol=1e-6)

python -m pytest "packages/omnibias-jax/tests/test_optim.py::test_natural_gradient_step_recovers_regression_in_one_step" -q

Regression example · gauss newton fisher bit parity

This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗

def test_gauss_newton_fisher_bit_parity(n: int, p: int) -> None:
    """The closed-form Gauss-Newton Fisher ``(F, g)`` agrees torch <-> jax on a shared residual."""
    rng = np.random.RandomState(n * 7 + p)
    x = rng.randn(n, p)
    y = rng.randn(n)
    theta0 = rng.randn(p)

    xt, yt = torch.tensor(x), torch.tensor(y)
    xj, yj = jnp.asarray(x), jnp.asarray(y)
    f_t, g_t = torch_gnf(lambda th: xt @ th - yt, torch.tensor(theta0))
    f_j, g_j = jax_gnf(lambda th: xj @ th - yj, jnp.asarray(theta0))
    assert np.allclose(f_t.numpy(), np.asarray(f_j), rtol=0, atol=1e-10)
    assert np.allclose(g_t.numpy(), np.asarray(g_j), rtol=0, atol=1e-10)

python -m pytest tests/test_optim_parity.py -k test_gauss_newton_fisher_bit_parity -q

Read the implementation

def gauss_newton_fisher(residual_fn: ResidualFn, params: Array) -> tuple[Array, Array]:
    r"""Dense Gauss-Newton Fisher ``F = (1/N) J^T J`` and gradient ``g = (1/N) J^T r``.

    The closed-form Fisher metric of the least-squares objective ``0.5 mean(r^2)``: with
    ``J = d r / d theta`` (one :func:`jax.jacrev`), ``F`` is the Gauss-Newton (and, at a zero
    residual, exact) Hessian and ``g`` its gradient. Feed the pair to
    :func:`natural_gradient_direction` for a Fisher-scoring / natural-gradient step (which on a
    residual linear in ``theta`` equals Newton and recovers the least-squares minimiser in one
    step). Bit-identical twin of :func:`omnibias.torch.optim.gauss_newton_fisher`.

    Returns ``(F, g)`` of shape ``((P, P), (P,))``.
    """
    res = residual_fn(params)
    if res.ndim != 1:
        raise ValueError(f"residual_fn must return a 1-D vector, got shape {tuple(res.shape)}")
    n_res = res.shape[0]
    jac = jax.jacrev(residual_fn)(params)  # (N, P)
    fisher = (jac.T @ jac) / n_res
    g = (jac.T @ res) / n_res
    return fisher, g

FUNCTION

gauss_newton_minimize

omnibias.jax.optim.gauss_newton_minimizeSource ↗

gauss_newton_minimize(residual_fn: ResidualFn, params: Array, *, steps: int, damping: float=0.001, **step_kwargs: Any) -> tuple[GaussNewtonState, list[float]]

Run gauss_newton_step steps times; return (state, loss_history).

Parameters

Name	Type & default	Contract
`residual_fn`positional-or-keyword	`ResidualFn`Required	See the operation contract and implementation below.
`params`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`steps`keyword-only	`int`Required	See the operation contract and implementation below.
`damping`keyword-only	`float`Default: 0.001	See the operation contract and implementation below.
`step_kwargs`variadic keyword	`Any`Required	See the operation contract and implementation below.

Returnstuple[GaussNewtonState, list[float]]

Executed regression · gauss newton beats adam on poisson pinn

Full file, imports & fixtures ↗

def test_gauss_newton_beats_adam_on_poisson_pinn() -> None:
    flat0, residual_fn = _poisson_residual_fn()
    _, history = gauss_newton_minimize(residual_fn, flat0, steps=40, damping=1e-2)
    gn_loss = history[-1]
    adam_loss = _adam_minimize(residual_fn, flat0, steps=800, lr=1e-3)
    assert gn_loss < 5e-4
    assert adam_loss > 2e-3
    assert gn_loss < adam_loss / 5.0

python -m pytest "packages/omnibias-jax/tests/test_optim.py::test_gauss_newton_beats_adam_on_poisson_pinn" -q

Regression example · gauss newton recovers nonlinear least squares

This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗

def test_gauss_newton_recovers_nonlinear_least_squares() -> None:
    t = jnp.linspace(0.0, 1.0, 24)
    true = jnp.array([2.0, -0.7])
    y = true[0] * jnp.exp(true[1] * t)

    def residual_fn(p: Array) -> Array:
        return p[0] * jnp.exp(p[1] * t) - y

    p0 = jnp.array([1.0, 0.0])
    state, history = gauss_newton_minimize(residual_fn, p0, steps=30, damping=1e-3)
    assert history[-1] < 1e-16
    assert jnp.allclose(state.params, true, atol=1e-6)
    assert all(history[i + 1] <= history[i] + 1e-18 for i in range(len(history) - 1))

python -m pytest packages/omnibias-jax/tests/test_optim.py -k test_gauss_newton_recovers_nonlinear_least_squares -q

Read the implementation

def gauss_newton_minimize(
    residual_fn: ResidualFn,
    params: Array,
    *,
    steps: int,
    damping: float = 1e-3,
    **step_kwargs: Any,
) -> tuple[GaussNewtonState, list[float]]:
    """Run :func:`gauss_newton_step` ``steps`` times; return ``(state, loss_history)``."""
    state = init_gauss_newton_state(params, damping=damping)
    history: list[float] = []
    for _ in range(steps):
        state = gauss_newton_step(residual_fn, state, **step_kwargs)
        history.append(state.loss)
    return state, history

FUNCTION

gauss_newton_step

omnibias.jax.optim.gauss_newton_stepSource ↗

gauss_newton_step(residual_fn: ResidualFn, state: GaussNewtonState, *, damping_increase: float=3.0, damping_decrease: float=0.5, min_damping: float=1e-12, max_damping: float=1000000000000.0, max_line_search: int=8) -> GaussNewtonState

One adaptive-damping Gauss-Newton (Levenberg-Marquardt) iteration.

Computes r and J = d r / d theta (one jax.jacrev), proposes the damped GN direction, and accepts it iff it decreases 1/2 mean(r^2); on success the damping is multiplied by damping_decrease, on failure by damping_increase (retried up to max_line_search times). The step is eager (data-dependent control flow); pass a jax.jit-compiled residual_fn for speed.

Parameters

Name	Type & default	Contract
`residual_fn`positional-or-keyword	`ResidualFn`Required	See the operation contract and implementation below.
`state`positional-or-keyword	`GaussNewtonState`Required	See the operation contract and implementation below.
`damping_increase`keyword-only	`float`Default: 3.0	See the operation contract and implementation below.
`damping_decrease`keyword-only	`float`Default: 0.5	See the operation contract and implementation below.
`min_damping`keyword-only	`float`Default: 1e-12	See the operation contract and implementation below.
`max_damping`keyword-only	`float`Default: 1000000000000.0	See the operation contract and implementation below.
`max_line_search`keyword-only	`int`Default: 8	See the operation contract and implementation below.

ReturnsGaussNewtonState

Executed regression · gauss newton beats adam on poisson pinn

Full file, imports & fixtures ↗

def test_gauss_newton_beats_adam_on_poisson_pinn() -> None:
    flat0, residual_fn = _poisson_residual_fn()
    _, history = gauss_newton_minimize(residual_fn, flat0, steps=40, damping=1e-2)
    gn_loss = history[-1]
    adam_loss = _adam_minimize(residual_fn, flat0, steps=800, lr=1e-3)
    assert gn_loss < 5e-4
    assert adam_loss > 2e-3
    assert gn_loss < adam_loss / 5.0

python -m pytest "packages/omnibias-jax/tests/test_optim.py::test_gauss_newton_beats_adam_on_poisson_pinn" -q

Read the implementation

def gauss_newton_step(
    residual_fn: ResidualFn,
    state: GaussNewtonState,
    *,
    damping_increase: float = 3.0,
    damping_decrease: float = 0.5,
    min_damping: float = 1e-12,
    max_damping: float = 1e12,
    max_line_search: int = 8,
) -> GaussNewtonState:
    r"""One adaptive-damping Gauss-Newton (Levenberg-Marquardt) iteration.

    Computes ``r`` and ``J = d r / d theta`` (one :func:`jax.jacrev`), proposes the
    damped GN direction, and accepts it iff it decreases ``1/2 mean(r^2)``; on success
    the damping is multiplied by ``damping_decrease``, on failure by ``damping_increase``
    (retried up to ``max_line_search`` times). The step is eager (data-dependent control
    flow); pass a ``jax.jit``-compiled ``residual_fn`` for speed.
    """
    params = state.params
    res = residual_fn(params)
    jac = jax.jacrev(residual_fn)(params)
    loss0 = _half_mean_sq(res)
    mu = state.damping
    for _ in range(max_line_search):
        delta = gauss_newton_direction(jac, res, mu)
        new_params = params + delta
        res_new = residual_fn(new_params)
        loss1 = _half_mean_sq(res_new)
        if jnp.isfinite(jnp.asarray(loss1)) and loss1 < loss0:
            new_mu = max(mu * damping_decrease, min_damping)
            return GaussNewtonState(new_params, new_mu, loss1, True, state.n_iter + 1)
        mu = min(mu * damping_increase, max_damping)
    return GaussNewtonState(params, mu, loss0, False, state.n_iter + 1)

FUNCTION

grad_norm_weights

omnibias.jax.optim.grad_norm_weightsSource ↗

grad_norm_weights(loss_fns: tuple[Callable[[Any], Array], ...], params: Any, prev_weights: Array, *, alpha: float=0.9, ref_index: int=0, eps: float=1e-12) -> Array

Self-adaptive loss weights equalising per-term gradient norms.

For each term L_k computes g_k = ||d L_k / d theta|| and forms the target weight lhat_k = ||g_ref|| / (||g_k|| + eps) (Wang-Teng-Perdikaris 2021 gradient balancing), then returns the EMA alpha * prev + (1 - alpha) * lhat. Applying the returned weights makes the weighted gradient norms lambda_k ||g_k|| all equal to the reference term's, curing the gradient-pathology stiffness of multi-term PINN losses.

Parameters

Name	Type & default	Contract
`loss_fns`positional-or-keyword	`tuple[Callable[[Any], Array], ...]`Required	See the operation contract and implementation below.
`params`positional-or-keyword	`Any`Required	See the operation contract and implementation below.
`prev_weights`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`alpha`keyword-only	`float`Default: 0.9	See the operation contract and implementation below.
`ref_index`keyword-only	`int`Default: 0	See the operation contract and implementation below.
`eps`keyword-only	`float`Default: 1e-12	See the operation contract and implementation below.

ReturnsArray

Explicit validation & failure cases

When not 0.0 <= alpha <= 1.0, raises ValueError(f'alpha must be in [0, 1], got {alpha}').
When not 0 <= ref_index < len(loss_fns), raises ValueError(f'ref_index {ref_index} out of range for {len(loss_fns)} terms').

These are explicit source guards. Backend operations may also raise for invalid shapes, dtypes, or unsupported devices.

Executed regression · grad norm weights validation

Full file, imports & fixtures ↗

def test_grad_norm_weights_validation() -> None:
    net = make_jet_mlp(1, 4, 1, depth=1, seed=0)

    def lf(n: object) -> Array:
        return jnp.mean(n.value(jnp.zeros((1, 1))) ** 2)  # type: ignore[attr-defined]

    with pytest.raises(ValueError):
        grad_norm_weights((lf,), net, jnp.ones(1), alpha=2.0)
    with pytest.raises(ValueError):
        grad_norm_weights((lf,), net, jnp.ones(1), ref_index=3)

python -m pytest "packages/omnibias-jax/tests/test_optim.py::test_grad_norm_weights_validation" -q

Regression example · grad norm weights equalises weighted norms

This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗

def test_grad_norm_weights_equalises_weighted_norms() -> None:
    net = make_jet_mlp(1, 16, 1, depth=2, seed=0)
    x_int = jnp.linspace(0.0, 1.0, 40)[1:-1][:, None]
    x_bc = jnp.array([[0.0], [1.0]])
    f_int = (-math.pi**2) * jnp.sin(math.pi * x_int).reshape(-1)

    def loss_pde(n: object) -> Array:
        _v, _g, h = n.value_grad_hessian(x_int)  # type: ignore[attr-defined]
        return jnp.mean((h[:, 0, 0, 0] - f_int) ** 2)

    def loss_bc(n: object) -> Array:
        return 1000.0 * jnp.mean(n.value(x_bc).reshape(-1) ** 2)  # type: ignore[attr-defined]

    from jax.flatten_util import ravel_pytree

    def gnorm(fn: Callable[[object], Array]) -> float:
        flat, _ = ravel_pytree(jax.grad(fn)(net))
        return float(jnp.linalg.norm(flat))

    weights = grad_norm_weights((loss_pde, loss_bc), net, jnp.ones(2), alpha=0.0, ref_index=0)
    weighted = jnp.array([weights[0] * gnorm(loss_pde), weights[1] * gnorm(loss_bc)])
    assert jnp.allclose(weighted, weighted[0], rtol=1e-6)
    assert float(weights[0]) == pytest.approx(1.0, abs=1e-9)

python -m pytest packages/omnibias-jax/tests/test_optim.py -k test_grad_norm_weights_equalises_weighted_norms -q

Read the implementation

def grad_norm_weights(
    loss_fns: tuple[Callable[[Any], Array], ...],
    params: Any,
    prev_weights: Array,
    *,
    alpha: float = 0.9,
    ref_index: int = 0,
    eps: float = 1e-12,
) -> Array:
    r"""Self-adaptive loss weights equalising per-term gradient norms.

    For each term ``L_k`` computes ``g_k = ||d L_k / d theta||`` and forms the target
    weight ``lhat_k = ||g_ref|| / (||g_k|| + eps)`` (Wang-Teng-Perdikaris 2021 gradient
    balancing), then returns the EMA ``alpha * prev + (1 - alpha) * lhat``. Applying the
    returned weights makes the weighted gradient norms ``lambda_k ||g_k||`` all equal to
    the reference term's, curing the gradient-pathology stiffness of multi-term PINN
    losses.
    """
    if not 0.0 <= alpha <= 1.0:
        raise ValueError(f"alpha must be in [0, 1], got {alpha}")
    if not 0 <= ref_index < len(loss_fns):
        raise ValueError(f"ref_index {ref_index} out of range for {len(loss_fns)} terms")
    norms = []
    for fn in loss_fns:
        grad = jax.grad(fn)(params)
        flat, _ = ravel_pytree(grad)
        norms.append(jnp.linalg.norm(flat))
    norm_vec = jnp.stack(norms)
    target = norm_vec[ref_index] / (norm_vec + eps)
    out: Array = alpha * prev_weights + (1.0 - alpha) * target
    return out

FUNCTION

homotopy_gauss_newton_minimize

omnibias.jax.optim.homotopy_gauss_newton_minimizeSource ↗

homotopy_gauss_newton_minimize(residual_fn: HomotopyResidualFn, params0: Array, *, config: HomotopyGNConfig | None=None) -> tuple[Array, dict[str, Any]]

Minimise 0.5 ||r(theta, t)||^2 at increasing coupling t.

residual_fn(theta, t) must already include every constraint that belongs in the Jacobian (hard gauge, anti-ghost). Extra Newton on a residual that omits the gauge will collapse amplitude.

Parameters

Name	Type & default	Contract
`residual_fn`positional-or-keyword	`HomotopyResidualFn`Required	See the operation contract and implementation below.
`params0`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`config`keyword-only	`HomotopyGNConfig \| None`Default: None	See the operation contract and implementation below.

Returnstuple[Array, dict[str, Any]]

Executed regression · homotopy gn tracks linear coupling

Full file, imports & fixtures ↗

def test_homotopy_gn_tracks_linear_coupling() -> None:
    from omnibias.jax.optim import HomotopyGNConfig, homotopy_gauss_newton_minimize

    def residual_fn(p: Array, t: float) -> Array:
        return p - (1.0 + t)

    p1, info = homotopy_gauss_newton_minimize(
        residual_fn,
        jnp.asarray([0.0]),
        config=HomotopyGNConfig(
            stages=(0.0, 0.5, 1.0),
            steps_per_stage=4,
            method="martens_grosse",
            damping=1e-8,
        ),
    )
    assert float(jnp.abs(p1[0] - 2.0)) < 1e-8
    assert info["stages"][-1]["t"] == 1.0

python -m pytest "packages/omnibias-jax/tests/test_optim.py::test_homotopy_gn_tracks_linear_coupling" -q

Read the implementation

def homotopy_gauss_newton_minimize(
    residual_fn: HomotopyResidualFn,
    params0: Array,
    *,
    config: HomotopyGNConfig | None = None,
) -> tuple[Array, dict[str, Any]]:
    """Minimise ``0.5 ||r(theta, t)||^2`` at increasing coupling ``t``.

    ``residual_fn(theta, t)`` must already include every constraint that
    belongs in the Jacobian (hard gauge, anti-ghost). Extra Newton on a
    residual that omits the gauge will collapse amplitude.
    """
    cfg = HomotopyGNConfig() if config is None else config
    params = params0
    history: list[dict[str, float]] = []
    for t in cfg.stages:

        def r_t(vec: Array, t_frozen: float = float(t)) -> Array:
            return residual_fn(vec, t_frozen)

        if cfg.method == "cubic":
            cubic_cfg = CubicRegularizedGNConfig(
                steps=int(cfg.steps_per_stage),
                sigma=float(cfg.cubic_sigma),
                krylov_dim=cfg.krylov_dim,
            )
            params, losses = cubic_regularized_gauss_newton_minimize(
                r_t, params, config=cubic_cfg
            )
        else:
            mg_cfg = MartensGrosseGNConfig(
                steps=int(cfg.steps_per_stage),
                damping=float(cfg.damping),
                solver=cfg.solver,
            )
            params, losses = martens_grosse_gauss_newton_minimize(
                r_t, params, config=mg_cfg
            )
        r_end = residual_fn(params, float(t))
        history.append(
            {
                "t": float(t),
                "loss0": float(losses[0]),
                "loss1": float(losses[-1]),
                "max_abs": float(jnp.max(jnp.abs(r_end))),
            }
        )
    return params, {"stages": history, "method": cfg.method}

FUNCTION

hvp

omnibias.jax.optim.hvpSource ↗

hvp(loss_fn: Callable[[Array], Array], params: Array, v: Array) -> Array

Hessian-vector product H v of a scalar loss_fn at params.

Twin of omnibias.torch.optim.hvp: the JVP of the gradient map in direction v. Exact through the closed-form sigma'' tower.

Parameters

Name	Type & default	Contract
`loss_fn`positional-or-keyword	`Callable[[Array], Array]`Required	See the operation contract and implementation below.
`params`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`v`positional-or-keyword	`Array`Required	See the operation contract and implementation below.

ReturnsArray

Read the implementation

def hvp(loss_fn: Callable[[Array], Array], params: Array, v: Array) -> Array:
    r"""Hessian-vector product ``H v`` of a scalar ``loss_fn`` at ``params``.

    Twin of :func:`omnibias.torch.optim.hvp`: the JVP of the gradient map
    in direction ``v``. Exact through the closed-form ``sigma''`` tower.
    """
    return cast(Array, jax.jvp(jax.grad(loss_fn), (params,), (v,))[1])

FUNCTION

init_gauss_newton_state

omnibias.jax.optim.init_gauss_newton_stateSource ↗

init_gauss_newton_state(params: Array, *, damping: float=0.001) -> GaussNewtonState

Seed a GaussNewtonState from a flat parameter vector.

Parameters

Name	Type & default	Contract
`params`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`damping`keyword-only	`float`Default: 0.001	See the operation contract and implementation below.

ReturnsGaussNewtonState

Explicit validation & failure cases

When damping <= 0.0, raises ValueError(f'damping must be > 0, got {damping}').

These are explicit source guards. Backend operations may also raise for invalid shapes, dtypes, or unsupported devices.

Executed regression · gauss newton beats adam on poisson pinn

Full file, imports & fixtures ↗

def test_gauss_newton_beats_adam_on_poisson_pinn() -> None:
    flat0, residual_fn = _poisson_residual_fn()
    _, history = gauss_newton_minimize(residual_fn, flat0, steps=40, damping=1e-2)
    gn_loss = history[-1]
    adam_loss = _adam_minimize(residual_fn, flat0, steps=800, lr=1e-3)
    assert gn_loss < 5e-4
    assert adam_loss > 2e-3
    assert gn_loss < adam_loss / 5.0

python -m pytest "packages/omnibias-jax/tests/test_optim.py::test_gauss_newton_beats_adam_on_poisson_pinn" -q

Read the implementation

def init_gauss_newton_state(params: Array, *, damping: float = 1e-3) -> GaussNewtonState:
    """Seed a :class:`GaussNewtonState` from a flat parameter vector."""
    if damping <= 0.0:
        raise ValueError(f"damping must be > 0, got {damping}")
    return GaussNewtonState(params=params, damping=float(damping), loss=float("inf"), accepted=False, n_iter=0)

FUNCTION

lanczos_tridiag

omnibias.jax.optim.lanczos_tridiagSource ↗

lanczos_tridiag(matvec: MatVec, b: Array, k: int, *, tol: float=1e-10) -> tuple[Array, Array]

k-step Lanczos on a symmetric operator with full reorthogonalisation.

Twin of omnibias.torch.optim.lanczos_tridiag. Returns (Q, T) with Q (n x m) orthonormal (first column b/||b||) and T = Q^T A Q symmetric tridiagonal (m <= min(k, n)).

Parameters

Name	Type & default	Contract
`matvec`positional-or-keyword	`MatVec`Required	See the operation contract and implementation below.
`b`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`k`positional-or-keyword	`int`Required	See the operation contract and implementation below.
`tol`keyword-only	`float`Default: 1e-10	See the operation contract and implementation below.

Returnstuple[Array, Array]

Executed regression · cubic gn default krylov is full dimension

Full file, imports & fixtures ↗

def test_cubic_gn_default_krylov_is_full_dimension() -> None:
    rng = np.random.RandomState(2)
    mat = jnp.asarray(rng.randn(40, 18))
    b = jnp.asarray(rng.randn(40))

    def residual_fn(p: Array) -> Array:
        return mat @ p - b

    p_star, _, _, _ = jnp.linalg.lstsq(mat, b, rcond=None)
    p, _hist = cubic_regularized_gauss_newton_minimize(
        residual_fn,
        jnp.zeros(18),
        config=CubicRegularizedGNConfig(steps=20, sigma=1.0),
    )
    assert jnp.allclose(p, p_star, atol=1e-6)

python -m pytest "packages/omnibias-jax/tests/test_optim.py::test_cubic_gn_default_krylov_is_full_dimension" -q

Read the implementation

def lanczos_tridiag(
    matvec: MatVec, b: Array, k: int, *, tol: float = 1e-10
) -> tuple[Array, Array]:
    r"""``k``-step Lanczos on a symmetric operator with full reorthogonalisation.

    Twin of :func:`omnibias.torch.optim.lanczos_tridiag`. Returns ``(Q, T)`` with
    ``Q`` (``n x m``) orthonormal (first column ``b/||b||``) and ``T = Q^T A Q``
    symmetric tridiagonal (``m <= min(k, n)``).
    """
    n = int(b.shape[0])
    m_max = min(int(k), n)
    beta0 = float(jnp.linalg.norm(b))
    if beta0 == 0.0:
        q = jnp.zeros_like(b).at[0].set(1.0)
    else:
        q = b / beta0
    qs: list[Array] = [q]
    alphas: list[Array] = []
    betas: list[Array] = []
    q_prev = jnp.zeros_like(b)
    beta_prev = jnp.zeros((), dtype=b.dtype)
    for _j in range(m_max):
        w = matvec(qs[-1])
        alpha = jnp.vdot(w, qs[-1])
        alphas.append(alpha)
        w = w - alpha * qs[-1] - beta_prev * q_prev
        for qi in qs:
            w = w - jnp.vdot(w, qi) * qi
        beta = jnp.linalg.norm(w)
        if float(beta) <= tol:
            break
        betas.append(beta)
        q_prev = qs[-1]
        beta_prev = beta
        qs.append(w / beta)
    m = len(alphas)
    q_basis = jnp.stack(qs[:m], axis=1)
    tri = jnp.diag(jnp.stack(alphas))
    if m > 1:
        off = jnp.stack(betas[: m - 1])
        tri = tri + jnp.diag(off, 1) + jnp.diag(off, -1)
    return q_basis, tri

FUNCTION

linearized_linf_direction

omnibias.jax.optim.linearized_linf_directionSource ↗

linearized_linf_direction(jac: Array, res: Array, *, box: float | None=None, irls_iters: int=16) -> Array

Minimax (Chebyshev) step on the linearized residual res + jac @ c.

One column is exact: bisection on the epigraph radius t with interval intersection. Several columns use Lawson IRLS (reweighted least squares), which is the L^∞ successor of gauss_newton_direction when L2 Newton raises max|r|. Optional box clips the step.

Parameters

Name	Type & default	Contract
`jac`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`res`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`box`keyword-only	`float \| None`Default: None	See the operation contract and implementation below.
`irls_iters`keyword-only	`int`Default: 16	See the operation contract and implementation below.

ReturnsArray

Explicit validation & failure cases

When jac.ndim != 2, raises ValueError(f'jac must be 2-D (N, P), got shape {jac.shape}').
When res.shape != (n,), raises ValueError(f'res must have shape ({n},) matching jac rows, got {res.shape}').

These are explicit source guards. Backend operations may also raise for invalid shapes, dtypes, or unsupported devices.

Executed regression · linearized linf direction is exact on one column

Full file, imports & fixtures ↗

def test_linearized_linf_direction_is_exact_on_one_column() -> None:
    from omnibias.jax.optim import linearized_linf_direction

    jac = jnp.ones((3, 1))
    res = jnp.asarray([1.0, -0.2, 0.4])
    step = linearized_linf_direction(jac, res)
    # min_c max_i |r_i + c| is the mid-range shift.
    assert float(step[0]) == pytest.approx(-0.4, abs=1e-8)
    linf = float(jnp.max(jnp.abs(res + jac @ step)))
    assert linf == pytest.approx(0.6, abs=1e-8)
    boxed = linearized_linf_direction(jac, res, box=0.1)
    assert abs(float(boxed[0])) <= 0.1 + 1e-12

python -m pytest "packages/omnibias-jax/tests/test_optim.py::test_linearized_linf_direction_is_exact_on_one_column" -q

Regression example · martens grosse jvp recovers exact line search on linear

This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗

def test_martens_grosse_jvp_recovers_exact_line_search_on_linear() -> None:
    from omnibias.jax.optim import martens_grosse_combine

    rng = np.random.RandomState(13)
    a = jnp.asarray(rng.randn(16, 4))
    b = jnp.asarray(rng.randn(16))

    def residual_fn(p: Array) -> Array:
        return a @ p - b

    p0 = jnp.zeros(4)
    # Any GN direction for linear residual is exact once alpha is optimal.
    delta = jnp.linalg.lstsq(a, b, rcond=None)[0]
    step, _ = martens_grosse_combine(residual_fn, p0, delta, None)
    # Exact alpha for linear: min_a ||r0 + a J d||^2 with Jd = A d, r0 = -b
    jd = a @ delta
    r0 = residual_fn(p0)
    alpha_star = float(-jnp.vdot(jd, r0) / (jnp.vdot(jd, jd) + 1e-30))
    assert jnp.allclose(step, alpha_star * delta, atol=1e-12)
    # With alpha_star=1 when delta is the exact LS step from 0: A delta = b => r0=-b, Jd=b
    assert jnp.allclose(p0 + step, jnp.linalg.lstsq(a, b, rcond=None)[0], atol=1e-9)

python -m pytest packages/omnibias-jax/tests/test_optim.py -k test_martens_grosse_jvp_recovers_exact_line_search_on_linear -q

Read the implementation

def linearized_linf_direction(
    jac: Array,
    res: Array,
    *,
    box: float | None = None,
    irls_iters: int = 16,
) -> Array:
    """Minimax (Chebyshev) step on the linearized residual ``res + jac @ c``.

    One column is exact: bisection on the epigraph radius ``t`` with interval
    intersection. Several columns use Lawson IRLS (reweighted least squares),
    which is the L^∞ successor of :func:`gauss_newton_direction` when L2
    Newton raises ``max|r|``. Optional ``box`` clips the step.
    """
    if jac.ndim != 2:
        raise ValueError(f"jac must be 2-D (N, P), got shape {jac.shape}")
    n, p = jac.shape
    if res.shape != (n,):
        raise ValueError(f"res must have shape ({n},) matching jac rows, got {res.shape}")
    if p == 1:
        j = jac[:, 0]
        lo = jnp.array(0.0, dtype=res.dtype)
        hi = jnp.max(jnp.abs(res))
        c_star = jnp.array(0.0, dtype=res.dtype)

        def body(carry: tuple[Array, Array, Array], _: None) -> tuple[tuple[Array, Array, Array], None]:
            lo_, hi_, c_ = carry
            t = 0.5 * (lo_ + hi_)
            # For each row, c in an interval so |r + j c| <= t.
            # Skip near-zero j: those rows only constrain t >= |r|.
            nz = jnp.abs(j) > 1e-15
            left = jnp.where(j >= 0.0, (-t - res) / (j + 1e-30), (t - res) / (j + 1e-30))
            right = jnp.where(j >= 0.0, (t - res) / (j + 1e-30), (-t - res) / (j + 1e-30))
            left = jnp.where(nz, left, -jnp.inf)
            right = jnp.where(nz, right, jnp.inf)
            c_lo = jnp.max(left)
            c_hi = jnp.min(right)
            feasible = (c_lo <= c_hi) & (jnp.max(jnp.abs(res) * (1.0 - nz.astype(res.dtype))) <= t)
            c_mid = 0.5 * (c_lo + c_hi)
            if box is not None:
                c_mid = jnp.clip(c_mid, -float(box), float(box))
                feasible = feasible & (c_lo <= float(box)) & (c_hi >= -float(box))
            lo_n = jnp.where(feasible, lo_, t)
            hi_n = jnp.where(feasible, t, hi_)
            c_n = jnp.where(feasible, c_mid, c_)
            return (lo_n, hi_n, c_n), None

        (_, _, c_star), _ = jax.lax.scan(body, (lo, hi, c_star), xs=None, length=48)
        step = c_star.reshape((1,))
    else:
        w = jnp.ones_like(res)
        c = jnp.zeros((p,), dtype=res.dtype)

        def irls(carry: tuple[Array, Array], _xs: None) -> tuple[tuple[Array, Array], None]:
            w_, _c = carry
            sw = jnp.sqrt(w_)
            c_n, _, _, _ = jnp.linalg.lstsq(jac * sw[:, None], -res * sw, rcond=None)
            if box is not None:
                c_n = jnp.clip(c_n, -float(box), float(box))
            resid = res + jac @ c_n
            w_n = w_ * (jnp.abs(resid) + 1e-15)
            w_n = w_n / jnp.mean(w_n)
            return (w_n, c_n), None

        (_, c), _ = jax.lax.scan(irls, (w, c), xs=None, length=int(irls_iters))
        step = c
    return step

FUNCTION

linf_minimax_minimize

omnibias.jax.optim.linf_minimax_minimizeSource ↗

linf_minimax_minimize(residual_fn: ResidualFn, params0: Array, *, config: LinfMinimaxConfig | None=None) -> tuple[Array, Array]

Minimise max_i |r_i(params)| by safeguarded linearized-epigraph steps.

The **L-infinity sibling** of martens_grosse_gauss_newton_minimize: repeatedly calls linf_minimax_step, which never accepts a step that increases the true max|r|, so the returned history is non-increasing by construction (unlike the L2 loss history of the Gauss-Newton drivers, which can rise between accepted steps only through momentum -- here there is no such channel). L2 Gauss-Newton and this L-infinity driver generally converge to *different* optima whenever a residual component is a genuine outlier; this function does not claim either is "better" in general, only that it minimises max|r| rather than 0.5 sum r_i^2.

This is a generic, independent comparison utility for any residual_fn; it makes no claim about any specific model, benchmark, or certified result.

Returns

(params, max_abs_history) where max_abs_history[k] is max|r| *before* step k and the final entry is max|r| at the returned params -- the same before/after history convention as martens_grosse_gauss_newton_minimize.

Parameters

Name	Type & default	Contract
`residual_fn`positional-or-keyword	`ResidualFn`Required	See the operation contract and implementation below.
`params0`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`config`keyword-only	`LinfMinimaxConfig \| None`Default: None	See the operation contract and implementation below.

Returnstuple[Array, Array]

Executed regression · linf minimax minimize reaches exact minimax on toy problem

Full file, imports & fixtures ↗

def test_linf_minimax_minimize_reaches_exact_minimax_on_toy_problem() -> None:
    from omnibias.jax.optim import LinfMinimaxConfig, linf_minimax_minimize

    c = jnp.asarray([1.0, -0.2, 0.4])

    def residual_fn(theta: Array) -> Array:
        return theta[0] + c

    theta0 = jnp.asarray([0.0])
    theta1, history = linf_minimax_minimize(
        residual_fn, theta0, config=LinfMinimaxConfig(steps=5, box0=1.0)
    )
    assert float(theta1[0]) == pytest.approx(-0.4, abs=1e-8)
    assert float(history[-1]) == pytest.approx(0.6, abs=1e-8)
    # The safeguard makes the recorded max|r| history non-increasing.
    assert all(
        float(history[i + 1]) <= float(history[i]) + 1e-12 for i in range(len(history) - 1)
    )

python -m pytest "packages/omnibias-jax/tests/test_optim.py::test_linf_minimax_minimize_reaches_exact_minimax_on_toy_problem" -q

Regression example · gn vs linf disagree on outlier heavy toy problem

This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗

def test_gn_vs_linf_disagree_on_outlier_heavy_toy_problem() -> None:
    """Fixed, deterministic, CPU-cheap regression: on a toy problem with a
    single outlier-heavy residual component, the L-infinity trainer reaches a
    strictly lower true max|r| than the existing Gauss-Newton (L2) trainer,
    under a fixed configuration (no seed search, no tuning against a live
    run). This is a generic comparison, a toy comparison: the
    residual is a 5-point sigmoid fit, unrelated to any champion network."""
    from omnibias.jax.optim import (
        LinfMinimaxConfig,
        MartensGrosseGNConfig,
        linf_minimax_minimize,
        martens_grosse_gauss_newton_minimize,
    )

    theta_gn, _losses_gn = martens_grosse_gauss_newton_minimize(
        _toy_outlier_residual,
        _TOY_SEED_THETA0,
        config=MartensGrosseGNConfig(steps=_TOY_GN_STEPS, damping=1e-3, solver="qr"),
    )
    maxr_gn = float(jnp.max(jnp.abs(_toy_outlier_residual(theta_gn))))

    theta_linf, history_linf = linf_minimax_minimize(
        _toy_outlier_residual,
        _TOY_SEED_THETA0,
        config=LinfMinimaxConfig(steps=_TOY_LINF_STEPS, box0=1.0),
    )
    maxr_linf = float(history_linf[-1])

    # Named absolute gates (literal constants, fixed ahead of time -- see the
    # closed-form optima recorded above the fixture).
    assert maxr_gn >= _TOY_GN_MAXR_GATE, f"GN max|r|={maxr_gn} below expected L2 basin"
    assert maxr_linf <= _TOY_LINF_MAXR_GATE, f"Linf max|r|={maxr_linf} above expected gate"
    # The headline comparison this test exists to make.
    assert maxr_linf < maxr_gn, (
        f"Linf trainer (max|r|={maxr_linf}) must beat GN (max|r|={maxr_gn}) "
        "on this outlier-heavy toy problem"
    )
    # Sanity: the two trainers land on genuinely different optima (not a tie).
    assert abs(float(theta_gn[0]) - float(theta_linf[0])) > 0.5

python -m pytest packages/omnibias-jax/tests/test_optim.py -k test_gn_vs_linf_disagree_on_outlier_heavy_toy_problem -q

Read the implementation

def linf_minimax_minimize(
    residual_fn: ResidualFn,
    params0: Array,
    *,
    config: LinfMinimaxConfig | None = None,
) -> tuple[Array, Array]:
    r"""Minimise ``max_i |r_i(params)|`` by safeguarded linearized-epigraph steps.

    The **L-infinity sibling** of :func:`martens_grosse_gauss_newton_minimize`:
    repeatedly calls :func:`linf_minimax_step`, which never accepts a step that
    increases the true ``max|r|``, so the returned history is non-increasing by
    construction (unlike the L2 loss history of the Gauss-Newton drivers, which can
    rise between accepted steps only through momentum -- here there is no such
    channel). L2 Gauss-Newton and this L-infinity driver generally converge to
    *different* optima whenever a residual component is a genuine outlier; this
    function does not claim either is "better" in general, only that it minimises
    ``max|r|`` rather than ``0.5 sum r_i^2``.

    This is a generic, independent comparison utility for any ``residual_fn``; it
    makes no claim about any specific model, benchmark, or certified result.

    Returns
    -------
    ``(params, max_abs_history)`` where ``max_abs_history[k]`` is ``max|r|``
    *before* step ``k`` and the final entry is ``max|r|`` at the returned
    ``params`` -- the same before/after history convention as
    :func:`martens_grosse_gauss_newton_minimize`.
    """
    cfg = LinfMinimaxConfig() if config is None else config
    params = params0
    box = float(cfg.box0)
    history: list[float] = []
    for _ in range(int(cfg.steps)):
        res0 = residual_fn(params)
        history.append(float(jnp.max(jnp.abs(res0))))
        params, box, _accepted, _max1 = linf_minimax_step(
            residual_fn,
            params,
            box,
            irls_iters=int(cfg.irls_iters),
            max_line_search=int(cfg.max_line_search),
            box_decrease=float(cfg.box_decrease),
            box_increase=float(cfg.box_increase),
            min_box=float(cfg.min_box),
            max_box=float(cfg.max_box),
            accept_tol=float(cfg.accept_tol),
        )
    history.append(float(jnp.max(jnp.abs(residual_fn(params)))))
    return params, jnp.asarray(history, dtype=jnp.float64)

FUNCTION

linf_minimax_step

omnibias.jax.optim.linf_minimax_stepSource ↗

linf_minimax_step(residual_fn: ResidualFn, params: Array, box: float, *, irls_iters: int=16, max_line_search: int=20, box_decrease: float=0.5, box_increase: float=1.5, min_box: float=1e-10, max_box: float=1000000.0, accept_tol: float=0.0) -> tuple[Array, float, bool, float]

One safeguarded linearized-epigraph L-infinity (minimax) step.

The **L-infinity sibling** of gauss_newton_step: it proposes a step from linearized_linf_direction on the live Jacobian J = d r / d theta (exact, one jax.jacfwd, no finite differences) inside a shrinking trust-region box, and accepts the step **only if it does not increase the true (non-linearized) max|r|**. On rejection the box shrinks by box_decrease (down to min_box) and the *same* live (res, jac) linearization is reused -- only the trust region shrinks, matching the LM-style backtracking of gauss_newton_step. On acceptance the box grows by box_increase (up to max_box) for the caller's next call.

This function does not compute or claim any application-specific result: it is a generic minimax step for any residual_fn: theta -> r (any shapes).

Parameters

residual_fn: Maps a flat parameter vector to a 1-D residual vector, exactly the ResidualFn convention used throughout this module. params: Current flat parameter vector. box: Trust-region radius for this call's linearized direction. irls_iters: Lawson-IRLS iterations forwarded to linearized_linf_direction (unused on the exact one-parameter path). max_line_search: Maximum number of box-shrink retries before giving up and rejecting. box_decrease, box_increase: Multiplicative shrink / grow factors for box. min_box, max_box: Hard floor / ceiling on box. accept_tol: Relative slack on the acceptance test; 0.0 requires the true max|r| to be non-increasing (up to float noise).

Returns

(new_params, new_box, accepted, max_abs_residual): on acceptance new_params is the improved point and max_abs_residual is its true max|r|; on rejection new_params is the unchanged params and max_abs_residual is max|r| at params.

Parameters

Name	Type & default	Contract
`residual_fn`positional-or-keyword	`ResidualFn`Required	Maps a flat parameter vector to a 1-D residual vector, exactly the `ResidualFn` convention used throughout this module.
`params`positional-or-keyword	`Array`Required	Current flat parameter vector.
`box`positional-or-keyword	`float`Required	Trust-region radius for this call's linearized direction.
`irls_iters`keyword-only	`int`Default: 16	Lawson-IRLS iterations forwarded to `linearized_linf_direction` (unused on the exact one-parameter path).
`max_line_search`keyword-only	`int`Default: 20	Maximum number of box-shrink retries before giving up and rejecting.
`box_decrease`keyword-only	`float`Default: 0.5	Multiplicative shrink / grow factors for `box`.
`box_increase`keyword-only	`float`Default: 1.5	See the operation contract and implementation below.
`min_box`keyword-only	`float`Default: 1e-10	Hard floor / ceiling on `box`.
`max_box`keyword-only	`float`Default: 1000000.0	See the operation contract and implementation below.
`accept_tol`keyword-only	`float`Default: 0.0	Relative slack on the acceptance test; `0.0` requires the true `max\|r\|` to be non-increasing (up to float noise).

Returnstuple[Array, float, bool, float]

Explicit validation & failure cases

When max_line_search < 1, raises ValueError(f'max_line_search must be >= 1, got {max_line_search}').

These are explicit source guards. Backend operations may also raise for invalid shapes, dtypes, or unsupported devices.

Executed regression · linf minimax minimize reaches exact minimax on toy problem

Full file, imports & fixtures ↗

def test_linf_minimax_minimize_reaches_exact_minimax_on_toy_problem() -> None:
    from omnibias.jax.optim import LinfMinimaxConfig, linf_minimax_minimize

    c = jnp.asarray([1.0, -0.2, 0.4])

    def residual_fn(theta: Array) -> Array:
        return theta[0] + c

    theta0 = jnp.asarray([0.0])
    theta1, history = linf_minimax_minimize(
        residual_fn, theta0, config=LinfMinimaxConfig(steps=5, box0=1.0)
    )
    assert float(theta1[0]) == pytest.approx(-0.4, abs=1e-8)
    assert float(history[-1]) == pytest.approx(0.6, abs=1e-8)
    # The safeguard makes the recorded max|r| history non-increasing.
    assert all(
        float(history[i + 1]) <= float(history[i]) + 1e-12 for i in range(len(history) - 1)
    )

python -m pytest "packages/omnibias-jax/tests/test_optim.py::test_linf_minimax_minimize_reaches_exact_minimax_on_toy_problem" -q

Regression example · linf minimax step reaches exact minimax in one step

This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗

def test_linf_minimax_step_reaches_exact_minimax_in_one_step() -> None:
    """On a linear residual, one safeguarded step already solves the exact
    minimax problem -- the same known-exact answer as
    ``test_linearized_linf_direction_is_exact_on_one_column``, but reached
    through the safeguarded step (not the raw direction primitive)."""
    from omnibias.jax.optim import linf_minimax_step

    c = jnp.asarray([1.0, -0.2, 0.4])

    def residual_fn(theta: Array) -> Array:
        return theta[0] + c

    theta0 = jnp.asarray([0.0])
    new_theta, new_box, accepted, max_abs = linf_minimax_step(residual_fn, theta0, box=1.0)
    assert accepted is True
    assert float(new_theta[0]) == pytest.approx(-0.4, abs=1e-8)
    assert max_abs == pytest.approx(0.6, abs=1e-8)
    assert new_box == pytest.approx(1.5, abs=1e-8)  # box_increase=1.5 on acceptance

python -m pytest packages/omnibias-jax/tests/test_optim.py -k test_linf_minimax_step_reaches_exact_minimax_in_one_step -q

Read the implementation

def linf_minimax_step(
    residual_fn: ResidualFn,
    params: Array,
    box: float,
    *,
    irls_iters: int = 16,
    max_line_search: int = 20,
    box_decrease: float = 0.5,
    box_increase: float = 1.5,
    min_box: float = 1e-10,
    max_box: float = 1e6,
    accept_tol: float = 0.0,
) -> tuple[Array, float, bool, float]:
    r"""One safeguarded linearized-epigraph L-infinity (minimax) step.

    The **L-infinity sibling** of :func:`gauss_newton_step`: it proposes a step from
    :func:`linearized_linf_direction` on the live Jacobian ``J = d r / d theta``
    (exact, one :func:`jax.jacfwd`, no finite differences) inside a shrinking
    trust-region ``box``, and accepts the step **only if it does not increase the
    true (non-linearized) ``max|r|``**. On rejection the box shrinks by ``box_decrease``
    (down to ``min_box``) and the *same* live ``(res, jac)`` linearization is
    reused -- only the trust region shrinks, matching the LM-style backtracking of
    :func:`gauss_newton_step`. On acceptance the box grows by ``box_increase`` (up
    to ``max_box``) for the caller's next call.

    This function does not compute or claim any application-specific result: it is
    a generic minimax step for any ``residual_fn: theta -> r`` (any shapes).

    Parameters
    ----------
    residual_fn:
        Maps a flat parameter vector to a 1-D residual vector, exactly the
        :data:`ResidualFn` convention used throughout this module.
    params:
        Current flat parameter vector.
    box:
        Trust-region radius for this call's linearized direction.
    irls_iters:
        Lawson-IRLS iterations forwarded to :func:`linearized_linf_direction`
        (unused on the exact one-parameter path).
    max_line_search:
        Maximum number of box-shrink retries before giving up and rejecting.
    box_decrease, box_increase:
        Multiplicative shrink / grow factors for ``box``.
    min_box, max_box:
        Hard floor / ceiling on ``box``.
    accept_tol:
        Relative slack on the acceptance test; ``0.0`` requires the true
        ``max|r|`` to be non-increasing (up to float noise).

    Returns
    -------
    ``(new_params, new_box, accepted, max_abs_residual)``: on acceptance
    ``new_params`` is the improved point and ``max_abs_residual`` is its true
    ``max|r|``; on rejection ``new_params`` is the unchanged ``params`` and
    ``max_abs_residual`` is ``max|r|`` at ``params``.
    """
    if max_line_search < 1:
        raise ValueError(f"max_line_search must be >= 1, got {max_line_search}")
    res0 = residual_fn(params)
    max0 = float(jnp.max(jnp.abs(res0)))
    jac = jax.jacfwd(residual_fn)(params)
    trial_box = float(box)
    for _ in range(int(max_line_search)):
        delta = linearized_linf_direction(jac, res0, box=trial_box, irls_iters=irls_iters)
        candidate = params + delta
        res1 = residual_fn(candidate)
        max1 = float(jnp.max(jnp.abs(res1)))
        if math.isfinite(max1) and max1 <= max0 * (1.0 + accept_tol) + 1e-15:
            new_box = min(trial_box * box_increase, max_box)
            return candidate, new_box, True, max1
        trial_box = max(trial_box * box_decrease, min_box)
    return params, trial_box, False, max0

FUNCTION

lstsq_gauss_newton_direction

omnibias.jax.optim.lstsq_gauss_newton_directionSource ↗

lstsq_gauss_newton_direction(jac: Array, res: Array, damping: float) -> Array

LM Gauss-Newton step via QR least-squares -- never squares the conditioning.

Solves the same damped problem as gauss_newton_direction as min_delta || [J; sqrt(mu) I] delta - [-r; 0] ||^2 via jnp.linalg.lstsq. Prefer this over forming J^T J whenever J is stiff (essentially all PINNs). mathematically equivalent twin of omnibias.torch.optim.lstsq_gauss_newton_direction.

Parameters

Name	Type & default	Contract
`jac`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`res`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`damping`positional-or-keyword	`float`Required	See the operation contract and implementation below.

ReturnsArray

Explicit validation & failure cases

When jac.ndim != 2, raises ValueError(f'jac must be 2-D (N, P), got shape {jac.shape}').
When res.shape != (n,), raises ValueError(f'res must have shape ({n},) matching jac rows, got {res.shape}').
When mu < 0.0, raises ValueError(f'damping must be >= 0, got {mu}').

These are explicit source guards. Backend operations may also raise for invalid shapes, dtypes, or unsupported devices.

Executed regression · lstsq qr matches dense on mild ls

Full file, imports & fixtures ↗

def test_lstsq_qr_matches_dense_on_mild_ls() -> None:
    from omnibias.jax.optim import lstsq_gauss_newton_direction

    rng = np.random.RandomState(7)
    jac = jnp.asarray(rng.randn(30, 4))
    res = jnp.asarray(rng.randn(30))
    mu = 1e-2
    d_qr = lstsq_gauss_newton_direction(jac, res, mu)
    d_dense = gauss_newton_direction(jac, res, mu)
    assert jnp.allclose(d_qr, d_dense, atol=1e-8, rtol=1e-7)

python -m pytest "packages/omnibias-jax/tests/test_optim.py::test_lstsq_qr_matches_dense_on_mild_ls" -q

Regression example · cgls matches qr on small problem

This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗

def test_cgls_matches_qr_on_small_problem() -> None:
    from omnibias.jax.optim import (
        gauss_newton_direction_cgls,
        lstsq_gauss_newton_direction,
    )

    rng = np.random.RandomState(11)
    a = jnp.asarray(rng.randn(20, 5))
    true = jnp.asarray(rng.randn(5))
    b = a @ true

    def residual_fn(p: Array) -> Array:
        return a @ p - b

    p0 = jnp.zeros(5)
    jac = a
    r0 = residual_fn(p0)
    d_qr = lstsq_gauss_newton_direction(jac, r0, 1e-8)
    d_cgls = gauss_newton_direction_cgls(residual_fn, p0, 1e-8, cgls_tol=1e-14)
    assert jnp.allclose(d_qr, d_cgls, atol=1e-6)

python -m pytest packages/omnibias-jax/tests/test_optim.py -k test_cgls_matches_qr_on_small_problem -q

Read the implementation

def lstsq_gauss_newton_direction(jac: Array, res: Array, damping: float) -> Array:
    r"""LM Gauss-Newton step via QR least-squares -- never squares the conditioning.

    Solves the same damped problem as :func:`gauss_newton_direction` as
    ``min_delta || [J; sqrt(mu) I] delta - [-r; 0] ||^2`` via :func:`jnp.linalg.lstsq`.
    Prefer this over forming ``J^T J`` whenever ``J`` is stiff (essentially all PINNs).
    Bit-identical twin of :func:`omnibias.torch.optim.lstsq_gauss_newton_direction`.
    """
    if jac.ndim != 2:
        raise ValueError(f"jac must be 2-D (N, P), got shape {jac.shape}")
    n, p = int(jac.shape[0]), int(jac.shape[1])
    if res.shape != (n,):
        raise ValueError(f"res must have shape ({n},) matching jac rows, got {res.shape}")
    mu = float(damping)
    if mu < 0.0:
        raise ValueError(f"damping must be >= 0, got {mu}")
    if mu == 0.0:
        a_aug = jac
        b_aug = -res
    else:
        eye = jnp.sqrt(jnp.asarray(mu, dtype=jac.dtype)) * jnp.eye(p, dtype=jac.dtype)
        a_aug = jnp.concatenate([jac, eye], axis=0)
        b_aug = jnp.concatenate([-res, jnp.zeros((p,), dtype=res.dtype)])
    sol, _, _, _ = jnp.linalg.lstsq(a_aug, b_aug, rcond=None)
    return sol.reshape(-1)

FUNCTION

make_residual_fn

omnibias.jax.optim.make_residual_fnSource ↗

make_residual_fn(build_residual: Callable[[Any], Array], params_pytree: Any) -> tuple[Array, ResidualFn]

Bridge a pytree model to a flat residual_fn.

Returns (flat0, residual_fn) where flat0 is the ravelled parameter vector and residual_fn(vec) = build_residual(unravel(vec)). build_residual receives a reconstructed pytree (e.g. a omnibias.jax.architectures.JetMLP) and returns the stacked residual vector.

Parameters

Name	Type & default	Contract
`build_residual`positional-or-keyword	`Callable[[Any], Array]`Required	See the operation contract and implementation below.
`params_pytree`positional-or-keyword	`Any`Required	See the operation contract and implementation below.

Returnstuple[Array, ResidualFn]

Executed regression · make residual fn roundtrip

Full file, imports & fixtures ↗

def test_make_residual_fn_roundtrip() -> None:
    flat0, residual_fn = _poisson_residual_fn()
    r = residual_fn(flat0)
    assert r.shape[0] == 40  # 38 interior + 2 boundary
    assert jnp.all(jnp.isfinite(r))

python -m pytest "packages/omnibias-jax/tests/test_optim.py::test_make_residual_fn_roundtrip" -q

Regression example · gauss newton solves poisson without boundary loss

This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗

def test_gauss_newton_solves_poisson_without_boundary_loss() -> None:
    net = make_jet_mlp(1, 24, 1, depth=2, seed=0)
    field = dirichlet_interval(net, 0.0, 1.0)
    x = jnp.linspace(0, 1, 40)[:, None]
    f = -(math.pi**2) * jnp.sin(math.pi * x[:, 0])

    def build_residual(fld: HardConstraintField) -> jnp.ndarray:
        return fld.hessian(x)[:, 0, 0, 0] - f

    flat0, rfn = make_residual_fn(build_residual, field)
    _, unravel = ravel_pytree(field)
    init_loss = 0.5 * float((rfn(flat0) ** 2).mean())
    state, history = gauss_newton_minimize(rfn, flat0, steps=40, damping=1e-2)
    assert history[-1] < init_loss * 1e-4

    field_opt = unravel(state.params)
    xx = jnp.linspace(0, 1, 101)[:, None]
    u = field_opt.value(xx)[:, 0]
    u_star = jnp.sin(math.pi * xx[:, 0])
    assert float(jnp.abs(u - u_star).max()) < 1e-3
    bc = field_opt.value(jnp.array([[0.0], [1.0]]))
    assert float(jnp.abs(bc).max()) < 1e-12

python -m pytest packages/omnibias-jax/tests/test_hardbc.py -k test_gauss_newton_solves_poisson_without_boundary_loss -q

Read the implementation

def make_residual_fn(
    build_residual: Callable[[Any], Array], params_pytree: Any
) -> tuple[Array, ResidualFn]:
    """Bridge a pytree model to a flat ``residual_fn``.

    Returns ``(flat0, residual_fn)`` where ``flat0`` is the ravelled parameter vector and
    ``residual_fn(vec) = build_residual(unravel(vec))``. ``build_residual`` receives a
    reconstructed pytree (e.g. a :class:`omnibias.jax.architectures.JetMLP`) and returns
    the stacked residual vector.
    """
    flat0, unravel = ravel_pytree(params_pytree)

    def residual_fn(vec: Array) -> Array:
        return build_residual(unravel(vec))

    return flat0, residual_fn

FUNCTION

martens_grosse_combine

omnibias.jax.optim.martens_grosse_combineSource ↗

martens_grosse_combine(residual_fn: ResidualFn, params: Array, delta_gn: Array, prev_delta: Array | None) -> tuple[Array, Array]

Martens–Grosse closed-form LR / momentum via **exact** JVPs.

Minimises the local quadratic model || r + J (alpha d + mu prev) ||^2 for (alpha, mu) using Jd = jvp(r, d) and Jp = jvp(r, prev) -- no finite-difference probes. Returns (step, momentum_state) where momentum_state is the accepted combined step (fed back as prev_delta on the next iteration).

Twin of omnibias.torch.optim.martens_grosse_combine.

Parameters

Name	Type & default	Contract
`residual_fn`positional-or-keyword	`ResidualFn`Required	See the operation contract and implementation below.
`params`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`delta_gn`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`prev_delta`positional-or-keyword	`Array \| None`Required	See the operation contract and implementation below.

Returnstuple[Array, Array]

Executed regression · homotopy gn tracks linear coupling

Full file, imports & fixtures ↗

def test_homotopy_gn_tracks_linear_coupling() -> None:
    from omnibias.jax.optim import HomotopyGNConfig, homotopy_gauss_newton_minimize

    def residual_fn(p: Array, t: float) -> Array:
        return p - (1.0 + t)

    p1, info = homotopy_gauss_newton_minimize(
        residual_fn,
        jnp.asarray([0.0]),
        config=HomotopyGNConfig(
            stages=(0.0, 0.5, 1.0),
            steps_per_stage=4,
            method="martens_grosse",
            damping=1e-8,
        ),
    )
    assert float(jnp.abs(p1[0] - 2.0)) < 1e-8
    assert info["stages"][-1]["t"] == 1.0

python -m pytest "packages/omnibias-jax/tests/test_optim.py::test_homotopy_gn_tracks_linear_coupling" -q

Regression example · martens grosse jvp recovers exact line search on linear

This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗

def test_martens_grosse_jvp_recovers_exact_line_search_on_linear() -> None:
    from omnibias.jax.optim import martens_grosse_combine

    rng = np.random.RandomState(13)
    a = jnp.asarray(rng.randn(16, 4))
    b = jnp.asarray(rng.randn(16))

    def residual_fn(p: Array) -> Array:
        return a @ p - b

    p0 = jnp.zeros(4)
    # Any GN direction for linear residual is exact once alpha is optimal.
    delta = jnp.linalg.lstsq(a, b, rcond=None)[0]
    step, _ = martens_grosse_combine(residual_fn, p0, delta, None)
    # Exact alpha for linear: min_a ||r0 + a J d||^2 with Jd = A d, r0 = -b
    jd = a @ delta
    r0 = residual_fn(p0)
    alpha_star = float(-jnp.vdot(jd, r0) / (jnp.vdot(jd, jd) + 1e-30))
    assert jnp.allclose(step, alpha_star * delta, atol=1e-12)
    # With alpha_star=1 when delta is the exact LS step from 0: A delta = b => r0=-b, Jd=b
    assert jnp.allclose(p0 + step, jnp.linalg.lstsq(a, b, rcond=None)[0], atol=1e-9)

python -m pytest packages/omnibias-jax/tests/test_optim.py -k test_martens_grosse_jvp_recovers_exact_line_search_on_linear -q

Read the implementation

def martens_grosse_combine(
    residual_fn: ResidualFn,
    params: Array,
    delta_gn: Array,
    prev_delta: Array | None,
) -> tuple[Array, Array]:
    r"""Martens–Grosse closed-form LR / momentum via **exact** JVPs.

    Minimises the local quadratic model
    ``|| r + J (alpha d + mu prev) ||^2`` for ``(alpha, mu)`` using
    ``Jd = jvp(r, d)`` and ``Jp = jvp(r, prev)`` -- no finite-difference probes.
    Returns ``(step, momentum_state)`` where ``momentum_state`` is the accepted
    combined step (fed back as ``prev_delta`` on the next iteration).

    Twin of :func:`omnibias.torch.optim.martens_grosse_combine`.
    """
    _, j_d = jax.jvp(residual_fn, (params,), (delta_gn,))
    r0 = residual_fn(params)
    if prev_delta is None:
        num = -jnp.vdot(j_d, r0)
        den = jnp.vdot(j_d, j_d) + 1e-30
        alpha = num / den
        step = alpha * delta_gn
        return step, step

    _, j_p = jax.jvp(residual_fn, (params,), (prev_delta,))
    a11 = jnp.vdot(j_d, j_d)
    a12 = jnp.vdot(j_d, j_p)
    a22 = jnp.vdot(j_p, j_p)
    b1 = -jnp.vdot(j_d, r0)
    b2 = -jnp.vdot(j_p, r0)
    det = a11 * a22 - a12 * a12
    det = jnp.where(jnp.abs(det) < 1e-30, 1e-30, det)
    alpha = (a22 * b1 - a12 * b2) / det
    mu = (-a12 * b1 + a11 * b2) / det
    step = alpha * delta_gn + mu * prev_delta
    return step, step

FUNCTION

martens_grosse_gauss_newton_minimize

omnibias.jax.optim.martens_grosse_gauss_newton_minimizeSource ↗

martens_grosse_gauss_newton_minimize(residual_fn: ResidualFn, params0: Array, *, config: MartensGrosseGNConfig | None=None) -> tuple[Array, Array]

Minimise 0.5 ||r(params)||^2 by damped GN + optional Martens–Grosse.

Default solver="qr" (non-squaring) and use_martens_grosse=True (exact JVP 2x2 LR/momentum). solver="dense" forms J^T J (legacy); "cgls" is matrix-free.

Parameters

Name	Type & default	Contract
`residual_fn`positional-or-keyword	`ResidualFn`Required	See the operation contract and implementation below.
`params0`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`config`keyword-only	`MartensGrosseGNConfig \| None`Default: None	See the operation contract and implementation below.

Returnstuple[Array, Array]

Explicit validation & failure cases

When cfg.solver not in ('dense', 'qr', 'cgls'), raises ValueError(f"solver must be 'dense', 'qr', or 'cgls', got {cfg.solver!r}").

These are explicit source guards. Backend operations may also raise for invalid shapes, dtypes, or unsupported devices.

Executed regression · homotopy gn tracks linear coupling

Full file, imports & fixtures ↗

def test_homotopy_gn_tracks_linear_coupling() -> None:
    from omnibias.jax.optim import HomotopyGNConfig, homotopy_gauss_newton_minimize

    def residual_fn(p: Array, t: float) -> Array:
        return p - (1.0 + t)

    p1, info = homotopy_gauss_newton_minimize(
        residual_fn,
        jnp.asarray([0.0]),
        config=HomotopyGNConfig(
            stages=(0.0, 0.5, 1.0),
            steps_per_stage=4,
            method="martens_grosse",
            damping=1e-8,
        ),
    )
    assert float(jnp.abs(p1[0] - 2.0)) < 1e-8
    assert info["stages"][-1]["t"] == 1.0

python -m pytest "packages/omnibias-jax/tests/test_optim.py::test_homotopy_gn_tracks_linear_coupling" -q

Regression example · martens grosse gauss newton minimize recovers nonlinear ls

This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗

def test_martens_grosse_gauss_newton_minimize_recovers_nonlinear_ls() -> None:
    from omnibias.jax.optim import (
        MartensGrosseGNConfig,
        martens_grosse_gauss_newton_minimize,
    )

    t = jnp.linspace(0.0, 1.0, 24)
    true = jnp.array([2.0, -0.7])
    y = true[0] * jnp.exp(true[1] * t)

    def residual_fn(p: Array) -> Array:
        return p[0] * jnp.exp(p[1] * t) - y

    p0 = jnp.array([1.0, 0.0])
    p1, hist = martens_grosse_gauss_newton_minimize(
        residual_fn,
        p0,
        config=MartensGrosseGNConfig(steps=40, damping=1e-3, solver="qr"),
    )
    assert float(hist[-1]) < 1e-16
    assert jnp.allclose(p1, true, atol=1e-6)

python -m pytest packages/omnibias-jax/tests/test_optim.py -k test_martens_grosse_gauss_newton_minimize_recovers_nonlinear_ls -q

Read the implementation

def martens_grosse_gauss_newton_minimize(
    residual_fn: ResidualFn,
    params0: Array,
    *,
    config: MartensGrosseGNConfig | None = None,
) -> tuple[Array, Array]:
    r"""Minimise ``0.5 ||r(params)||^2`` by damped GN + optional Martens–Grosse.

    Default ``solver="qr"`` (non-squaring) and ``use_martens_grosse=True`` (exact JVP
    2x2 LR/momentum). ``solver="dense"`` forms ``J^T J`` (legacy); ``"cgls"`` is
    matrix-free.
    """
    cfg = MartensGrosseGNConfig() if config is None else config
    if cfg.solver not in ("dense", "qr", "cgls"):
        raise ValueError(f"solver must be 'dense', 'qr', or 'cgls', got {cfg.solver!r}")
    params = params0
    gamma = float(cfg.damping)
    losses: list[float] = []
    prev_delta: Array | None = None

    for _ in range(int(cfg.steps)):
        r0 = residual_fn(params)
        loss0 = _half_sum_sq(r0)
        losses.append(loss0)

        if cfg.solver == "cgls":
            delta_gn = gauss_newton_direction_cgls(
                residual_fn,
                params,
                gamma,
                cgls_max_iter=cfg.cgls_max_iter,
                cgls_tol=cfg.cgls_tol,
            )
        else:
            jac = jax.jacfwd(residual_fn)(params)
            if cfg.solver == "qr":
                delta_gn = lstsq_gauss_newton_direction(jac, r0, gamma)
            else:
                delta_gn = gauss_newton_direction(jac, r0, gamma)

        if cfg.use_martens_grosse:
            step, prev_delta = martens_grosse_combine(
                residual_fn, params, delta_gn, prev_delta
            )
        else:
            step = delta_gn
            prev_delta = delta_gn

        candidate = params + step
        loss1 = _half_sum_sq(residual_fn(candidate))
        if loss1 <= loss0 * (1.0 + cfg.accept_tol) + 1e-15:
            params = candidate
            gamma = max(cfg.min_damping, gamma * cfg.damping_decrease)
        else:
            gamma = min(cfg.max_damping, gamma * cfg.damping_increase)
            prev_delta = None

    losses.append(_half_sum_sq(residual_fn(params)))
    return params, jnp.asarray(losses, dtype=jnp.float64)

FUNCTION

natural_gradient_direction

omnibias.jax.optim.natural_gradient_directionSource ↗

natural_gradient_direction(metric: Array, grad: Array, *, damping: float=0.001) -> Array

Natural-gradient direction delta = (M + damping I)^{-1} grad.

The metric-preconditioned counterpart of the raw gradient: for a Riemannian metric M on parameter space the *natural* gradient M^{-1} grad makes the descent step invariant to smooth reparametrisations. metric is a symmetric positive-(semi)definite (P, P) matrix and grad a (P,) vector; damping >= 0 (Tikhonov) keeps the solve well-posed when M is singular or ill-conditioned.

mathematically equivalent twin of omnibias.torch.optim.natural_gradient_direction (dense path) and identical in form to omnibias.curvature.natural_gradient.damped_solve. Pair it with the closed-form gauss_newton_fisher (Fisher scoring / Newton on a least-squares residual) or a geometry pullback metric (omnibias.geometry.jax.ops.pullback_metric).

Parameters

Name	Type & default	Contract
`metric`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`grad`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`damping`keyword-only	`float`Default: 0.001	See the operation contract and implementation below.

ReturnsArray

Explicit validation & failure cases

When metric.ndim != 2 or metric.shape[0] != metric.shape[1], raises ValueError(f'metric must be a square (P, P) matrix, got {tuple(metric.shape)}').
When grad.ndim != 1 or grad.shape[0] != metric.shape[0], raises ValueError(f'grad must be (P,) with P = {metric.shape[0]}, got {tuple(grad.shape)}').
When damping < 0.0, raises ValueError(f'damping must be >= 0, got {damping}').

These are explicit source guards. Backend operations may also raise for invalid shapes, dtypes, or unsupported devices.

Executed regression · natural gradient direction validation

Full file, imports & fixtures ↗

def test_natural_gradient_direction_validation() -> None:
    with pytest.raises(ValueError):
        natural_gradient_direction(jnp.zeros((3, 4)), jnp.zeros(3))
    with pytest.raises(ValueError):
        natural_gradient_direction(jnp.eye(3), jnp.zeros(4))
    with pytest.raises(ValueError):
        natural_gradient_direction(jnp.eye(3), jnp.zeros(3), damping=-1.0)

python -m pytest "packages/omnibias-jax/tests/test_optim.py::test_natural_gradient_direction_validation" -q

Regression example · gn direction solves linear least squares in one step

This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗

def test_gn_direction_solves_linear_least_squares_in_one_step() -> None:
    rng = np.random.RandomState(0)
    a = jnp.asarray(rng.randn(20, 5))
    b = jnp.asarray(rng.randn(20))
    p0 = jnp.zeros(5)
    delta = gauss_newton_direction(a, a @ p0 - b, damping=1e-10)
    p_star, _, _, _ = jnp.linalg.lstsq(a, b, rcond=None)
    assert jnp.allclose(p0 + delta, p_star, atol=1e-6)

python -m pytest packages/omnibias-jax/tests/test_optim.py -k test_gn_direction_solves_linear_least_squares_in_one_step -q

Read the implementation

def natural_gradient_direction(metric: Array, grad: Array, *, damping: float = 1e-3) -> Array:
    r"""Natural-gradient direction ``delta = (M + damping I)^{-1} grad``.

    The metric-preconditioned counterpart of the raw gradient: for a Riemannian metric ``M``
    on parameter space the *natural* gradient ``M^{-1} grad`` makes the descent step invariant
    to smooth reparametrisations. ``metric`` is a symmetric positive-(semi)definite ``(P, P)``
    matrix and ``grad`` a ``(P,)`` vector; ``damping >= 0`` (Tikhonov) keeps the solve
    well-posed when ``M`` is singular or ill-conditioned.

    Bit-identical twin of :func:`omnibias.torch.optim.natural_gradient_direction` (dense path)
    and identical in form to :func:`omnibias.curvature.natural_gradient.damped_solve`. Pair
    it with the closed-form :func:`gauss_newton_fisher` (Fisher scoring / Newton on a
    least-squares residual) or a geometry pullback metric
    (:func:`omnibias.geometry.jax.ops.pullback_metric`).
    """
    if metric.ndim != 2 or metric.shape[0] != metric.shape[1]:
        raise ValueError(f"metric must be a square (P, P) matrix, got {tuple(metric.shape)}")
    if grad.ndim != 1 or grad.shape[0] != metric.shape[0]:
        raise ValueError(
            f"grad must be (P,) with P = {metric.shape[0]}, got {tuple(grad.shape)}"
        )
    if damping < 0.0:
        raise ValueError(f"damping must be >= 0, got {damping}")
    p = metric.shape[0]
    damped = metric + jnp.asarray(damping, dtype=metric.dtype) * jnp.eye(p, dtype=metric.dtype)
    delta: Array = jnp.linalg.solve(damped, grad)
    return delta

FUNCTION

natural_gradient_step

omnibias.jax.optim.natural_gradient_stepSource ↗

natural_gradient_step(params: Array, grad: Array, metric: Array, *, learning_rate: float=1.0, damping: float=0.001) -> Array

Preconditioned parameter update theta - lr (M + damping I)^{-1} grad.

The metric-aware (natural-gradient / Riemannian) counterpart of vanilla gradient descent: descending along M^{-1} grad instead of grad makes the step invariant to smooth reparametrisations. params and grad are flat (P,) vectors and metric the (P, P) Riemannian metric (Fisher via gauss_newton_fisher, or a geometry pullback). Mirrors omnibias.curvature.natural_gradient.natural_gradient_step and is the functional twin of the torch omnibias.torch.optim.NaturalGradient step.

Parameters

Name	Type & default	Contract
`params`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`grad`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`metric`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`learning_rate`keyword-only	`float`Default: 1.0	See the operation contract and implementation below.
`damping`keyword-only	`float`Default: 0.001	See the operation contract and implementation below.

ReturnsArray

Executed regression · natural gradient step newton on quadratic

Full file, imports & fixtures ↗

def test_natural_gradient_step_newton_on_quadratic() -> None:
    """With the exact Hessian as metric the natural step is Newton: exact on a quadratic."""
    a = _spd(6, 32)
    b = jnp.asarray(np.random.RandomState(33).randn(6))
    theta_star = jnp.linalg.solve(a, b)
    theta = jnp.zeros(6)
    for _ in range(2):
        grad = a @ theta - b  # grad of 0.5 theta^T A theta - b^T theta
        theta = natural_gradient_step(theta, grad, a, learning_rate=1.0, damping=0.0)
    assert jnp.allclose(theta, theta_star, atol=1e-8)

python -m pytest "packages/omnibias-jax/tests/test_optim.py::test_natural_gradient_step_newton_on_quadratic" -q

Regression example · natural gradient step recovers regression in one step

This example belongs to the pinned upstream test suite. It may use fixtures and imports from its full file. Run it from that checkout with its test dependencies installed.

Full file, imports & fixtures ↗

def test_natural_gradient_step_recovers_regression_in_one_step() -> None:
    """The Fisher-scoring parameter update lands on the least-squares solution in one step."""
    rng = np.random.RandomState(31)
    x = jnp.asarray(rng.randn(24, 5))
    y = jnp.asarray(rng.randn(24))

    def residual_fn(theta: Array) -> Array:
        return x @ theta - y

    theta0 = jnp.zeros(5)
    fisher, g = gauss_newton_fisher(residual_fn, theta0)
    theta1 = natural_gradient_step(theta0, g, fisher, learning_rate=1.0, damping=0.0)
    p_star, _, _, _ = jnp.linalg.lstsq(x, y, rcond=None)
    assert jnp.allclose(theta1, p_star, atol=1e-6)

python -m pytest packages/omnibias-jax/tests/test_optim.py -k test_natural_gradient_step_recovers_regression_in_one_step -q

Read the implementation

def natural_gradient_step(
    params: Array,
    grad: Array,
    metric: Array,
    *,
    learning_rate: float = 1.0,
    damping: float = 1e-3,
) -> Array:
    r"""Preconditioned parameter update ``theta - lr (M + damping I)^{-1} grad``.

    The metric-aware (natural-gradient / Riemannian) counterpart of vanilla gradient descent:
    descending along ``M^{-1} grad`` instead of ``grad`` makes the step invariant to smooth
    reparametrisations. ``params`` and ``grad`` are flat ``(P,)`` vectors and ``metric`` the
    ``(P, P)`` Riemannian metric (Fisher via :func:`gauss_newton_fisher`, or a geometry
    pullback). Mirrors :func:`omnibias.curvature.natural_gradient.natural_gradient_step` and is
    the functional twin of the torch :class:`omnibias.torch.optim.NaturalGradient` step.
    """
    delta = natural_gradient_direction(metric, grad, damping=damping)
    out: Array = params - learning_rate * delta
    return out

FUNCTION

peak_weighted_residual

omnibias.jax.optim.peak_weighted_residualSource ↗

peak_weighted_residual(res: Array, power: float) -> Array

Emphasize residual peaks without changing the zero set.

r |r|^p is still zero iff r is. Used as a differentiable L^∞ proxy so Gauss–Newton cannot hide a large max-norm behind a small L2.

Parameters

Name	Type & default	Contract
`res`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`power`positional-or-keyword	`float`Required	See the operation contract and implementation below.

ReturnsArray

Executed regression · peak weighted residual preserves zeros

Full file, imports & fixtures ↗

def test_peak_weighted_residual_preserves_zeros() -> None:
    from omnibias.jax.optim import peak_weighted_residual

    r = jnp.asarray([0.0, -0.2, 0.5])
    w = peak_weighted_residual(r, 3.0)
    assert float(w[0]) == 0.0
    assert float(jnp.sign(w[1])) == float(jnp.sign(r[1]))

python -m pytest "packages/omnibias-jax/tests/test_optim.py::test_peak_weighted_residual_preserves_zeros" -q

Read the implementation

def peak_weighted_residual(res: Array, power: float) -> Array:
    """Emphasize residual peaks without changing the zero set.

    ``r |r|^p`` is still zero iff ``r`` is. Used as a differentiable L^∞ proxy
    so Gauss–Newton cannot hide a large max-norm behind a small L2.
    """
    if power == 0.0:
        return res
    return res * jnp.power(1.0 + jnp.abs(res), float(power))

← Package overview Runnable learning examples →

omnibias.jax.optim

Imports & aliases​

CubicRegularizedGNConfig​

Parameters

Fields & methods

GNSolver​

GaussNewtonState​

Parameters

Fields & methods

HomotopyGNConfig​

Parameters

Fields & methods

HomotopyResidualFn​

LinfMinimaxConfig​

Parameters

Fields & methods

MartensGrosseGNConfig​

Parameters

Fields & methods

MatVec​

ResidualFn​

cgls​

Parameters

champ_barrier_residual​

Parameters

cubic_regularized_gauss_newton_minimize​

Parameters

cubic_regularized_newton_step​

Parameters

gauss_newton_direction​

Parameters

gauss_newton_direction_cgls​

Parameters

gauss_newton_fisher​

Parameters

gauss_newton_minimize​

Parameters

gauss_newton_step​

Parameters

grad_norm_weights​

Parameters

homotopy_gauss_newton_minimize​

Parameters

hvp​

Parameters

init_gauss_newton_state​

Parameters

lanczos_tridiag​

Parameters

linearized_linf_direction​

Parameters

linf_minimax_minimize​

Parameters

linf_minimax_step​

Parameters

lstsq_gauss_newton_direction​

Parameters

make_residual_fn​

Parameters

martens_grosse_combine​

Parameters

martens_grosse_gauss_newton_minimize​

Parameters

natural_gradient_direction​

Parameters

natural_gradient_step​

Parameters

peak_weighted_residual​

Parameters

Imports & aliases

CubicRegularizedGNConfig

GNSolver

GaussNewtonState

HomotopyGNConfig

HomotopyResidualFn

LinfMinimaxConfig

MartensGrosseGNConfig

MatVec

ResidualFn

cgls

champ_barrier_residual

cubic_regularized_gauss_newton_minimize

cubic_regularized_newton_step

gauss_newton_direction

gauss_newton_direction_cgls

gauss_newton_fisher

gauss_newton_minimize

gauss_newton_step

grad_norm_weights

homotopy_gauss_newton_minimize

hvp

init_gauss_newton_state

lanczos_tridiag

linearized_linf_direction

linf_minimax_minimize

linf_minimax_step

lstsq_gauss_newton_direction

make_residual_fn

martens_grosse_combine

martens_grosse_gauss_newton_minimize

natural_gradient_direction

natural_gradient_step

peak_weighted_residual