Skip to main content
API REFERENCE/omnibias-fields

omnibias.fields.jax.ops.basic

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.basic
Module contract & mathematical conventions

JAX basic ops (mirrors omnibias.fields.torch.ops.basic).

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

derivative​

omnibias.fields.jax.ops.derivativeSource ↗
derivative(state: FieldState, name: str, *, axis: int | str, order: int=1) -> Array

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

Parameters

NameType & defaultContract
statepositional-or-keywordFieldStateRequiredSee the operation contract and implementation below.
namepositional-or-keywordstrRequiredSee the operation contract and implementation below.
axiskeyword-onlyint | strRequiredSee the operation contract and implementation below.
orderkeyword-onlyintDefault: 1See the operation contract and implementation below.
ReturnsArray
Explicit validation & failure cases
  • When order < 1, raises ValueError(f'derivative order must be >= 0, got {order}').

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

Executed regression · cross backend bit 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_cross_backend_bit_parity(which, opfn):
    nodes = _nodes(2 if which == "2d" else 3, seed=7)
    ts = _state("torch", which, nodes)
    js = _state("jax", which, nodes)
    t = _np(opfn(ts.ops, ts))
    j = _np(opfn(js.ops, js))
    assert np.allclose(t, j, rtol=1e-12, atol=1e-12)
python -m pytest "packages/omnibias-fields/tests/test_vector_calculus.py::test_cross_backend_bit_parity[2d-<lambda>]" -q
Read the implementation
def derivative(
    state: FieldState, name: str, *, axis: int | str, order: int = 1,
) -> Array:
    if order < 1:
        if order == 0:
            return value(state, name)
        raise ValueError(f"derivative order must be >= 0, got {order}")
    a = _resolve_axis(state, axis)
    if _is_one_layer(state):
        sigma_n = _sigma_of_order(state, order)
        return state.field.nth_partial(sigma_n, name, a, order)
    if _is_spectral(state):
        return state.field.derivative(state, name, axis=a, order=order)  # type: ignore[attr-defined]
    if _is_chebyshev(state):
        return state.field.derivative(state, name, axis=a, order=order)  # type: ignore[attr-defined]
    if _is_cage(state) or _is_partitioned(state) or _is_jet_mlp(state):
        return state.field.derivative(state, name, axis=a, order=order)  # type: ignore[attr-defined]
    raise NotImplementedError(
        f"derivative op not implemented for field type {type(state.field).__name__}"
    )
Related APIsomnibias.fields.jax.ops.value↗
Other public import paths
  • omnibias.fields.jax.ops.derivative
  • omnibias.fields.jax.ops.basic.derivative
FUNCTION

divergence​

