Jraph API¶

Portable jraph: GraphsTuple, batching/padding, and GraphNetwork models on NumPy / JAX / PyTorch / TF.

Segment ops on this module require num_segments (AnyTensor contract). None node/edge/global features are empty pytrees (jraph / jax.tree).

Narrative docs: Overview. Recipes: Examples.

anytensor.jraph ¶

Portable jraph: GraphsTuple, batching/padding, and GNN models on any tensor.

Public names match :mod:jraph. Segment ops require num_segments (AnyTensor contract). Feature nests use :mod:anytensor.tree (jax.tree API).

GraphsTuple ¶

Bases: NamedTuple

An ordered collection of graphs in a sparse format.

A port of :class:jraph.GraphsTuple. nodes, edges and globals may be None or an ArrayTree of features; senders / receivers are integer index arrays (or None when there are no edges). n_node and n_edge are integer vectors with one entry per graph in the batch.

Sender and receiver indices are absolute in the batched node array (offset by the nodes of earlier graphs). See the jraph docs for the full field layout.

__tree_batch__ / __tree_unbatch__ implement graph batching (not fieldwise array concat). :func:anytensor.tree.batch / unbatch are the same functions as :func:~anytensor.jraph.batch / unbatch; this type owns the offsetting logic. Custom node/edge/global objects may define the same methods so feature batching uses their logic.

Source code in anytensor/jraph/graph.py

class GraphsTuple(NamedTuple):
    """An ordered collection of graphs in a sparse format.

    A port of :class:`jraph.GraphsTuple`. ``nodes``, ``edges`` and ``globals``
    may be ``None`` or an ``ArrayTree`` of features; ``senders`` / ``receivers``
    are integer index arrays (or ``None`` when there are no edges). ``n_node``
    and ``n_edge`` are integer vectors with one entry per graph in the batch.

    Sender and receiver indices are **absolute** in the batched node array
    (offset by the nodes of earlier graphs). See the jraph docs for the
    full field layout.

    ``__tree_batch__`` / ``__tree_unbatch__`` implement graph batching (not
    fieldwise array concat). :func:`anytensor.tree.batch` / ``unbatch`` are
    the same functions as :func:`~anytensor.jraph.batch` / ``unbatch``; this
    type owns the offsetting logic. Custom node/edge/global objects may
    define the same methods so feature batching uses their logic.
    """

    nodes: Optional[ArrayTree]
    edges: Optional[ArrayTree]
    receivers: Optional[Any]
    senders: Optional[Any]
    globals: Optional[ArrayTree]
    n_node: Any
    n_edge: Any

    @classmethod
    def __tree_batch__(cls, xs, axis: int = 0):
        """Batch graphs. Senders/receivers are offset; not a fieldwise concat."""
        if axis != 0:
            raise ValueError("GraphsTuple batch only supports axis=0")
        from .utils import _batch_graphs

        return _batch_graphs(xs)

    def __tree_unbatch__(self, axis: int = 0):
        """Unbatch into one :class:`GraphsTuple` per graph (jraph ``unbatch``)."""
        if axis != 0:
            raise ValueError("GraphsTuple unbatch only supports axis=0")
        from .utils import _unbatch_graphs

        return _unbatch_graphs(self)

batch ¶

batch(trees, axis=0)

Batch a sequence of pytrees along axis (leading axis by default).

Same function as :func:anytensor.jraph.batch. If the type defines __tree_batch__(xs, axis=0), that method is used and children are not walked. Otherwise leaves are concatenated with :func:anytensor.concatenate. All-None stays None.

import numpy as np import anytensor.tree as tree tree.batch([np.array([1, 2]), np.array([3])]) array([1, 2, 3])

Source code in anytensor/tree.py

