API REFERENCE/omnibias-fields

omnibias.fields.jax.ops.high_order

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.jax.ops.high_order

Module contract & mathematical conventions

JAX high-order ops (mirrors omnibias.fields.torch.ops.high_order).

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 →

FUNCTION

biharmonic

omnibias.fields.jax.ops.biharmonicSource ↗

biharmonic(state: FieldState, name: str) -> 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
statepositional-or-keyword	FieldStateRequired	See the operation contract and implementation below.
namepositional-or-keyword	strRequired	See the operation contract and implementation below.

ReturnsArray

Read the implementation

def biharmonic(state: FieldState, name: str) -> Array:
    if _is_one_layer(state):
        sigma_4 = _sigma_of_order(state, 4)
        return state.field.polylaplacian(sigma_4, name, k=2)
    if _is_spectral(state):
        return state.field.biharmonic(state, name)  # type: ignore[attr-defined]
    if _is_chebyshev(state):
        return state.field.biharmonic(state, name)  # type: ignore[attr-defined]
    if _is_jet_mlp(state):
        return state.field.biharmonic(state, name)  # type: ignore[attr-defined]
    raise NotImplementedError(
        f"biharmonic not implemented for field type {type(state.field).__name__}"
    )

Other public import paths

• omnibias.fields.jax.ops.biharmonic
• omnibias.fields.jax.ops.high_order.biharmonic

FUNCTION

gradient_of_derivative

omnibias.fields.jax.ops.gradient_of_derivativeSource ↗

gradient_of_derivative(state: FieldState, name: str, *, axis: int | str) -> 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
statepositional-or-keyword	FieldStateRequired	See the operation contract and implementation below.
namepositional-or-keyword	strRequired	See the operation contract and implementation below.
axiskeyword-only	int | strRequired	See the operation contract and implementation below.

ReturnsArray

Executed regression · sobolev norm cross backend

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_sobolev_norm_cross_backend(make_torch_field, make_jax_field):  # type: ignore[no-untyped-def]
    rule = gauss_legendre(BOUNDS, (8, 8))
    ts = make_torch_field()(torch.as_tensor(rule.nodes, dtype=torch.float64))
    js = make_jax_field()(jnp.asarray(rule.nodes, dtype=jnp.float64))
    for k in (0, 1, 2):
        t = _np(ts.u.sobolev_norm(rule=rule, k=k))
        j = _np(js.u.sobolev_norm(rule=rule, k=k))
        assert np.allclose(t, j, rtol=1e-12, atol=1e-12), k

python -m pytest "packages/omnibias-fields/tests/test_integration_norms.py::test_sobolev_norm_cross_backend" -q

Read the implementation

def gradient_of_derivative(
    state: FieldState, name: str, *, axis: int | str,
) -> Array:
    a = _resolve_axis(state, axis)
    D = state.coordinate_spec.ndim
    cols = []
    for j in range(D):
        if a == j:
            cols.append(derivative(state, name, axis=a, order=2))
        else:
            from omnibias.fields.jax.ops.basic import mixed_partial
            cols.append(mixed_partial(state, name, (a, j), (1, 1)))
    return jnp.stack(cols, axis=-1)

Related APIsomnibias.fields.jax.ops.derivative↗omnibias.fields.jax.ops.mixed_partial↗

Other public import paths

• omnibias.fields.jax.ops.gradient_of_derivative
• omnibias.fields.jax.ops.high_order.gradient_of_derivative

FUNCTION

hessian

omnibias.fields.jax.ops.hessianSource ↗

hessian(state: FieldState, name: str, *, axes: tuple[int | str, ...] | None=None) -> 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
statepositional-or-keyword	FieldStateRequired	See the operation contract and implementation below.
namepositional-or-keyword	strRequired	See the operation contract and implementation below.
axeskeyword-only	tuple[int | str, ...] | NoneDefault: None	See the operation contract and implementation below.

ReturnsArray

Executed regression · sobolev norm cross backend

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_sobolev_norm_cross_backend(make_torch_field, make_jax_field):  # type: ignore[no-untyped-def]
    rule = gauss_legendre(BOUNDS, (8, 8))
    ts = make_torch_field()(torch.as_tensor(rule.nodes, dtype=torch.float64))
    js = make_jax_field()(jnp.asarray(rule.nodes, dtype=jnp.float64))
    for k in (0, 1, 2):
        t = _np(ts.u.sobolev_norm(rule=rule, k=k))
        j = _np(js.u.sobolev_norm(rule=rule, k=k))
        assert np.allclose(t, j, rtol=1e-12, atol=1e-12), k

python -m pytest "packages/omnibias-fields/tests/test_integration_norms.py::test_sobolev_norm_cross_backend" -q