omnibias.fields.jax.ops.divergenceSource ↗
divergence(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

NameType & defaultContract
statepositional-or-keywordFieldStateRequiredSee the operation contract and implementation below.
namespositional-or-keywordtuple[str, ...]RequiredSee the operation contract and implementation below.
ReturnsArray
Explicit validation & failure cases
  • When len(names) != len(sa), raises ValueError(f'divergence: number of components {len(names)} must equal number of spatial axes {len(sa)} ({sa!r})').

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

Executed regression · cross backend bit 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_cross_backend_bit_parity(which, opfn):
    nodes = _nodes(2 if which == "2d" else 3, seed=7)
    ts = _state("torch", which, nodes)
    js = _state("jax", which, nodes)
    t = _np(opfn(ts.ops, ts))
    j = _np(opfn(js.ops, js))
    assert np.allclose(t, j, rtol=1e-12, atol=1e-12)
python -m pytest "packages/omnibias-fields/tests/test_vector_calculus.py::test_cross_backend_bit_parity[3d-<lambda>6]" -q
Read the implementation
def divergence(state: FieldState, names: tuple[str, ...]) -> Array:
    sa = state.coordinate_spec.spatial_axes
    if len(names) != len(sa):
        raise ValueError(
            f"divergence: number of components {len(names)} must equal "
            f"number of spatial axes {len(sa)} ({sa!r})"
        )
    out = None
    for n, a in zip(names, sa, strict=False):
        d = derivative(state, n, axis=a, order=1)
        out = d if out is None else out + d
    assert out is not None
    return out
Related APIsomnibias.fields.jax.ops.derivative↗
Other public import paths
  • omnibias.fields.jax.ops.divergence
  • omnibias.fields.jax.ops.basic.divergence
FUNCTION

gradient​

omnibias.fields.jax.ops.gradientSource ↗
gradient(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

NameType & defaultContract
statepositional-or-keywordFieldStateRequiredSee the operation contract and implementation below.
namepositional-or-keywordstrRequiredSee the operation contract and implementation below.
axeskeyword-onlytuple[int | str, ...] | NoneDefault: NoneSee the operation contract and implementation below.
ReturnsArray
Executed regression · cross backend bit 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_cross_backend_bit_parity(which, opfn):
    nodes = _nodes(2 if which == "2d" else 3, seed=7)
    ts = _state("torch", which, nodes)
    js = _state("jax", which, nodes)
    t = _np(opfn(ts.ops, ts))
    j = _np(opfn(js.ops, js))
    assert np.allclose(t, j, rtol=1e-12, atol=1e-12)
python -m pytest "packages/omnibias-fields/tests/test_vector_calculus.py::test_cross_backend_bit_parity[3d-<lambda>2]" -q
Read the implementation
def gradient(
    state: FieldState,
    name: str,
    *,
    axes: tuple[int | str, ...] | None = None,
) -> Array:
    if axes is None:
        axis_idx = tuple(
            state.coordinate_spec.axis_index(a)
            for a in state.coordinate_spec.spatial_axes
        )
    else:
        axis_idx = tuple(_resolve_axis(state, a) for a in axes)
    if _is_one_layer(state):
        sigma_p = _sigma_of_order(state, 1)
        full = state.field.gradient_full(sigma_p, name)
        if axis_idx == tuple(range(full.shape[-1])):
            return full
        return full[..., list(axis_idx)]
    if _is_jet_mlp(state):
        # One jet already holds every first partial; read them all in one go.
        full = state.field.gradient_full(state, name)  # type: ignore[attr-defined]
        if axis_idx == tuple(range(full.shape[-1])):
            return full
        return full[..., list(axis_idx)]
    if _is_spectral(state) or _is_chebyshev(state) or _is_cage(state) or _is_partitioned(state):
        cols = [
            derivative(state, name, axis=a, order=1) for a in axis_idx
        ]
        return jnp.stack(cols, axis=-1)
    raise NotImplementedError(
        f"gradient op not implemented for field type {type(state.field).__name__}"
    )
Related APIsomnibias.fields.jax.ops.derivative↗
Other public import paths
  • omnibias.fields.jax.ops.gradient
  • omnibias.fields.jax.ops.basic.gradient
FUNCTION

laplacian​

omnibias.fields.jax.ops.laplacianSource ↗
laplacian(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

NameType & defaultContract
statepositional-or-keywordFieldStateRequiredSee the operation contract and implementation below.
namepositional-or-keywordstrRequiredSee the operation contract and implementation below.
axeskeyword-onlytuple[int | str, ...] | NoneDefault: NoneSee the operation contract and implementation below.
ReturnsArray
Executed regression · cross backend bit 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_cross_backend_bit_parity(which, opfn):
    nodes = _nodes(2 if which == "2d" else 3, seed=7)
    ts = _state("torch", which, nodes)
    js = _state("jax", which, nodes)
    t = _np(opfn(ts.ops, ts))
    j = _np(opfn(js.ops, js))
    assert np.allclose(t, j, rtol=1e-12, atol=1e-12)
python -m pytest "packages/omnibias-fields/tests/test_vector_calculus.py::test_cross_backend_bit_parity[3d-<lambda>5]" -q
Read the implementation
def laplacian(
    state: FieldState,
    name: str,
    *,
    axes: tuple[int | str, ...] | None = None,
) -> Array:
    if axes is None and _is_one_layer(state):
        sigma_pp = _sigma_of_order(state, 2)
        return state.field.laplacian(sigma_pp, name)
    if axes is None and _is_jet_mlp(state):
        # Fast-path: the spatial trace of the Hessian block of a single jet.
        return state.field.laplacian(state, name)  # type: ignore[attr-defined]
    if axes is None:
        axes = tuple(state.coordinate_spec.spatial_axes)
    out = None
    for a in axes:
        d2 = derivative(state, name, axis=a, order=2)
        out = d2 if out is None else out + d2
    assert out is not None
    return out
Related APIsomnibias.fields.jax.ops.derivative↗
Other public import paths
  • omnibias.fields.jax.ops.laplacian
  • omnibias.fields.jax.ops.basic.laplacian
FUNCTION

mixed_partial​

omnibias.fields.jax.ops.mixed_partialSource ↗
mixed_partial(state: FieldState, name: str, axes: tuple[int | str, ...], orders: tuple[int, ...]) -> Array

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

Parameters

NameType & defaultContract
statepositional-or-keywordFieldStateRequiredSee the operation contract and implementation below.
namepositional-or-keywordstrRequiredSee the operation contract and implementation below.
axespositional-or-keywordtuple[int | str, ...]RequiredSee the operation contract and implementation below.
orderspositional-or-keywordtuple[int, ...]RequiredSee the operation contract and implementation below.
ReturnsArray
Explicit validation & failure cases
  • When len(axes) != len(orders), raises ValueError(f'axes and orders must have the same length; got {len(axes)} and {len(orders)}').

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

Executed regression · cross backend bit 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_cross_backend_bit_parity(which, opfn):
    nodes = _nodes(2 if which == "2d" else 3, seed=7)
    ts = _state("torch", which, nodes)
    js = _state("jax", which, nodes)
    t = _np(opfn(ts.ops, ts))
    j = _np(opfn(js.ops, js))
    assert np.allclose(t, j, rtol=1e-12, atol=1e-12)
python -m pytest "packages/omnibias-fields/tests/test_vector_calculus.py::test_cross_backend_bit_parity[2d-<lambda>]" -q
Read the implementation
def mixed_partial(
    state: FieldState,
    name: str,
    axes: tuple[int | str, ...],
    orders: tuple[int, ...],
) -> Array:
    if len(axes) != len(orders):
        raise ValueError(
            f"axes and orders must have the same length; got {len(axes)} and {len(orders)}"
        )
    if not axes:
        return value(state, name)
    folded: dict[int, int] = {}
    for a, o in zip(axes, orders, strict=False):
        if o < 1:
            continue
        ai = _resolve_axis(state, a)
        folded[ai] = folded.get(ai, 0) + int(o)
    if not folded:
        return value(state, name)
    int_axes = tuple(folded)
    int_orders = tuple(folded[a] for a in int_axes)
    total_order = sum(int_orders)
    if _is_one_layer(state):
        sigma_n = _sigma_of_order(state, total_order)
        return state.field.mixed_partial(sigma_n, name, int_axes, int_orders)
    if _is_spectral(state):
        return state.field.mixed_partial(state, name, int_axes, int_orders)  # type: ignore[attr-defined]
    if _is_chebyshev(state):
        return state.field.mixed_partial(state, name, int_axes, int_orders)  # type: ignore[attr-defined]
    if _is_cage(state):
        return state.field.mixed_partial(state, name, int_axes, int_orders)  # type: ignore[attr-defined]
    if _is_partitioned(state):
        return state.field.mixed_partial(state, name, int_axes, int_orders)  # type: ignore[attr-defined]
    if _is_jet_mlp(state):
        return state.field.mixed_partial(state, name, int_axes, int_orders)  # type: ignore[attr-defined]
    raise NotImplementedError(
        f"mixed_partial not implemented for field type {type(state.field).__name__}"
    )
Related APIsomnibias.fields.jax.ops.value↗
Other public import paths
  • omnibias.fields.jax.ops.mixed_partial
  • omnibias.fields.jax.ops.basic.mixed_partial
FUNCTION

stack_components​

omnibias.fields.jax.ops.stack_componentsSource ↗
stack_components(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

NameType & defaultContract
statepositional-or-keywordFieldStateRequiredSee the operation contract and implementation below.
namespositional-or-keywordtuple[str, ...]RequiredSee the operation contract and implementation below.
ReturnsArray
Executed regression · cross backend bit 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_cross_backend_bit_parity(which, opfn):
    nodes = _nodes(2 if which == "2d" else 3, seed=7)
    ts = _state("torch", which, nodes)
    js = _state("jax", which, nodes)
    t = _np(opfn(ts.ops, ts))
    j = _np(opfn(js.ops, js))
    assert np.allclose(t, j, rtol=1e-12, atol=1e-12)
python -m pytest "packages/omnibias-fields/tests/test_vector_calculus.py::test_cross_backend_bit_parity[3d-<lambda>6]" -q
Read the implementation
def stack_components(state: FieldState, names: tuple[str, ...]) -> Array:
    cols = [value(state, n) for n in names]
    return jnp.stack(cols, axis=-1)
Related APIsomnibias.fields.jax.ops.value↗
Other public import paths
  • omnibias.fields.jax.ops.stack_components
  • omnibias.fields.jax.ops.basic.stack_components
FUNCTION

value​

omnibias.fields.jax.ops.valueSource ↗
value(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

NameType & defaultContract
statepositional-or-keywordFieldStateRequiredSee the operation contract and implementation below.
namepositional-or-keywordstrRequiredSee the operation contract and implementation below.
ReturnsArray
Executed regression · cross backend bit 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_cross_backend_bit_parity(which, opfn):
    nodes = _nodes(2 if which == "2d" else 3, seed=7)
    ts = _state("torch", which, nodes)
    js = _state("jax", which, nodes)
    t = _np(opfn(ts.ops, ts))
    j = _np(opfn(js.ops, js))
    assert np.allclose(t, j, rtol=1e-12, atol=1e-12)
python -m pytest "packages/omnibias-fields/tests/test_vector_calculus.py::test_cross_backend_bit_parity[3d-<lambda>4]" -q
Read the implementation
def value(state: FieldState, name: str) -> Array:
    if _is_one_layer(state):
        sigma_z = _sigma_of_order(state, 0)
        return state.field.value(sigma_z, name)
    if _is_spectral(state):
        return state.field.value_component(state, name)  # type: ignore[attr-defined]
    if _is_chebyshev(state):
        return state.field.value_component(state, name)  # type: ignore[attr-defined]
    if _is_cage(state) or _is_partitioned(state) or _is_jet_mlp(state):
        return state.field.value_component(state, name)  # type: ignore[attr-defined]
    raise NotImplementedError(
        f"value op not implemented for field type {type(state.field).__name__}"
    )
Other public import paths
  • omnibias.fields.jax.ops.value
  • omnibias.fields.jax.ops.basic.value
FUNCTION

vector_derivative​

omnibias.fields.jax.ops.vector_derivativeSource ↗
vector_derivative(state: FieldState, names: tuple[str, ...], *, axis: int | str, order: int=1) -> Array

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

Parameters

NameType & defaultContract
statepositional-or-keywordFieldStateRequiredSee the operation contract and implementation below.
namespositional-or-keywordtuple[str, ...]RequiredSee the operation contract and implementation below.
axiskeyword-onlyint | strRequiredSee the operation contract and implementation below.
orderkeyword-onlyintDefault: 1See the operation contract and implementation below.
ReturnsArray
Read the implementation
def vector_derivative(
    state: FieldState,
    names: tuple[str, ...],
    *,
    axis: int | str,
    order: int = 1,
) -> Array:
    cols = [derivative(state, n, axis=axis, order=order) for n in names]
    return jnp.stack(cols, axis=-1)
Related APIsomnibias.fields.jax.ops.derivative↗
Other public import paths
  • omnibias.fields.jax.ops.vector_derivative
  • omnibias.fields.jax.ops.basic.vector_derivative
← Package overviewRunnable learning examples →