def batch(trees, axis: int = 0):
    """Batch a sequence of pytrees along ``axis`` (leading axis by default).

    Same function as :func:`anytensor.jraph.batch`. If the type defines
    ``__tree_batch__(xs, axis=0)``, that method is used and children are
    **not** walked. Otherwise leaves are concatenated with
    :func:`anytensor.concatenate`. All-``None`` stays ``None``.

    >>> import numpy as np
    >>> import anytensor.tree as tree
    >>> tree.batch([np.array([1, 2]), np.array([3])])
    array([1, 2, 3])
    """
    xs = list(trees)
    if not xs:
        raise ValueError("batch() requires at least one structure")
    return _batch_impl(xs, axis=axis)

unbatch ¶

unbatch(structure, axis=0)

Unbatch structure along axis into unit slices.

Same function as :func:anytensor.jraph.unbatch. __tree_unbatch__(axis=0) on the object wins (GraphsTuple yields one graph per n_node entry). Nested containers recurse. Top-level None cannot infer a batch size.

import numpy as np import anytensor.tree as tree head, tail = tree.unbatch(np.arange(1, 3)) tuple(int(x) for x in head), tuple(int(x) for x in tail) ((1,), (2,))

Source code in anytensor/tree.py

def unbatch(structure, axis: int = 0):
    """Unbatch ``structure`` along ``axis`` into unit slices.

    Same function as :func:`anytensor.jraph.unbatch`. ``__tree_unbatch__(axis=0)``
    on the object wins (``GraphsTuple`` yields one graph per ``n_node`` entry).
    Nested containers recurse. Top-level ``None`` cannot infer a batch size.

    >>> import numpy as np
    >>> import anytensor.tree as tree
    >>> head, tail = tree.unbatch(np.arange(1, 3))
    >>> tuple(int(x) for x in head), tuple(int(x) for x in tail)
    ((1,), (2,))
    """
    if structure is not None and hasattr(type(structure), "__tree_unbatch__"):
        return list(_call_unbatch(structure, axis))
    if structure is None:
        raise ValueError("Cannot unbatch None")
    n = _leading_length(structure, axis)
    return _unbatch_units(structure, n, axis)

GAT ¶

GAT(
    attention_query_fn,
    attention_logit_fn,
    node_update_fn=None,
)

Graph Attention Network layer (Veličković et al.). Expects self-edges.

Apply is @cache (same pattern as GraphNetwork). Destination size is :func:~anytensor.shape of nodes.

Source code in anytensor/jraph/models.py

def GAT(
    attention_query_fn: GATAttentionQueryFn,
    attention_logit_fn: GATAttentionLogitFn,
    node_update_fn: Optional[GATNodeUpdateFn] = None,
):
    """Graph Attention Network layer (Veličković et al.). Expects self-edges.

    Apply is ``@cache`` (same pattern as GraphNetwork). Destination size is
    :func:`~anytensor.shape` of nodes.
    """
    if node_update_fn is None:

        def node_update_fn(x):
            y = _leaky_relu(x)
            return reshape(y, (shape(y)[0], -1))

    @cache
    def _ApplyGAT(graph: GraphsTuple) -> GraphsTuple:
        nodes, edges, receivers, senders, _, _, _ = graph
        if nodes is None:
            raise IndexError("GAT requires node features")
        sum_n_node = shape(nodes)[0]
        nodes = attention_query_fn(nodes)
        sent_attributes = take(nodes, senders)
        received_attributes = take(nodes, receivers)
        softmax_logits = attention_logit_fn(sent_attributes, received_attributes, edges)
        weights = utils.segment_softmax(
            softmax_logits, segment_ids=receivers, num_segments=sum_n_node
        )
        messages = sent_attributes * weights
        nodes = utils.segment_sum(messages, receivers, num_segments=sum_n_node)
        nodes = node_update_fn(nodes)
        return graph._replace(nodes=nodes)

    return _ApplyGAT

DeepSets ¶

DeepSets(
    update_node_fn,
    update_global_fn,
    aggregate_nodes_for_globals_fn=segment_sum,
)

DeepSets layer as a configured GraphNetwork.

Source code in anytensor/jraph/models.py

