Skip to main content
API REFERENCE/omnibias-graph

omnibias.graph.arrangement.jax

0.1.0a2AlphaPython >=3.10AGPL / commercial

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

import omnibias.graph.arrangement.jax
Module contract & mathematical conventions

Face-Net (JAX twin; theory 02-02).

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

facenet_apply​

omnibias.graph.arrangement.jax.facenet_applySource ↗
facenet_apply(graph: ArrangementGraph, arr: Arrangement, *, lift_w: Array, lift_b: Array, msg_w: Array, msg_b: Array, out_w: Array, out_b: Array, rounds: int, beta: float) -> Array

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

Parameters

NameType & defaultContract
graphpositional-or-keywordArrangementGraphRequiredSee the operation contract and implementation below.
arrpositional-or-keywordArrangementRequiredSee the operation contract and implementation below.
lift_wkeyword-onlyArrayRequiredSee the operation contract and implementation below.
lift_bkeyword-onlyArrayRequiredSee the operation contract and implementation below.
msg_wkeyword-onlyArrayRequiredSee the operation contract and implementation below.
msg_bkeyword-onlyArrayRequiredSee the operation contract and implementation below.
out_wkeyword-onlyArrayRequiredSee the operation contract and implementation below.
out_bkeyword-onlyArrayRequiredSee the operation contract and implementation below.
roundskeyword-onlyintRequiredSee the operation contract and implementation below.
betakeyword-onlyfloatRequiredSee the operation contract and implementation below.
ReturnsArray
Read the implementation
def facenet_apply(
    graph: ArrangementGraph,
    arr: Arrangement,
    *,
    lift_w: Array,
    lift_b: Array,
    msg_w: Array,
    msg_b: Array,
    out_w: Array,
    out_b: Array,
    rounds: int,
    beta: float,
) -> Array:
    feats = jnp.asarray(node_features(arr, graph, beta=beta), dtype=lift_w.dtype)
    h = jnp.tanh(feats @ lift_w.T + lift_b)
    n = h.shape[0]
    adj = jnp.zeros((n, n), dtype=h.dtype)
    crossed = jnp.zeros((n, n), dtype=h.dtype)
    for u, v, k in graph.edges:
        adj = adj.at[u, v].set(1.0).at[v, u].set(1.0)
        crossed = crossed.at[u, v].set(float(k)).at[v, u].set(float(k))
    for _ in range(int(rounds)):
        msgs = []
        for i in range(n):
            acc = jnp.zeros_like(h[i])
            deg = 0.0
            for j in range(n):
                if float(adj[i, j]) <= 0.0:
                    continue
                cat = jnp.concatenate([h[j], crossed[i, j].reshape(1)])
                acc = acc + jnp.tanh(cat @ msg_w.T + msg_b)
                deg += 1.0
            msgs.append(acc / max(deg, 1.0))
        h = h + jnp.stack(msgs)
    return (h @ out_w.T + out_b).reshape(-1)
Related APIsomnibias.graph.arrangement.node_features↗
← Package overviewRunnable learning examples →