API REFERENCE/omnibias-fields

omnibias.fields.locus.jax

0.2.0rc1Release candidatePython >=3.10Apache-2.0

Structured scalar, vector, and tensor fields with derivative caches, calculus operators, and coordinate conventions.

import omnibias.fields.locus.jax

Module contract & mathematical conventions

JAX Newton projection on an equality locus (theory 01-09).

Forward is Gauss-Newton (host, matching the torch landing point in float64). The implicit-function-theorem VJP is newton_ift_vjp; differentiating through the unrolled loop is refused.

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 →

CLASS

LocusOutput

omnibias.fields.locus.jax.LocusOutputSource ↗

class LocusOutput

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.

The module contract describes the context for this operation. Its signature, explicit validation, return expressions, and implementation are recorded below.

Parameters

Name	Type & default	Contract
pointdataclass field	ArrayRequired	See the operation contract and implementation below.
branchdataclass field	ArrayRequired	See the operation contract and implementation below.
conditiondataclass field	ArrayRequired	See the operation contract and implementation below.
convergeddataclass field	ArrayRequired	See the operation contract and implementation below.

Fields & methods

pointfield

point: Array

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

View source ↗

branchfield

branch: Array

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

View source ↗

conditionfield

condition: Array

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

View source ↗

convergedfield

converged: Array

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

View source ↗

Read the implementation

class LocusOutput:
    point: Array
    branch: Array
    condition: Array
    converged: Array

FUNCTION

equality_locus_apply

omnibias.fields.locus.jax.equality_locus_applySource ↗

equality_locus_apply(sys: EqualitySystem, x0: Array, *, max_iter: int=20, tol: float=1e-12, require_transversal: bool=True) -> LocusOutput

The module contract describes the context for this operation. Its signature, explicit validation, return expressions, and implementation are recorded below.

Parameters

Name	Type & default	Contract
syspositional-or-keyword	EqualitySystemRequired	See the operation contract and implementation below.
x0positional-or-keyword	ArrayRequired	See the operation contract and implementation below.
max_iterkeyword-only	intDefault: 20	See the operation contract and implementation below.
tolkeyword-only	floatDefault: 1e-12	See the operation contract and implementation below.
require_transversalkeyword-only	boolDefault: True	See the operation contract and implementation below.

ReturnsLocusOutput

Executed regression · g6 torch jax parity

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_g6_torch_jax_parity() -> None:
    pytest.importorskip("jax")
    import jax
    import jax.numpy as jnp
    from omnibias.fields.locus.jax import equality_locus_apply

    jax.config.update("jax_enable_x64", True)
    torch.set_default_dtype(torch.float64)
    sys = _matched()
    x0 = torch.tensor([0.4, -0.1], dtype=torch.float64)
    out_t = EqualityLocusLayer(sys, dtype=torch.float64)(x0)
    out_j = equality_locus_apply(sys, jnp.asarray([0.4, -0.1], dtype=jnp.float64))
    assert out_t.point.detach().cpu().numpy() == pytest.approx(out_j.point.tolist(), rel=0, abs=0)

python -m pytest "packages/omnibias-fields/tests/test_locus_layer.py::test_g6_torch_jax_parity" -q

Read the implementation

def equality_locus_apply(
    sys: EqualitySystem,
    x0: Array,
    *,
    max_iter: int = 20,
    tol: float = 1e-12,
    require_transversal: bool = True,
) -> LocusOutput:
    result = newton_result(sys, x0, max_iter=max_iter, tol=tol)
    point = newton_project(sys, x0, max_iter=max_iter, tol=tol)
    sigma_min = float(result.condition)
    cond = 1.0 / max(sigma_min, 1e-30)
    ok = bool(result.converged)
    if require_transversal:
        ok = ok and bool(result.transversal)
    return LocusOutput(
        point=point,
        branch=jnp.asarray(result.branch, dtype=jnp.int64),
        condition=jnp.asarray(cond, dtype=x0.dtype),
        converged=jnp.asarray(ok),
    )

Related APIsomnibias.fields.locus.jax.LocusOutput↗omnibias.fields.locus.jax.newton_project↗omnibias.fields.locus.jax.newton_result↗