def DeepSets(
    update_node_fn: Callable[[NodeFeatures, Globals], NodeFeatures],
    update_global_fn: Callable[[NodeFeatures], Globals],
    aggregate_nodes_for_globals_fn: AggregateNodesToGlobalsFn = utils.segment_sum,
):
    """DeepSets layer as a configured GraphNetwork."""
    return GraphNetwork(
        update_edge_fn=None,
        update_node_fn=lambda n, s, r, g: update_node_fn(n, g),
        update_global_fn=lambda n, e, g: update_global_fn(n),
        aggregate_nodes_for_globals_fn=aggregate_nodes_for_globals_fn,
    )

GraphConvolution ¶

GraphConvolution(
    update_node_fn,
    aggregate_nodes_fn=segment_sum,
    add_self_edges=False,
    symmetric_normalization=True,
)

GCN layer (Kipf & Welling). No activation after aggregation.

Apply uses the public cache pattern: @cache plus :meth:~anytensor.cache.lookup / :meth:~anytensor.cache.store on "gcn" keyed by senders and the constructor flags so stacked applies and ONNX do not duplicate Shape / Range / Concat.

Source code in anytensor/jraph/models.py

def GraphConvolution(
    update_node_fn: Callable[[NodeFeatures], NodeFeatures],
    aggregate_nodes_fn: AggregateEdgesToNodesFn = utils.segment_sum,
    add_self_edges: bool = False,
    symmetric_normalization: bool = True,
):
    """GCN layer (Kipf & Welling). No activation after aggregation.

    Apply uses the public cache pattern: ``@cache`` plus
    :meth:`~anytensor.cache.lookup` / :meth:`~anytensor.cache.store` on
    ``"gcn"`` keyed by senders and the constructor flags so stacked applies
    and ONNX do not duplicate ``Shape`` / ``Range`` / ``Concat``.
    """

    @cache
    def _ApplyGCN(graph: GraphsTuple) -> GraphsTuple:
        nodes, _, receivers, senders, _, n_node, _ = graph
        nodes = update_node_fn(nodes)
        conv_senders, conv_receivers, total_num_nodes, sender_degree, receiver_degree = (
            _gcn_structure(
                nodes,
                senders,
                receivers,
                n_node,
                add_self_edges=add_self_edges,
                symmetric_normalization=symmetric_normalization,
            )
        )
        if symmetric_normalization:
            one = at_ones((), dtype=tree.leaves(nodes)[0].dtype, like=tree.leaves(nodes)[0])
            nodes = tree.map(
                lambda x: x
                * reshape(
                    rsqrt(maximum(sender_degree, one)),
                    (shape(sender_degree)[0],) + (1,) * (x.ndim - 1),
                ),
                nodes,
            )
            nodes = tree.map(
                lambda x: aggregate_nodes_fn(
                    take(x, conv_senders), conv_receivers, total_num_nodes
                ),
                nodes,
            )
            nodes = tree.map(
                lambda x: x
                * reshape(
                    rsqrt(maximum(receiver_degree, one)),
                    (shape(receiver_degree)[0],) + (1,) * (x.ndim - 1),
                ),
                nodes,
            )
        else:
            nodes = tree.map(
                lambda x: aggregate_nodes_fn(
                    take(x, conv_senders), conv_receivers, total_num_nodes
                ),
                nodes,
            )
        return graph._replace(nodes=nodes)

    return _ApplyGCN

GraphMapFeatures ¶

GraphMapFeatures(
    embed_edge_fn=None,
    embed_node_fn=None,
    embed_global_fn=None,
)

Embed nodes, edges, and globals independently.

Apply is @cache (same pattern as GraphNetwork).

Source code in anytensor/jraph/models.py