Read the implementation

def hessian(
    state: FieldState,
    name: str,
    *,
    axes: tuple[int | str, ...] | None = None,
) -> Array:
    axis_idx = (
        tuple(range(state.coordinate_spec.ndim))
        if axes is None
        else tuple(_resolve_axis(state, a) for a in axes)
    )
    if _is_one_layer(state):
        sigma_pp = _sigma_of_order(state, 2)
        full = state.field.hessian_full(sigma_pp, name)
        if axis_idx == tuple(range(full.shape[-1])):
            return full
        idx = jnp.array(list(axis_idx))
        return full[..., idx, :][..., :, idx]
    if _is_jet_mlp(state):
        # The order-2 block of a single jet is the whole Hessian.
        full = state.field.hessian_full(state, name)  # type: ignore[attr-defined]
        if axis_idx == tuple(range(full.shape[-1])):
            return full
        idx = jnp.array(list(axis_idx))
        return full[..., idx, :][..., :, idx]
    rows = []
    idx = jnp.array(list(axis_idx))
    for i in axis_idx:
        gi = gradient_of_derivative(state, name, axis=i)
        rows.append(gi[..., idx])
    return jnp.stack(rows, axis=-2)

Related APIsomnibias.fields.jax.ops.gradient_of_derivative↗

Other public import paths

• omnibias.fields.jax.ops.hessian
• omnibias.fields.jax.ops.high_order.hessian

FUNCTION

jacobian

omnibias.fields.jax.ops.jacobianSource ↗

jacobian(state: FieldState, names: tuple[str, ...]) -> 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
statepositional-or-keyword	FieldStateRequired	See the operation contract and implementation below.
namespositional-or-keyword	tuple[str, ...]Required	See the operation contract and implementation below.

ReturnsArray

Read the implementation

def jacobian(state: FieldState, names: tuple[str, ...]) -> Array:
    if _is_one_layer(state):
        sigma_p = _sigma_of_order(state, 1)
        rows = [
            state.field.gradient_full(sigma_p, n) for n in names
        ]
        return jnp.stack(rows, axis=-2)
    if _is_jet_mlp(state):
        rows = [state.field.gradient_full(state, n) for n in names]  # type: ignore[attr-defined]
        return jnp.stack(rows, axis=-2)
    rows = []
    for n in names:
        rows.append(gradient(state, n, axes=tuple(range(state.coordinate_spec.ndim))))
    return jnp.stack(rows, axis=-2)

Related APIsomnibias.fields.jax.ops.gradient↗

Other public import paths

• omnibias.fields.jax.ops.jacobian
• omnibias.fields.jax.ops.high_order.jacobian

FUNCTION

polylaplacian

omnibias.fields.jax.ops.polylaplacianSource ↗

polylaplacian(state: FieldState, name: str, *, k: int) -> 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
statepositional-or-keyword	FieldStateRequired	See the operation contract and implementation below.
namepositional-or-keyword	strRequired	See the operation contract and implementation below.
kkeyword-only	intRequired	See the operation contract and implementation below.

ReturnsArray

Explicit validation & failure cases

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

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

Read the implementation

def polylaplacian(state: FieldState, name: str, *, k: int) -> Array:
    if k < 1:
        raise ValueError(f"polylaplacian k must be >= 1, got {k}")
    if k == 1:
        from omnibias.fields.jax.ops.basic import laplacian
        return laplacian(state, name)
    if _is_one_layer(state):
        sigma_2k = _sigma_of_order(state, 2 * k)
        return state.field.polylaplacian(sigma_2k, name, k=k)
    if _is_spectral(state):
        return state.field.polylaplacian(state, name, k=k)  # type: ignore[attr-defined]
    if _is_chebyshev(state):
        return state.field.polylaplacian(state, name, k=k)  # type: ignore[attr-defined]
    if _is_jet_mlp(state):
        # Multinomial expansion of (sum_i d_i^2)^k over one order-2k jet.
        return state.field.polylaplacian(state, name, k=k)  # type: ignore[attr-defined]
    raise NotImplementedError(
        f"polylaplacian not implemented for field type {type(state.field).__name__}"
    )

Other public import paths

• omnibias.fields.jax.ops.polylaplacian
• omnibias.fields.jax.ops.high_order.polylaplacian

FUNCTION

spatial_hessian

omnibias.fields.jax.ops.spatial_hessianSource ↗

spatial_hessian(state: FieldState, name: str) -> 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
statepositional-or-keyword	FieldStateRequired	See the operation contract and implementation below.
namepositional-or-keyword	strRequired	See the operation contract and implementation below.

ReturnsArray

