Skip to main content
API REFERENCE/omnibias-fields

omnibias.fields.torch.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.torch.ops.basic
Module contract & mathematical conventions

Basic ops: value, derivative, gradient, divergence, laplacian.

Each op consumes a FieldState, lazily fills the state's SigmaCache with the required sigma^(n)(z) order, and asks the field for the corresponding closed-form chain-rule reduction.

Field-type dispatch lives here: each op picks the right kernel based on type(state.field). The dispatch is intentionally explicit (no multi-dispatch decorator) -- the field types in v0.1 are small and the explicit branches are easy to read in the parity tests.

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.torch.ops.derivativeSource ↗
derivative(state: FieldState, name: str, *, axis: int | str, order: int=1) -> Tensor

d^order f_name / dx_axis^order of shape (B,).

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.
ReturnsTensor
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,
) -> Tensor:
    """``d^order f_name / dx_axis^order`` of shape ``(B,)``."""
    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.torch.ops.value↗
Other public import paths
  • omnibias.fields.torch.ops.derivative
  • omnibias.fields.torch.ops.basic.derivative
FUNCTION

divergence​

omnibias.fields.torch.ops.divergenceSource ↗
divergence(state: FieldState, names: tuple[str, ...]) -> Tensor

nabla . u for a vector field of components names.

Shape (B,). The i-th component is differentiated along the i-th *spatial* axis, in the order given by CoordinateSpec.spatial_axes. len(names) must equal the number of spatial axes.

Parameters

NameType & defaultContract
statepositional-or-keywordFieldStateRequiredSee the operation contract and implementation below.
namespositional-or-keywordtuple[str, ...]RequiredSee the operation contract and implementation below.
ReturnsTensor
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, ...]) -> Tensor:
    """``nabla . u`` for a vector field of components ``names``.

    Shape ``(B,)``. The ``i``-th component is differentiated along the
    ``i``-th *spatial* axis, in the order given by
    :attr:`CoordinateSpec.spatial_axes`. ``len(names)`` must equal the
    number of spatial axes.
    """
    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.torch.ops.derivative↗
Other public import paths
  • omnibias.fields.torch.ops.divergence
  • omnibias.fields.torch.ops.basic.divergence
FUNCTION

gradient​

omnibias.fields.torch.ops.gradientSource ↗
gradient(state: FieldState, name: str, *, axes: tuple[int | str, ...] | None=None) -> Tensor

nabla f_name of shape (B, len(axes)).

Default axes=None means *all spatial* axes (the time axis is excluded). Pass axes=tuple(range(state.coordinate_spec.ndim)) to include time.

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.
ReturnsTensor
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,
) -> Tensor:
    """``nabla f_name`` of shape ``(B, len(axes))``.

    Default ``axes=None`` means *all spatial* axes (the time axis is
    excluded). Pass ``axes=tuple(range(state.coordinate_spec.ndim))``
    to include time.
    """
    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 torch.stack(cols, dim=-1)
    raise NotImplementedError(
        f"gradient op not implemented for field type {type(state.field).__name__}"
    )
Related APIsomnibias.fields.torch.ops.derivative↗
Other public import paths
  • omnibias.fields.torch.ops.gradient
  • omnibias.fields.torch.ops.basic.gradient
FUNCTION

laplacian​

omnibias.fields.torch.ops.laplacianSource ↗
laplacian(state: FieldState, name: str, *, axes: tuple[int | str, ...] | None=None) -> Tensor

Delta f_name = sum_a d^2 f / dx_a^2 over the spatial axes.

Default axes=None reduces over the spatial axes. Pass an explicit axes tuple to over-ride (e.g. include the time axis to get the Wave-equation operator).

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.
ReturnsTensor
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,
) -> Tensor:
    """``Delta f_name = sum_a d^2 f / dx_a^2`` over the spatial axes.

    Default ``axes=None`` reduces over the spatial axes. Pass an
    explicit axes tuple to over-ride (e.g. include the time axis to get
    the Wave-equation operator).
    """
    if axes is None and _is_one_layer(state):
        # Fast-path: closed-form via row-norm-sq.
        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]
    # Fallback: sum of pure 2nd partials along the requested axes.
    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, "axes must be non-empty"
    return out
Related APIsomnibias.fields.torch.ops.derivative↗
Other public import paths
  • omnibias.fields.torch.ops.laplacian
  • omnibias.fields.torch.ops.basic.laplacian
FUNCTION

mixed_partial​

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

d^|orders| f_name / prod_a dx_a^{orders_a} of shape (B,).

Equivalent to repeated single-axis differentiation. Distinct axes produce a true mixed partial; repeated axes are folded by summing orders.

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.
ReturnsTensor
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, ...],
) -> Tensor:
    """``d^|orders| f_name / prod_a dx_a^{orders_a}`` of shape ``(B,)``.

    Equivalent to repeated single-axis differentiation. Distinct axes
    produce a true mixed partial; repeated axes are folded by summing
    orders.
    """
    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)
    # Fold repeated axes into a single (axis, total_order) pair.
    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) or _is_partitioned(state) or _is_jet_mlp(state):
        return state.field.mixed_partial(state, name, int_axes, int_orders)  # type: ignore[attr-defined]
    raise NotImplementedError(
        f"mixed_partial op not implemented for field type {type(state.field).__name__}"
    )
Related APIsomnibias.fields.torch.ops.value↗
Other public import paths
  • omnibias.fields.torch.ops.mixed_partial
  • omnibias.fields.torch.ops.basic.mixed_partial
FUNCTION

stack_components​

omnibias.fields.torch.ops.stack_componentsSource ↗
stack_components(state: FieldState, names: tuple[str, ...]) -> Tensor

Stack component values into a (B, len(names)) tensor.

Parameters

NameType & defaultContract
statepositional-or-keywordFieldStateRequiredSee the operation contract and implementation below.
namespositional-or-keywordtuple[str, ...]RequiredSee the operation contract and implementation below.
ReturnsTensor
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, ...]) -> Tensor:
    """Stack component values into a ``(B, len(names))`` tensor."""
    cols = [value(state, n) for n in names]
    return torch.stack(cols, dim=-1)
Related APIsomnibias.fields.torch.ops.value↗
Other public import paths
  • omnibias.fields.torch.ops.stack_components
  • omnibias.fields.torch.ops.basic.stack_components
FUNCTION

value​

omnibias.fields.torch.ops.valueSource ↗
value(state: FieldState, name: str) -> Tensor

Return f_name(coords) of shape (B,).

Parameters

NameType & defaultContract
statepositional-or-keywordFieldStateRequiredSee the operation contract and implementation below.
namepositional-or-keywordstrRequiredSee the operation contract and implementation below.
ReturnsTensor
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) -> Tensor:
    """Return ``f_name(coords)`` of shape ``(B,)``."""
    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.torch.ops.value
  • omnibias.fields.torch.ops.basic.value
FUNCTION

vector_derivative​

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

Stack of per-component derivatives, shape (B, len(names)).

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.
ReturnsTensor
Read the implementation
def vector_derivative(
    state: FieldState,
    names: tuple[str, ...],
    *,
    axis: int | str,
    order: int = 1,
) -> Tensor:
    """Stack of per-component derivatives, shape ``(B, len(names))``."""
    cols = [derivative(state, n, axis=axis, order=order) for n in names]
    return torch.stack(cols, dim=-1)
Related APIsomnibias.fields.torch.ops.derivative↗
Other public import paths
  • omnibias.fields.torch.ops.vector_derivative
  • omnibias.fields.torch.ops.basic.vector_derivative
← Package overviewRunnable learning examples →