def GraphMapFeatures(
    embed_edge_fn: Optional[EmbedEdgeFn] = None,
    embed_node_fn: Optional[EmbedNodeFn] = None,
    embed_global_fn: Optional[EmbedGlobalFn] = None,
):
    """Embed nodes, edges, and globals independently.

    Apply is ``@cache`` (same pattern as GraphNetwork).
    """
    identity = lambda x: x
    embed_edges_fn = embed_edge_fn if embed_edge_fn else identity
    embed_nodes_fn = embed_node_fn if embed_node_fn else identity
    embed_globals_fn = embed_global_fn if embed_global_fn else identity

    @cache
    def Embed(graphs_tuple: GraphsTuple) -> GraphsTuple:
        return graphs_tuple._replace(
            nodes=embed_nodes_fn(graphs_tuple.nodes),
            edges=embed_edges_fn(graphs_tuple.edges),
            globals=embed_globals_fn(graphs_tuple.globals),
        )

    return Embed

GraphNetGAT ¶

GraphNetGAT(
    update_edge_fn,
    update_node_fn,
    attention_logit_fn,
    attention_reduce_fn,
    update_global_fn=None,
    aggregate_edges_for_nodes_fn=segment_sum,
    aggregate_nodes_for_globals_fn=segment_sum,
    aggregate_edges_for_globals_fn=segment_sum,
)

GraphNet with required attention on edge features.

Source code in anytensor/jraph/models.py

def GraphNetGAT(
    update_edge_fn: GNUpdateEdgeFn,
    update_node_fn: GNUpdateNodeFn,
    attention_logit_fn: AttentionLogitFn,
    attention_reduce_fn: AttentionReduceFn,
    update_global_fn: Optional[GNUpdateGlobalFn] = None,
    aggregate_edges_for_nodes_fn: AggregateEdgesToNodesFn = utils.segment_sum,
    aggregate_nodes_for_globals_fn: AggregateNodesToGlobalsFn = utils.segment_sum,
    aggregate_edges_for_globals_fn: AggregateEdgesToGlobalsFn = utils.segment_sum,
):
    """GraphNet with required attention on edge features."""
    if (attention_logit_fn is None) or (attention_reduce_fn is None):
        raise ValueError(
            "`None` value not supported for `attention_logit_fn` or "
            "`attention_reduce_fn` in a Graph Attention network."
        )
    return GraphNetwork(
        update_edge_fn=update_edge_fn,
        update_node_fn=update_node_fn,
        update_global_fn=update_global_fn,
        attention_logit_fn=attention_logit_fn,
        attention_reduce_fn=attention_reduce_fn,
        aggregate_edges_for_nodes_fn=aggregate_edges_for_nodes_fn,
        aggregate_nodes_for_globals_fn=aggregate_nodes_for_globals_fn,
        aggregate_edges_for_globals_fn=aggregate_edges_for_globals_fn,
    )

GraphNetwork ¶

GraphNetwork(
    update_edge_fn,
    update_node_fn,
    update_global_fn=None,
    aggregate_edges_for_nodes_fn=segment_sum,
    aggregate_nodes_for_globals_fn=segment_sum,
    aggregate_edges_for_globals_fn=segment_sum,
    attention_logit_fn=None,
    attention_normalize_fn=segment_softmax,
    attention_reduce_fn=None,
)

Returns a method that applies a configured GraphNetwork.

Follows Algorithm 1 of https://arxiv.org/abs/1806.01261, with separate sender/receiver aggregations and optional softmax attention. Same call signature as :func:jraph.GraphNetwork. Apply uses the public cache pattern: @cache (sticky) plus :func:~anytensor.partition_ids (cache.lookup / store on "partition", keyed by n_node / n_edge). Callers write the same @cache apply.

Flattened totals (official sum_n_node / sum_n_edge) are :func:~anytensor.shape of the node / sender axis — not sum(n_node) — so they stay symbolic sizes on ONNX export.

Source code in anytensor/jraph/models.py