FUNCTION

jacobian

omnibias.fields.locus.jax.jacobianSource ↗

jacobian(sys: EqualitySystem, x: Array) -> Array

The module contract describes the context for this operation. Its signature, explicit validation, return expressions, and implementation are recorded below.

Parameters

Name	Type & default	Contract
syspositional-or-keyword	EqualitySystemRequired	See the operation contract and implementation below.
xpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.

ReturnsArray

Read the implementation

def jacobian(sys: EqualitySystem, x: Array) -> Array:
    return jnp.asarray(jacobian_core(sys, _as_tuple(x)), dtype=x.dtype)

Related APIsomnibias.core.locus.jacobian↗

FUNCTION

newton_ift_vjp

omnibias.fields.locus.jax.newton_ift_vjpSource ↗

newton_ift_vjp(sys: EqualitySystem, x_star: Array, weights: Array, g: Array) -> Array

dL/dc via dx*/dc = -DF^+ dF/dc. Memory is independent of iteration count.

Parameters

Name	Type & default	Contract
syspositional-or-keyword	EqualitySystemRequired	See the operation contract and implementation below.
x_starpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
weightspositional-or-keyword	ArrayRequired	See the operation contract and implementation below.
gpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.

ReturnsArray

Read the implementation

def newton_ift_vjp(
    sys: EqualitySystem,
    x_star: Array,
    weights: Array,
    g: Array,
) -> Array:
    """``dL/dc`` via ``dx*/dc = -DF^+ dF/dc``. Memory is independent of iteration count."""
    sys_w = _with_weights(sys, _as_tuple(weights))
    xs = _as_tuple(x_star)
    df = jacobian_core(sys_w, xs)
    pinv = pseudoinverse(df)
    dfdc = dF_d_weights(sys_w, xs)
    gv = _as_tuple(g)
    m = len(df)
    pinv_t_g = tuple(sum(pinv[k][j] * gv[k] for k in range(sys_w.dim)) for j in range(m))
    grad_w = []
    for j in range(sys_w.n_terms):
        acc = 0.0
        for i in range(m):
            acc -= dfdc[i][j] * pinv_t_g[i]
        grad_w.append(acc)
    return jnp.asarray(grad_w, dtype=weights.dtype)

Related APIsomnibias.core.locus.dF_d_weights↗omnibias.core.locus.jacobian↗omnibias.core.locus.pseudoinverse↗

FUNCTION

newton_project

omnibias.fields.locus.jax.newton_projectSource ↗

newton_project(sys: EqualitySystem, x0: Array, *, weights: Array | None=None, max_iter: int=20, tol: float=1e-12) -> Array

Gauss-Newton projection onto the locus.

Parameters

Name	Type & default	Contract
syspositional-or-keyword	EqualitySystemRequired	See the operation contract and implementation below.
x0positional-or-keyword	ArrayRequired	See the operation contract and implementation below.
weightskeyword-only	Array | NoneDefault: None	See the operation contract and implementation below.
max_iterkeyword-only	intDefault: 20	See the operation contract and implementation below.
tolkeyword-only	floatDefault: 1e-12	See the operation contract and implementation below.

ReturnsArray

Executed regression · g6 torch jax parity

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_g6_torch_jax_parity() -> None:
    pytest.importorskip("jax")
    import jax
    import jax.numpy as jnp
    from omnibias.fields.locus.jax import equality_locus_apply

    jax.config.update("jax_enable_x64", True)
    torch.set_default_dtype(torch.float64)
    sys = _matched()
    x0 = torch.tensor([0.4, -0.1], dtype=torch.float64)
    out_t = EqualityLocusLayer(sys, dtype=torch.float64)(x0)
    out_j = equality_locus_apply(sys, jnp.asarray([0.4, -0.1], dtype=jnp.float64))
    assert out_t.point.detach().cpu().numpy() == pytest.approx(out_j.point.tolist(), rel=0, abs=0)

python -m pytest "packages/omnibias-fields/tests/test_locus_layer.py::test_g6_torch_jax_parity" -q

