Skip to main content
API REFERENCE/omnibias-graph

omnibias.graph.jax.ops.spectral

0.1.0a2AlphaPython >=3.10AGPL / commercial

Spectral graph operators and differentiable relaxations for ordering, assignment, and graph-structured learning.

import omnibias.graph.jax.ops.spectral
Module contract & mathematical conventions

Differentiable spectral graph operators (jax).

Bit-identical twin of omnibias.graph.torch.ops.spectral; see that module for the mathematical definitions. All operators take a symmetric non-negative adjacency A of shape (n, n) and are differentiable in A.

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 →

External type aliases

These public names refer to external types. Their implementations and inherited members belong to the originating library.

Public nameExternal definition
omnibias.graph.jax.ops.spectral.Arraytyping.Any
CLASS

CutRelaxation​

omnibias.graph.jax.ops.CutRelaxationSource ↗
class CutRelaxation

Dataclass constructor fields; inherited fields, if any, are declared by the base class. A field(...) expression describes a default factory or field options, not a literal argument default.

Rayleigh-Ritz relaxation of the k-way ratio / normalized cut.

Parameters

NameType & defaultContract
embeddingdataclass fieldArrayRequiredSee the operation contract and implementation below.
relaxed_cut_valuedataclass fieldArrayRequiredSee the operation contract and implementation below.
eigenvaluesdataclass fieldArrayRequiredSee the operation contract and implementation below.

Fields & methods

embeddingfield
embedding: Array

Declared Array field on CutRelaxation. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.

View source ↗
relaxed_cut_valuefield
relaxed_cut_value: Array

Declared Array field on CutRelaxation. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.

View source ↗
eigenvaluesfield
eigenvalues: Array

Declared Array field on CutRelaxation. A value must be provided when constructing the configuration, unless an inherited constructor supplies it.

View source ↗
Read the implementation
class CutRelaxation:
    """Rayleigh-Ritz relaxation of the ``k``-way ratio / normalized cut."""

    embedding: Array
    relaxed_cut_value: Array
    eigenvalues: Array
Other public import paths
  • omnibias.graph.jax.ops.CutRelaxation
  • omnibias.graph.jax.ops.spectral.CutRelaxation
FUNCTION

degree​

omnibias.graph.jax.ops.degreeSource ↗
degree(adjacency: Array) -> Array

Weighted node degrees d_i = sum_j A_{ij} -> shape (n,).

Parameters

NameType & defaultContract
adjacencypositional-or-keywordArrayRequiredSee the operation contract and implementation below.
ReturnsArray
Read the implementation
def degree(adjacency: Array) -> Array:
    r"""Weighted node degrees ``d_i = sum_j A_{ij}`` -> shape ``(n,)``."""
    _check_adjacency(adjacency)
    return adjacency.sum(axis=-1)
Other public import paths
  • omnibias.graph.jax.ops.degree
  • omnibias.graph.jax.ops.spectral.degree
FUNCTION

fiedler_vector​

omnibias.graph.jax.ops.fiedler_vectorSource ↗
fiedler_vector(adjacency: Array, *, normalized: bool=False) -> Array

The eigenvector of the second-smallest Laplacian eigenvalue -> (n,).

Parameters

NameType & defaultContract
adjacencypositional-or-keywordArrayRequiredSee the operation contract and implementation below.
normalizedkeyword-onlyboolDefault: FalseSee the operation contract and implementation below.
ReturnsArray
Executed regression · fiedler separates blocks

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_fiedler_separates_blocks() -> None:
    v = np.asarray(G.fiedler_vector(two_block_sbm(5, 1.0, 0.02)))
    b0, b1 = v[:5], v[5:]
    assert np.all(np.sign(b0) == np.sign(b0[0]))
    assert np.all(np.sign(b1) == np.sign(b1[0]))
    assert np.sign(b0[0]) != np.sign(b1[0])