def GraphNetwork(
    update_edge_fn: Optional[GNUpdateEdgeFn],
    update_node_fn: Optional[GNUpdateNodeFn],
    update_global_fn: Optional[GNUpdateGlobalFn] = None,
    aggregate_edges_for_nodes_fn: AggregateEdgesToNodesFn = utils.segment_sum,
    aggregate_nodes_for_globals_fn: AggregateNodesToGlobalsFn = utils.segment_sum,
    aggregate_edges_for_globals_fn: AggregateEdgesToGlobalsFn = utils.segment_sum,
    attention_logit_fn: Optional[AttentionLogitFn] = None,
    attention_normalize_fn: Optional[AttentionNormalizeFn] = utils.segment_softmax,
    attention_reduce_fn: Optional[AttentionReduceFn] = None,
):
    """Returns a method that applies a configured GraphNetwork.

    Follows Algorithm 1 of https://arxiv.org/abs/1806.01261, with separate
    sender/receiver aggregations and optional softmax attention. Same call
    signature as :func:`jraph.GraphNetwork`. Apply uses the public cache
    pattern: ``@cache`` (sticky) plus :func:`~anytensor.partition_ids`
    (``cache.lookup`` / ``store`` on ``"partition"``, keyed by ``n_node`` /
    ``n_edge``). Callers write the same ``@cache`` apply.

    Flattened totals (official ``sum_n_node`` / ``sum_n_edge``) are
    :func:`~anytensor.shape` of the node / sender axis — not
    ``sum(n_node)`` — so they stay symbolic sizes on ONNX export.
    """
    not_both_supplied = lambda x, y: (x != y) and ((x is None) or (y is None))
    if not_both_supplied(attention_reduce_fn, attention_logit_fn):
        raise ValueError(
            "attention_logit_fn and attention_reduce_fn must both be supplied."
        )

    @cache
    def _ApplyGraphNet(graph: GraphsTuple) -> GraphsTuple:
        nodes, edges, receivers, senders, globals_, n_node, n_edge = graph
        node_leaves = tree.leaves(nodes)
        # Official jraph uses ``sum(n_node)`` / ``sum(n_edge)``. That is a data
        # reduction (``ReduceSum`` in ONNX) and ``int()`` of it bakes a host
        # constant. ``shape(nodes)[0]`` is the same integer when the
        # GraphsTuple invariant holds, and is a dim_param on export.
        if node_leaves:
            sum_n_node = shape(node_leaves[0])[0]
        else:
            # No node tensor to read. Eager-only host int (cannot be a
            # symbolic node axis).
            sum_n_node = int(np_sum_n_node(n_node))
        sum_n_edge = 0 if senders is None else shape(senders)[0]
        # ``int(size)`` is rewritten by TF Autograph into a graph op, so a
        # symbolic leading dim looks "concrete". Only compare nest lengths
        # when the size is already a Python int (eager NumPy / JAX / TF).
        if (
            node_leaves
            and type(sum_n_node) is int
            and not utils._tree_all(  # noqa: SLF001
                tree.map(lambda n: n.shape[0] == sum_n_node, nodes)
            )
        ):
            raise ValueError(
                "All node arrays in nest must contain the same number of nodes."
            )

        sent_attributes = _take_index(nodes, senders)
        received_attributes = _take_index(nodes, receivers)
        if globals_ is not None and n_edge is not None and sum_n_edge is not None:
            global_edge_attributes = tree.map(
                lambda g: _repeat_by(g, n_edge, sum_n_edge), globals_
            )
        else:
            global_edge_attributes = None

        if update_edge_fn:
            edges = update_edge_fn(
                edges, sent_attributes, received_attributes, global_edge_attributes
            )

        if attention_logit_fn:
            logits = attention_logit_fn(
                edges, sent_attributes, received_attributes, global_edge_attributes
            )
            tree_calculate_weights = functools.partial(
                attention_normalize_fn, segment_ids=receivers, num_segments=sum_n_node
            )
            weights = tree.map(tree_calculate_weights, logits)
            edges = attention_reduce_fn(edges, weights)

        if update_node_fn:
            sent_attributes = (
                None
                if edges is None
                else tree.map(
                    lambda e: aggregate_edges_for_nodes_fn(e, senders, sum_n_node), edges
                )
            )
            received_attributes = (
                None
                if edges is None
                else tree.map(
                    lambda e: aggregate_edges_for_nodes_fn(e, receivers, sum_n_node),
                    edges,
                )
            )
            if globals_ is not None:
                global_attributes = tree.map(
                    lambda g: _repeat_by(g, n_node, sum_n_node), globals_
                )
            else:
                global_attributes = None
            nodes = update_node_fn(
                nodes, sent_attributes, received_attributes, global_attributes
            )

        if update_global_fn:
            n_graph = shape(n_node)[0]
            node_gr_idx = partition_ids(n_node, sum_n_node)
            edge_gr_idx = (
                None
                if senders is None
                else partition_ids(n_edge, sum_n_edge)
            )
            node_attributes = (
                None
                if nodes is None
                else tree.map(
                    lambda n: aggregate_nodes_for_globals_fn(n, node_gr_idx, n_graph),
                    nodes,
                )
            )
            edge_attributes = (
                None
                if edges is None or edge_gr_idx is None
                else tree.map(
                    lambda e: aggregate_edges_for_globals_fn(e, edge_gr_idx, n_graph),
                    edges,
                )
            )
            globals_ = update_global_fn(node_attributes, edge_attributes, globals_)
        return GraphsTuple(
            nodes=nodes,
            edges=edges,
            receivers=receivers,
            senders=senders,
            globals=globals_,
            n_node=n_node,
            n_edge=n_edge,
        )

    return _ApplyGraphNet