Executed regression · sobolev norm cross backend

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_sobolev_norm_cross_backend(make_torch_field, make_jax_field):  # type: ignore[no-untyped-def]
    rule = gauss_legendre(BOUNDS, (8, 8))
    ts = make_torch_field()(torch.as_tensor(rule.nodes, dtype=torch.float64))
    js = make_jax_field()(jnp.asarray(rule.nodes, dtype=jnp.float64))
    for k in (0, 1, 2):
        t = _np(ts.u.sobolev_norm(rule=rule, k=k))
        j = _np(js.u.sobolev_norm(rule=rule, k=k))
        assert np.allclose(t, j, rtol=1e-12, atol=1e-12), k

python -m pytest "packages/omnibias-fields/tests/test_integration_norms.py::test_sobolev_norm_cross_backend" -q

Read the implementation

def spatial_hessian(state: FieldState, name: str) -> Array:
    return hessian(state, name, axes=tuple(state.coordinate_spec.spatial_axes))

Related APIsomnibias.fields.jax.ops.hessian↗

Other public import paths

• omnibias.fields.jax.ops.spatial_hessian
• omnibias.fields.jax.ops.high_order.spatial_hessian

FUNCTION

vector_biharmonic

omnibias.fields.jax.ops.vector_biharmonicSource ↗

vector_biharmonic(state: FieldState, names: tuple[str, ...]) -> 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
statepositional-or-keyword	FieldStateRequired	See the operation contract and implementation below.
namespositional-or-keyword	tuple[str, ...]Required	See the operation contract and implementation below.

ReturnsArray

Read the implementation

def vector_biharmonic(state: FieldState, names: tuple[str, ...]) -> Array:
    cols = [biharmonic(state, n) for n in names]
    return jnp.stack(cols, axis=-1)

Related APIsomnibias.fields.jax.ops.biharmonic↗

Other public import paths

• omnibias.fields.jax.ops.vector_biharmonic
• omnibias.fields.jax.ops.high_order.vector_biharmonic

FUNCTION

vector_hessian

omnibias.fields.jax.ops.vector_hessianSource ↗

vector_hessian(state: FieldState, names: tuple[str, ...], *, axes: tuple[int | str, ...] | None=None) -> 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
statepositional-or-keyword	FieldStateRequired	See the operation contract and implementation below.
namespositional-or-keyword	tuple[str, ...]Required	See the operation contract and implementation below.
axeskeyword-only	tuple[int | str, ...] | NoneDefault: None	See the operation contract and implementation below.

ReturnsArray

Read the implementation

def vector_hessian(
    state: FieldState,
    names: tuple[str, ...],
    *,
    axes: tuple[int | str, ...] | None = None,
) -> Array:
    rows = [hessian(state, n, axes=axes) for n in names]
    return jnp.stack(rows, axis=-3)

Related APIsomnibias.fields.jax.ops.hessian↗

Other public import paths

• omnibias.fields.jax.ops.vector_hessian
• omnibias.fields.jax.ops.high_order.vector_hessian

FUNCTION

vector_laplacian

omnibias.fields.jax.ops.vector_laplacianSource ↗

vector_laplacian(state: FieldState, names: tuple[str, ...]) -> 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
statepositional-or-keyword	FieldStateRequired	See the operation contract and implementation below.
namespositional-or-keyword	tuple[str, ...]Required	See the operation contract and implementation below.

ReturnsArray

Read the implementation

def vector_laplacian(state: FieldState, names: tuple[str, ...]) -> Array:
    from omnibias.fields.jax.ops.basic import laplacian
    cols = [laplacian(state, n) for n in names]
    return jnp.stack(cols, axis=-1)

Other public import paths

• omnibias.fields.jax.ops.vector_laplacian
• omnibias.fields.jax.ops.high_order.vector_laplacian

FUNCTION

vector_polylaplacian

omnibias.fields.jax.ops.vector_polylaplacianSource ↗

vector_polylaplacian(state: FieldState, names: tuple[str, ...], *, k: int) -> 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
statepositional-or-keyword	FieldStateRequired	See the operation contract and implementation below.
namespositional-or-keyword	tuple[str, ...]Required	See the operation contract and implementation below.
kkeyword-only	intRequired	See the operation contract and implementation below.

ReturnsArray

Read the implementation

def vector_polylaplacian(state: FieldState, names: tuple[str, ...], *, k: int) -> Array:
    cols = [polylaplacian(state, n, k=k) for n in names]
    return jnp.stack(cols, axis=-1)

Related APIsomnibias.fields.jax.ops.polylaplacian↗

Other public import paths

• omnibias.fields.jax.ops.vector_polylaplacian
• omnibias.fields.jax.ops.high_order.vector_polylaplacian

← Package overviewRunnable learning examples →