omnibias.jax.weight_loss_jet

0.5.0rc1Release candidatePython >=3.10Apache-2.0

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

import omnibias.jax.weight_loss_jet

Module contract & mathematical conventions

Closed-form one-layer weight-space loss jets, JAX twin.

Host-side conversion into omnibias.core.weight_loss_jet. Bit-identical to the torch twin on the same floats. Do not jit the driver.

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
`WeightLossJetSpec`	`omnibias.core.WeightLossJetSpec`
`honesty_payload`	`omnibias.core.weight_loss_jet.honesty_payload`
`one_layer_param_count`	`omnibias.core.one_layer_param_count`
`pack_one_layer_params`	`omnibias.core.pack_one_layer_params`
`unpack_one_layer_params`	`omnibias.core.unpack_one_layer_params`

FUNCTION

one_layer_loss

omnibias.jax.one_layer_lossSource ↗

one_layer_loss(features: Array, targets: Array, params: Array, hidden: int, dim: int, activation: str='tanh') -> float

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
`features`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`targets`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`params`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`hidden`positional-or-keyword	`int`Required	See the operation contract and implementation below.
`dim`positional-or-keyword	`int`Required	See the operation contract and implementation below.
`activation`positional-or-keyword	`str`Default: 'tanh'	See the operation contract and implementation below.

Returnsfloat

Read the implementation

def one_layer_loss(
    features: Array,
    targets: Array,
    params: Array,
    hidden: int,
    dim: int,
    activation: str = "tanh",
) -> float:
    return core.one_layer_loss(
        _rows(features), _vec(targets), _vec(params), hidden, dim, activation
    )

Other public import paths

omnibias.jax.one_layer_loss
omnibias.jax.weight_loss_jet.one_layer_loss

FUNCTION

one_layer_loss_grad

omnibias.jax.one_layer_loss_gradSource ↗

one_layer_loss_grad(features: Array, targets: Array, params: Array, hidden: int, dim: int, activation: str='tanh') -> list[float]

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
`features`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`targets`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`params`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`hidden`positional-or-keyword	`int`Required	See the operation contract and implementation below.
`dim`positional-or-keyword	`int`Required	See the operation contract and implementation below.
`activation`positional-or-keyword	`str`Default: 'tanh'	See the operation contract and implementation below.

Returnslist[float]

Read the implementation

def one_layer_loss_grad(
    features: Array,
    targets: Array,
    params: Array,
    hidden: int,
    dim: int,
    activation: str = "tanh",
) -> list[float]:
    return core.one_layer_loss_grad(
        _rows(features), _vec(targets), _vec(params), hidden, dim, activation
    )

Other public import paths

omnibias.jax.one_layer_loss_grad
omnibias.jax.weight_loss_jet.one_layer_loss_grad

FUNCTION

one_layer_loss_hessian

omnibias.jax.one_layer_loss_hessianSource ↗

one_layer_loss_hessian(features: Array, targets: Array, params: Array, hidden: int, dim: int, activation: str='tanh', *, max_params: int=256) -> list[list[float]]

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
`features`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`targets`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`params`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`hidden`positional-or-keyword	`int`Required	See the operation contract and implementation below.
`dim`positional-or-keyword	`int`Required	See the operation contract and implementation below.
`activation`positional-or-keyword	`str`Default: 'tanh'	See the operation contract and implementation below.
`max_params`keyword-only	`int`Default: 256	See the operation contract and implementation below.

Returnslist[list[float]]

Read the implementation

def one_layer_loss_hessian(
    features: Array,
    targets: Array,
    params: Array,
    hidden: int,
    dim: int,
    activation: str = "tanh",
    *,
    max_params: int = 256,
) -> list[list[float]]:
    return core.one_layer_loss_hessian(
        _rows(features),
        _vec(targets),
        _vec(params),
        hidden,
        dim,
        activation,
        max_params=max_params,
    )

Other public import paths

omnibias.jax.one_layer_loss_hessian
omnibias.jax.weight_loss_jet.one_layer_loss_hessian

FUNCTION

one_layer_loss_jet

omnibias.jax.one_layer_loss_jetSource ↗

one_layer_loss_jet(features: Array, targets: Array, params: Array, direction: Array, *, order: int, hidden: int, dim: int, activation: str='tanh') -> list[float]