python -m pytest "packages/omnibias-graph/tests/test_spectral_jax.py::test_fiedler_separates_blocks" -q
Read the implementation
def fiedler_vector(adjacency: Array, *, normalized: bool = False) -> Array:
    r"""The eigenvector of the second-smallest Laplacian eigenvalue -> ``(n,)``."""
    _, evecs = laplacian_spectrum(adjacency, normalized=normalized)
    return evecs[:, 1]
Related APIsomnibias.graph.jax.ops.laplacian_spectrum↗
Other public import paths
  • omnibias.graph.jax.ops.fiedler_vector
  • omnibias.graph.jax.ops.spectral.fiedler_vector
FUNCTION

graph_heat_kernel​

omnibias.graph.jax.ops.graph_heat_kernelSource ↗
graph_heat_kernel(adjacency: Array, t: float, *, normalized: bool=False) -> Array

Heat kernel H = exp(-t L) via the symmetric eigendecomposition.

Parameters

NameType & defaultContract
adjacencypositional-or-keywordArrayRequiredSee the operation contract and implementation below.
tpositional-or-keywordfloatRequiredSee the operation contract and implementation below.
normalizedkeyword-onlyboolDefault: FalseSee the operation contract and implementation below.
ReturnsArray
Explicit validation & failure cases
  • When t < 0, raises ValueError(f'diffusion time t must be >= 0; got {t}').

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

Executed regression · heat kernel 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_heat_kernel_parity(t: float) -> None:
    m = _sym_adjacency(5)
    _assert_close(GT.graph_heat_kernel(torch.tensor(m), t), GJ.graph_heat_kernel(jnp.asarray(m), t))
python -m pytest "packages/omnibias-graph/tests/test_cross_backend.py::test_heat_kernel_parity[0.1]" -q
Regression example · heat kernel semigroup and constant

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_heat_kernel_semigroup_and_constant() -> None:
    a = two_block_sbm(3, 1.0, 0.3)
    h_half = np.asarray(G.graph_heat_kernel(a, 0.25))
    h_full = np.asarray(G.graph_heat_kernel(a, 0.5))
    assert np.allclose(h_half @ h_half, h_full, atol=1e-10)
    ones = np.ones(6)
    assert np.allclose(np.asarray(G.graph_heat_kernel(a, 1.0)) @ ones, ones, atol=1e-10)
python -m pytest packages/omnibias-graph/tests/test_spectral_jax.py -k test_heat_kernel_semigroup_and_constant -q
Read the implementation
def graph_heat_kernel(adjacency: Array, t: float, *, normalized: bool = False) -> Array:
    r"""Heat kernel ``H = exp(-t L)`` via the symmetric eigendecomposition."""
    if t < 0:
        raise ValueError(f"diffusion time t must be >= 0; got {t}")
    evals, evecs = laplacian_spectrum(adjacency, normalized=normalized)
    damp = jnp.exp(-t * evals)
    return (evecs * damp[None, :]) @ evecs.T
Related APIsomnibias.graph.jax.ops.laplacian_spectrum↗
Other public import paths
  • omnibias.graph.jax.ops.graph_heat_kernel
  • omnibias.graph.jax.ops.spectral.graph_heat_kernel
FUNCTION

graph_laplacian​

omnibias.graph.jax.ops.graph_laplacianSource ↗
graph_laplacian(adjacency: Array) -> Array

Combinatorial Laplacian L = D - A.

Parameters

NameType & defaultContract
adjacencypositional-or-keywordArrayRequiredSee the operation contract and implementation below.
ReturnsArray
Executed regression · non square raises

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_non_square_raises() -> None:
    with pytest.raises(ValueError, match="square"):
        G.graph_laplacian(jnp.zeros((3, 4)))
