omnibias.fields.jax.ops.tensor

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.tensor

Module contract & mathematical conventions

Divergence of a second-order tensor field (jax).

Bit-identical twin of omnibias.fields.torch.ops.tensor.

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

tensor_cofactor

omnibias.fields.jax.ops.tensor_cofactorSource ↗

tensor_cofactor(a: Array) -> Array

Cofactor matrix \operatorname{cof}A = (\det A)\,A^{-\top}, shape (..., d, d).

Satisfies A @ cof(A)^T = det(A) I (the adjugate/Cramer identity); used for Nanson's formula and the derivative of the Jacobian dJ/dF = cof(F).

Parameters

Name	Type & default	Contract
`a`positional-or-keyword	`Array`Required	See the operation contract and implementation below.

ReturnsArray

Read the implementation

def tensor_cofactor(a: Array) -> Array:
    r"""Cofactor matrix :math:`\operatorname{cof}A = (\det A)\,A^{-\top}`, shape ``(..., d, d)``.

    Satisfies ``A @ cof(A)^T = det(A) I`` (the adjugate/Cramer identity); used
    for Nanson's formula and the derivative of the Jacobian ``dJ/dF = cof(F)``.
    """
    a = jnp.asarray(a)
    _require_square(a, "tensor_cofactor")
    det = jnp.linalg.det(a)
    inv_t = jnp.swapaxes(jnp.linalg.inv(a), -1, -2)
    return det[..., None, None] * inv_t

Other public import paths

omnibias.fields.jax.ops.tensor_cofactor
omnibias.fields.jax.ops.tensor.tensor_cofactor

FUNCTION

tensor_determinant

omnibias.fields.jax.ops.tensor_determinantSource ↗

tensor_determinant(a: Array) -> Array

Determinant of the trailing block \det A, shape (...,).

Parameters

Name	Type & default	Contract
`a`positional-or-keyword	`Array`Required	See the operation contract and implementation below.

ReturnsArray

Read the implementation

def tensor_determinant(a: Array) -> Array:
    r"""Determinant of the trailing block :math:`\det A`, shape ``(...,)``."""
    a = jnp.asarray(a)
    _require_square(a, "tensor_determinant")
    return jnp.linalg.det(a)

Other public import paths

omnibias.fields.jax.ops.tensor_determinant
omnibias.fields.jax.ops.tensor.tensor_determinant

FUNCTION

tensor_divergence

omnibias.fields.jax.ops.tensor_divergenceSource ↗

tensor_divergence(state: FieldState, sigma_names: Sequence[Sequence[str]]) -> Array

Row-wise divergence (\nabla\cdot\sigma)_i = \partial_j \sigma_{ij}.

See omnibias.fields.torch.ops.tensor.tensor_divergence.

Parameters

Name	Type & default	Contract
`state`positional-or-keyword	`FieldState`Required	See the operation contract and implementation below.
`sigma_names`positional-or-keyword	`Sequence[Sequence[str]]`Required	See the operation contract and implementation below.

ReturnsArray

Explicit validation & failure cases

When len(rows) != d, raises ValueError(f'tensor_divergence expects a {d}x{d} layout (d = number of spatial axes {spatial!r}); got {len(rows)} rows').
When len(row) != d, raises ValueError(f'row {i} of sigma_names must have {d} entries, got {len(row)}').

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

Executed regression · tensor divergence 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_tensor_divergence_cross_backend(make_torch_field, make_jax_field, grid_nodes):  # type: ignore[no-untyped-def]
    ts = make_torch_field()(torch.as_tensor(grid_nodes, dtype=torch.float64))
    js = make_jax_field()(jnp.asarray(grid_nodes, dtype=jnp.float64))
    layout = (("u", "v"), ("v", "u"))
    t = _np(ts.ops.tensor_divergence(ts, layout))
    j = _np(js.ops.tensor_divergence(js, layout))
    assert np.allclose(t, j, rtol=1e-12, atol=1e-12)

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

Read the implementation

def tensor_divergence(
    state: FieldState,
    sigma_names: Sequence[Sequence[str]],
) -> Array:
    r"""Row-wise divergence :math:`(\nabla\cdot\sigma)_i = \partial_j \sigma_{ij}`.

    See :func:`omnibias.fields.torch.ops.tensor.tensor_divergence`.
    """
    spatial = state.coordinate_spec.spatial_axes
    d = len(spatial)
    rows = list(sigma_names)
    if len(rows) != d:
        raise ValueError(
            f"tensor_divergence expects a {d}x{d} layout (d = number of spatial "
            f"axes {spatial!r}); got {len(rows)} rows"
        )
    cols: list[Array] = []
    for i, row in enumerate(rows):
        row = list(row)
        if len(row) != d:
            raise ValueError(
                f"row {i} of sigma_names must have {d} entries, got {len(row)}"
            )
        acc: Array | None = None
        for j in range(d):
            term = derivative(state, row[j], axis=spatial[j], order=1)
            acc = term if acc is None else acc + term
        assert acc is not None
        cols.append(acc)
    return jnp.stack(cols, axis=-1)

Other public import paths

omnibias.fields.jax.ops.tensor_divergence
omnibias.fields.jax.ops.tensor.tensor_divergence

FUNCTION

tensor_double_dot

omnibias.fields.jax.ops.tensor_double_dotSource ↗

tensor_double_dot(a: Array, b: Array) -> Array