InteractionNetwork ¶

InteractionNetwork(
    update_edge_fn,
    update_node_fn,
    aggregate_edges_for_nodes_fn=segment_sum,
    include_sent_messages_in_node_update=False,
)

Interaction network (Battaglia et al.) as a configured GraphNetwork.

Source code in anytensor/jraph/models.py

def InteractionNetwork(
    update_edge_fn: InteractionUpdateEdgeFn,
    update_node_fn: Union[InteractionUpdateNodeFn, InteractionUpdateNodeFnNoSentEdges],
    aggregate_edges_for_nodes_fn: AggregateEdgesToNodesFn = utils.segment_sum,
    include_sent_messages_in_node_update: bool = False,
):
    """Interaction network (Battaglia et al.) as a configured GraphNetwork."""
    wrapped_update_edge_fn = lambda e, s, r, g: update_edge_fn(e, s, r)
    if include_sent_messages_in_node_update:
        wrapped_update_node_fn = lambda n, s, r, g: update_node_fn(n, s, r)
    else:
        wrapped_update_node_fn = lambda n, s, r, g: update_node_fn(n, r)
    return GraphNetwork(
        update_edge_fn=wrapped_update_edge_fn,
        update_node_fn=wrapped_update_node_fn,
        aggregate_edges_for_nodes_fn=aggregate_edges_for_nodes_fn,
    )

RelationNetwork ¶

RelationNetwork(
    update_edge_fn,
    update_global_fn,
    aggregate_edges_for_globals_fn=segment_sum,
)

Relation network as a configured GraphNetwork.

Source code in anytensor/jraph/models.py

def RelationNetwork(
    update_edge_fn: Callable[[SenderFeatures, ReceiverFeatures], EdgeFeatures],
    update_global_fn: Callable[[EdgeFeatures], NodeFeatures],
    aggregate_edges_for_globals_fn: AggregateEdgesToGlobalsFn = utils.segment_sum,
):
    """Relation network as a configured GraphNetwork."""
    return GraphNetwork(
        update_edge_fn=lambda e, s, r, g: update_edge_fn(s, r),
        update_node_fn=None,
        update_global_fn=lambda n, e, g: update_global_fn(e),
        attention_logit_fn=None,
        aggregate_edges_for_globals_fn=aggregate_edges_for_globals_fn,
    )