python -m pytest "packages/omnibias-graph/tests/test_spectral_jax.py::test_non_square_raises" -q
Read the implementation
def graph_laplacian(adjacency: Array) -> Array:
    r"""Combinatorial Laplacian ``L = D - A``."""
    _check_adjacency(adjacency)
    d = adjacency.sum(axis=-1)
    return jnp.diag(d) - adjacency
Other public import paths
  • omnibias.graph.jax.ops.graph_laplacian
  • omnibias.graph.jax.ops.spectral.graph_laplacian
FUNCTION

laplacian_spectrum​

omnibias.graph.jax.ops.laplacian_spectrumSource ↗
laplacian_spectrum(adjacency: Array, *, normalized: bool=False) -> tuple[Array, Array]

Ascending eigenvalues / eigenvectors of the (symmetric) Laplacian.

Parameters

NameType & defaultContract
adjacencypositional-or-keywordArrayRequiredSee the operation contract and implementation below.
normalizedkeyword-onlyboolDefault: FalseSee the operation contract and implementation below.
Returnstuple[Array, Array]
Executed regression · heat kernel 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_heat_kernel_parity(t: float) -> None:
    m = _sym_adjacency(5)
    _assert_close(GT.graph_heat_kernel(torch.tensor(m), t), GJ.graph_heat_kernel(jnp.asarray(m), t))
python -m pytest "packages/omnibias-graph/tests/test_cross_backend.py::test_heat_kernel_parity[0.1]" -q
Regression example · ring eigenvalues

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_ring_eigenvalues(n: int) -> None:
    evals, _ = G.laplacian_spectrum(ring_adjacency(n))
    exact = np.sort([2 - 2 * math.cos(2 * math.pi * k / n) for k in range(n)])
    assert np.allclose(np.asarray(evals), exact, atol=1e-10)
python -m pytest packages/omnibias-graph/tests/test_spectral_jax.py -k test_ring_eigenvalues -q
Read the implementation
def laplacian_spectrum(
    adjacency: Array, *, normalized: bool = False
) -> tuple[Array, Array]:
    r"""Ascending eigenvalues / eigenvectors of the (symmetric) Laplacian."""
    lap = _symmetric_laplacian(adjacency, normalized=normalized)
    lap = 0.5 * (lap + lap.T)
    evals, evecs = jnp.linalg.eigh(lap)
    return evals, evecs
Other public import paths
  • omnibias.graph.jax.ops.laplacian_spectrum
  • omnibias.graph.jax.ops.spectral.laplacian_spectrum
FUNCTION

normalized_laplacian​

omnibias.graph.jax.ops.normalized_laplacianSource ↗
normalized_laplacian(adjacency: Array) -> Array

Symmetric normalized Laplacian L_sym = I - D^{-1/2} A D^{-1/2}.

Parameters

NameType & defaultContract
adjacencypositional-or-keywordArrayRequiredSee the operation contract and implementation below.
ReturnsArray
Executed regression · laplacian 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_laplacian_parity(seed: int, name: str) -> None:
    m = _sym_adjacency(seed)
    _assert_close(getattr(GT, name)(torch.tensor(m)), getattr(GJ, name)(jnp.asarray(m)))
python -m pytest "packages/omnibias-graph/tests/test_cross_backend.py::test_laplacian_parity[normalized_laplacian-0]" -q
Regression example · normalized spectrum in 0 2

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_normalized_spectrum_in_0_2() -> None:
    ln = G.normalized_laplacian(ring_adjacency(10))
    ev = np.linalg.eigvalsh(np.asarray(ln))
    assert ev.min() > -1e-10 and ev.max() < 2.0 + 1e-10
python -m pytest packages/omnibias-graph/tests/test_spectral_jax.py -k test_normalized_spectrum_in_0_2 -q
Read the implementation
def normalized_laplacian(adjacency: Array) -> Array:
    r"""Symmetric normalized Laplacian ``L_sym = I - D^{-1/2} A D^{-1/2}``."""
    _check_adjacency(adjacency)
    d = adjacency.sum(axis=-1)
    d_inv_sqrt = jnp.where(d > 0, 1.0 / jnp.sqrt(d), jnp.zeros_like(d))
    a_norm = d_inv_sqrt[:, None] * adjacency * d_inv_sqrt[None, :]
    n = adjacency.shape[0]
    return jnp.eye(n, dtype=adjacency.dtype) - a_norm