φ^(k)(0) via the shared core assembler.

Parameters

Name	Type & default	Contract
`features`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`targets`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`params`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`direction`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`order`keyword-only	`int`Required	See the operation contract and implementation below.
`hidden`keyword-only	`int`Required	See the operation contract and implementation below.
`dim`keyword-only	`int`Required	See the operation contract and implementation below.
`activation`keyword-only	`str`Default: 'tanh'	See the operation contract and implementation below.

Returnslist[float]

Regression example · torch jax core loss jet 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_torch_jax_core_loss_jet_parity() -> None:
    torch = pytest.importorskip("torch")
    jax = pytest.importorskip("jax")
    jax.config.update("jax_enable_x64", True)
    import jax.numpy as jnp
    from omnibias.jax.weight_loss_jet import one_layer_loss_jet as jax_jet
    from omnibias.torch.weight_loss_jet import one_layer_loss_jet as torch_jet

    torch.set_default_dtype(torch.float64)
    expected = core_jet(
        _XS,
        _YS,
        _THETA,
        _DIR,
        order=3,
        hidden=2,
        dim=2,
        activation="tanh",
    )
    t_out = torch_jet(
        torch.tensor(_XS, dtype=torch.float64),
        torch.tensor(_YS, dtype=torch.float64),
        torch.tensor(_THETA, dtype=torch.float64),
        torch.tensor(_DIR, dtype=torch.float64),
        order=3,
        hidden=2,
        dim=2,
        activation="tanh",
    )
    j_out = jax_jet(
        jnp.asarray(_XS, dtype=jnp.float64),
        jnp.asarray(_YS, dtype=jnp.float64),
        jnp.asarray(_THETA, dtype=jnp.float64),
        jnp.asarray(_DIR, dtype=jnp.float64),
        order=3,
        hidden=2,
        dim=2,
        activation="tanh",
    )
    assert t_out == expected
    assert j_out == expected

python -m pytest tests/test_weight_loss_jet_parity.py -k test_torch_jax_core_loss_jet_parity -q

Read the implementation

def one_layer_loss_jet(
    features: Array,
    targets: Array,
    params: Array,
    direction: Array,
    *,
    order: int,
    hidden: int,
    dim: int,
    activation: str = "tanh",
) -> list[float]:
    """``φ^(k)(0)`` via the shared core assembler."""
    return core.one_layer_loss_jet(
        _rows(features),
        _vec(targets),
        _vec(params),
        _vec(direction),
        order=order,
        hidden=hidden,
        dim=dim,
        activation=activation,
    )

Other public import paths

omnibias.jax.one_layer_loss_jet
omnibias.jax.weight_loss_jet.one_layer_loss_jet

FUNCTION

one_layer_newton_direction

omnibias.jax.one_layer_newton_directionSource ↗

one_layer_newton_direction(grad: Array | Sequence[float], hess: Array | Sequence[Sequence[float]], *, damping: float=1e-06) -> list[float]

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
`grad`positional-or-keyword	`Array \| Sequence[float]`Required	See the operation contract and implementation below.
`hess`positional-or-keyword	`Array \| Sequence[Sequence[float]]`Required	See the operation contract and implementation below.
`damping`keyword-only	`float`Default: 1e-06	See the operation contract and implementation below.

Returnslist[float]

Read the implementation

def one_layer_newton_direction(
    grad: Array | Sequence[float],
    hess: Array | Sequence[Sequence[float]],
    *,
    damping: float = 1e-6,
) -> list[float]:
    if isinstance(hess, Array):
        rows = [[float(v) for v in row] for row in hess]
    else:
        rows = [[float(v) for v in row] for row in hess]
    return core.one_layer_newton_direction(_vec(grad), rows, damping=damping)

Other public import paths

omnibias.jax.one_layer_newton_direction
omnibias.jax.weight_loss_jet.one_layer_newton_direction

← Package overview Runnable learning examples →

omnibias.jax.weight_loss_jet

Imports & aliases​

one_layer_loss​

Parameters

one_layer_loss_grad​

Parameters

one_layer_loss_hessian​

Parameters

one_layer_loss_jet​

Parameters

one_layer_newton_direction​

Parameters

Imports & aliases

one_layer_loss

one_layer_loss_grad

one_layer_loss_hessian

one_layer_loss_jet

one_layer_newton_direction