Frobenius double contraction A:B = \sum_{ij} A_{ij} B_{ij}.

Operates on the trailing two axes of two equally-shaped tensor batches (e.g. stress : strain_rate for viscous dissipation). Returns shape (B,).

Parameters

Name	Type & default	Contract
`a`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`b`positional-or-keyword	`Array`Required	See the operation contract and implementation below.

ReturnsArray

Explicit validation & failure cases

When a.shape != b.shape, raises ValueError(f'tensor_double_dot requires equal shapes; got {tuple(a.shape)} and {tuple(b.shape)}').
When a.ndim < 2, raises ValueError('tensor_double_dot requires at least 2 trailing axes').

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

Read the implementation

def tensor_double_dot(a: Array, b: Array) -> Array:
    r"""Frobenius double contraction :math:`A:B = \sum_{ij} A_{ij} B_{ij}`.

    Operates on the trailing two axes of two equally-shaped tensor batches
    (e.g. ``stress : strain_rate`` for viscous dissipation). Returns shape ``(B,)``.
    """
    a = jnp.asarray(a)
    b = jnp.asarray(b)
    if a.shape != b.shape:
        raise ValueError(
            f"tensor_double_dot requires equal shapes; got {tuple(a.shape)} "
            f"and {tuple(b.shape)}"
        )
    if a.ndim < 2:
        raise ValueError("tensor_double_dot requires at least 2 trailing axes")
    return (a * b).sum(axis=(-2, -1))

Other public import paths

omnibias.fields.jax.ops.tensor_double_dot
omnibias.fields.jax.ops.tensor.tensor_double_dot

FUNCTION

tensor_inverse

omnibias.fields.jax.ops.tensor_inverseSource ↗

tensor_inverse(a: Array) -> Array

Inverse of the trailing block A^{-1}, shape (..., d, d).

Parameters

Name	Type & default	Contract
`a`positional-or-keyword	`Array`Required	See the operation contract and implementation below.

ReturnsArray

Read the implementation

def tensor_inverse(a: Array) -> Array:
    r"""Inverse of the trailing block :math:`A^{-1}`, shape ``(..., d, d)``."""
    a = jnp.asarray(a)
    _require_square(a, "tensor_inverse")
    return jnp.linalg.inv(a)

Other public import paths

omnibias.fields.jax.ops.tensor_inverse
omnibias.fields.jax.ops.tensor.tensor_inverse

FUNCTION

tensor_matmul

omnibias.fields.jax.ops.tensor_matmulSource ↗

tensor_matmul(a: Array, b: Array) -> Array

Batched matrix product (AB)_{ik} = A_{ij} B_{jk}, shape (..., d, d).

Parameters

Name	Type & default	Contract
`a`positional-or-keyword	`Array`Required	See the operation contract and implementation below.
`b`positional-or-keyword	`Array`Required	See the operation contract and implementation below.

ReturnsArray

Read the implementation

def tensor_matmul(a: Array, b: Array) -> Array:
    r"""Batched matrix product :math:`(AB)_{ik} = A_{ij} B_{jk}`, shape ``(..., d, d)``."""
    return jnp.matmul(a, b)

Other public import paths

omnibias.fields.jax.ops.tensor_matmul
omnibias.fields.jax.ops.tensor.tensor_matmul

FUNCTION

tensor_trace

omnibias.fields.jax.ops.tensor_traceSource ↗

tensor_trace(a: Array) -> Array

Trace of the trailing block \operatorname{tr}A = A_{ii}, shape (...,).

Parameters

Name	Type & default	Contract
`a`positional-or-keyword	`Array`Required	See the operation contract and implementation below.

ReturnsArray

Read the implementation

def tensor_trace(a: Array) -> Array:
    r"""Trace of the trailing block :math:`\operatorname{tr}A = A_{ii}`, shape ``(...,)``."""
    a = jnp.asarray(a)
    _require_square(a, "tensor_trace")
    return jnp.trace(a, axis1=-2, axis2=-1)

Other public import paths

omnibias.fields.jax.ops.tensor_trace
omnibias.fields.jax.ops.tensor.tensor_trace

FUNCTION

tensor_transpose

omnibias.fields.jax.ops.tensor_transposeSource ↗

tensor_transpose(a: Array) -> Array

Transpose the trailing two axes: A^\top, shape (..., d, d).

Parameters

Name	Type & default	Contract
`a`positional-or-keyword	`Array`Required	See the operation contract and implementation below.

ReturnsArray

Read the implementation

def tensor_transpose(a: Array) -> Array:
    r"""Transpose the trailing two axes: :math:`A^\top`, shape ``(..., d, d)``."""
    a = jnp.asarray(a)
    _require_square(a, "tensor_transpose")
    return jnp.swapaxes(a, -1, -2)

Other public import paths

omnibias.fields.jax.ops.tensor_transpose
omnibias.fields.jax.ops.tensor.tensor_transpose

← Package overview Runnable learning examples →

omnibias.fields.jax.ops.tensor

tensor_cofactor​

Parameters

tensor_determinant​

Parameters

tensor_divergence​

Parameters

tensor_double_dot​

Parameters

tensor_inverse​

Parameters

tensor_matmul​

Parameters

tensor_trace​

Parameters

tensor_transpose​

Parameters

tensor_cofactor

tensor_determinant

tensor_divergence

tensor_double_dot

tensor_inverse

tensor_matmul

tensor_trace

tensor_transpose