Other public import paths
  • omnibias.graph.jax.ops.normalized_laplacian
  • omnibias.graph.jax.ops.spectral.normalized_laplacian
FUNCTION

random_walk_laplacian​

omnibias.graph.jax.ops.random_walk_laplacianSource ↗
random_walk_laplacian(adjacency: Array) -> Array

Random-walk Laplacian L_rw = I - D^{-1} A.

Parameters

NameType & defaultContract
adjacencypositional-or-keywordArrayRequiredSee the operation contract and implementation below.
ReturnsArray
Executed regression · laplacian 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_laplacian_parity(seed: int, name: str) -> None:
    m = _sym_adjacency(seed)
    _assert_close(getattr(GT, name)(torch.tensor(m)), getattr(GJ, name)(jnp.asarray(m)))
python -m pytest "packages/omnibias-graph/tests/test_cross_backend.py::test_laplacian_parity[random_walk_laplacian-0]" -q
Read the implementation
def random_walk_laplacian(adjacency: Array) -> Array:
    r"""Random-walk Laplacian ``L_rw = I - D^{-1} A``."""
    _check_adjacency(adjacency)
    d = adjacency.sum(axis=-1)
    d_inv = jnp.where(d > 0, 1.0 / d, jnp.zeros_like(d))
    n = adjacency.shape[0]
    return jnp.eye(n, dtype=adjacency.dtype) - d_inv[:, None] * adjacency
Other public import paths
  • omnibias.graph.jax.ops.random_walk_laplacian
  • omnibias.graph.jax.ops.spectral.random_walk_laplacian
FUNCTION

spectral_clustering_relaxation​

omnibias.graph.jax.ops.spectral_clustering_relaxationSource ↗
spectral_clustering_relaxation(adjacency: Array, k: int, *, normalized: bool=True) -> CutRelaxation

Continuous eigenvector relaxation of the k-way graph cut.

Parameters

NameType & defaultContract
adjacencypositional-or-keywordArrayRequiredSee the operation contract and implementation below.
kpositional-or-keywordintRequiredSee the operation contract and implementation below.
normalizedkeyword-onlyboolDefault: TrueSee the operation contract and implementation below.
ReturnsCutRelaxation
Explicit validation & failure cases
  • When k < 1, raises ValueError(f'k must be >= 1; got {k}').
  • When k > adjacency.shape[0], raises ValueError(f'k={k} exceeds the number of nodes {adjacency.shape[0]}').

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

Executed regression · differentiable through jax grad

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_differentiable_through_jax_grad() -> None:
    a = two_block_sbm(3, 1.0, 0.2)

    def cut(x):
        return G.spectral_clustering_relaxation(x, k=2, normalized=False).relaxed_cut_value

    grad = jax.grad(cut)(a)
    assert np.all(np.isfinite(np.asarray(grad)))
python -m pytest "packages/omnibias-graph/tests/test_spectral_jax.py::test_differentiable_through_jax_grad" -q
Regression example · cut relaxation value

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_cut_relaxation_value() -> None:
    a = two_block_sbm(4, 1.0, 0.1)
    rel = G.spectral_clustering_relaxation(a, k=2, normalized=True)
    evals, _ = G.laplacian_spectrum(a, normalized=True)
    assert abs(float(rel.relaxed_cut_value) - float(evals[:2].sum())) < 1e-10