Regression example · g6 torch jax newton 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_g6_torch_jax_newton_parity() -> None:
    pytest.importorskip("jax")
    import jax
    import jax.numpy as jnp
    from omnibias.fields.locus.jax import newton_project as newton_jax

    jax.config.update("jax_enable_x64", True)
    sys = _case_b()
    x0_t = torch.tensor([0.0, 0.20], dtype=torch.float64)
    x_t = newton_project(sys, x0_t)
    x_j = newton_jax(sys, jnp.asarray([0.0, 0.20], dtype=jnp.float64))
    assert x_t.detach().cpu().numpy() == pytest.approx(x_j.tolist(), rel=0, abs=0)
    core = newton_core(sys, (0.0, 0.20))
    assert x_t[1].item() == pytest.approx(core.point[1], rel=0, abs=0)
    assert math.isfinite(float(x_t[1].detach()))

python -m pytest packages/omnibias-fields/tests/test_locus.py -k test_g6_torch_jax_newton_parity -q

Read the implementation

def newton_project(
    sys: EqualitySystem,
    x0: Array,
    *,
    weights: Array | None = None,
    max_iter: int = 20,
    tol: float = 1e-12,
) -> Array:
    """Gauss-Newton projection onto the locus."""
    w = [t.weight for t in sys.terms] if weights is None else _as_tuple(weights)
    sys_w = _with_weights(sys, w)
    result = newton_project_core(sys_w, _as_tuple(x0), max_iter=max_iter, tol=tol)
    return jnp.asarray(result.point, dtype=x0.dtype)

Related APIsomnibias.core.locus.newton_project↗

FUNCTION

newton_result

omnibias.fields.locus.jax.newton_resultSource ↗

newton_result(sys: EqualitySystem, x0: Array, *, max_iter: int=20, tol: float=1e-12) -> NewtonResult

The module contract describes the context for this operation. Its signature, explicit validation, return expressions, and implementation are recorded below.

Parameters

Name	Type & default	Contract
syspositional-or-keyword	EqualitySystemRequired	See the operation contract and implementation below.
x0positional-or-keyword	ArrayRequired	See the operation contract and implementation below.
max_iterkeyword-only	intDefault: 20	See the operation contract and implementation below.
tolkeyword-only	floatDefault: 1e-12	See the operation contract and implementation below.

ReturnsNewtonResult

Executed regression · g6 torch jax parity

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_g6_torch_jax_parity() -> None:
    pytest.importorskip("jax")
    import jax
    import jax.numpy as jnp
    from omnibias.fields.locus.jax import equality_locus_apply

    jax.config.update("jax_enable_x64", True)
    torch.set_default_dtype(torch.float64)
    sys = _matched()
    x0 = torch.tensor([0.4, -0.1], dtype=torch.float64)
    out_t = EqualityLocusLayer(sys, dtype=torch.float64)(x0)
    out_j = equality_locus_apply(sys, jnp.asarray([0.4, -0.1], dtype=jnp.float64))
    assert out_t.point.detach().cpu().numpy() == pytest.approx(out_j.point.tolist(), rel=0, abs=0)

python -m pytest "packages/omnibias-fields/tests/test_locus_layer.py::test_g6_torch_jax_parity" -q

Read the implementation

def newton_result(
    sys: EqualitySystem, x0: Array, *, max_iter: int = 20, tol: float = 1e-12
) -> NewtonResult:
    return newton_project_core(sys, _as_tuple(x0), max_iter=max_iter, tol=tol)

Related APIsomnibias.core.locus.newton_project↗

FUNCTION

residual

omnibias.fields.locus.jax.residualSource ↗

residual(sys: EqualitySystem, x: Array) -> Array

The module contract describes the context for this operation. Its signature, explicit validation, return expressions, and implementation are recorded below.

Parameters

Name	Type & default	Contract
syspositional-or-keyword	EqualitySystemRequired	See the operation contract and implementation below.
xpositional-or-keyword	ArrayRequired	See the operation contract and implementation below.

ReturnsArray

Read the implementation

def residual(sys: EqualitySystem, x: Array) -> Array:
    return jnp.asarray(residual_core(sys, _as_tuple(x)), dtype=x.dtype)

Related APIsomnibias.core.locus.residual↗

← Package overviewRunnable learning examples →