python -m pytest packages/omnibias-graph/tests/test_spectral_jax.py -k test_cut_relaxation_value -q
Read the implementation
def spectral_clustering_relaxation(
    adjacency: Array, k: int, *, normalized: bool = True
) -> CutRelaxation:
    r"""Continuous eigenvector relaxation of the ``k``-way graph cut."""
    if k < 1:
        raise ValueError(f"k must be >= 1; got {k}")
    if k > adjacency.shape[0]:
        raise ValueError(f"k={k} exceeds the number of nodes {adjacency.shape[0]}")
    evals, evecs = laplacian_spectrum(adjacency, normalized=normalized)
    return CutRelaxation(
        embedding=evecs[:, :k],
        relaxed_cut_value=evals[:k].sum(),
        eigenvalues=evals[:k],
    )
Related APIsomnibias.graph.jax.ops.CutRelaxation↗omnibias.graph.jax.ops.laplacian_spectrum↗
Other public import paths
  • omnibias.graph.jax.ops.spectral_clustering_relaxation
  • omnibias.graph.jax.ops.spectral.spectral_clustering_relaxation
FUNCTION

spectral_embedding​

omnibias.graph.jax.ops.spectral_embeddingSource ↗
spectral_embedding(adjacency: Array, n_components: int, *, normalized: bool=True, drop_first: bool=True) -> Array

Laplacian-eigenmaps embedding: the smallest non-trivial eigenvectors.

Parameters

NameType & defaultContract
adjacencypositional-or-keywordArrayRequiredSee the operation contract and implementation below.
n_componentspositional-or-keywordintRequiredSee the operation contract and implementation below.
normalizedkeyword-onlyboolDefault: TrueSee the operation contract and implementation below.
drop_firstkeyword-onlyboolDefault: TrueSee the operation contract and implementation below.
ReturnsArray
Explicit validation & failure cases
  • When n_components < 1, raises ValueError(f'n_components must be >= 1; got {n_components}').
  • When start + n_components > n, raises ValueError(f'cannot take {n_components} components (drop_first={drop_first}) from an {n}-node graph').

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

Executed regression · embedding subspace projector 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_embedding_subspace_projector_parity() -> None:
    # A graph with a clear spectral gap so the smallest-k subspace is unique.
    n_per, k = 5, 2
    m = np.full((2 * n_per, 2 * n_per), 0.02)
    m[:n_per, :n_per] = 1.0
    m[n_per:, n_per:] = 1.0
    np.fill_diagonal(m, 0.0)
    yt = GT.spectral_embedding(torch.tensor(m), k, normalized=True, drop_first=False).detach().numpy()
    yj = np.asarray(GJ.spectral_embedding(jnp.asarray(m), k, normalized=True, drop_first=False))
    # Y Y^T (the projector onto the smallest-k eigenspace) is rotation-invariant.
    assert np.allclose(yt @ yt.T, yj @ yj.T, rtol=RTOL, atol=1e-9)
python -m pytest "packages/omnibias-graph/tests/test_cross_backend.py::test_embedding_subspace_projector_parity" -q
Read the implementation
def spectral_embedding(
    adjacency: Array,
    n_components: int,
    *,
    normalized: bool = True,
    drop_first: bool = True,
) -> Array:
    r"""Laplacian-eigenmaps embedding: the smallest non-trivial eigenvectors."""
    if n_components < 1:
        raise ValueError(f"n_components must be >= 1; got {n_components}")
    n = adjacency.shape[0]
    start = 1 if drop_first else 0
    if start + n_components > n:
        raise ValueError(
            f"cannot take {n_components} components (drop_first={drop_first}) "
            f"from an {n}-node graph"
        )
    _, evecs = laplacian_spectrum(adjacency, normalized=normalized)
    return evecs[:, start : start + n_components]
Related APIsomnibias.graph.jax.ops.laplacian_spectrum↗
Other public import paths
  • omnibias.graph.jax.ops.spectral_embedding
  • omnibias.graph.jax.ops.spectral.spectral_embedding
← Package overviewRunnable learning examples →