jraphx.utils

Utility functions and operations for graph processing and manipulation.

Scatter Operations

The scatter module provides efficient implementations of scatter operations for aggregating node features.

scatter_add

scatter_add(src: Array, index: Array, dim_size: int | None = None, dim: int = -2) → Array

Sums all values from the src tensor at the indices specified in the index tensor along a given dimension dim.

Uses JAX’s optimized segment_sum for better performance.

Note

dim_size determines the output shape and is therefore a static argument. When it is None it is inferred from index, which requires index to be concrete and hence is not available under jax.jit.

Parameters:

• src (Array) – The source tensor.

• index (Array) – The index tensor.

• dim_size (int | None, default: None) – The size of the output tensor at dimension dim. If set to None, will create a minimal-sized output tensor according to index.max() + 1. (default: None)

• dim (int, default: -2) – The dimension along which to index. (default: -2)

Returns:

Tensor with scattered values summed at each index.

Returns:

Array –

Scatter sum operation for aggregating values by index.

Example:

from jraphx.utils.scatter import scatter_add
import jax.numpy as jnp

src = jnp.array([1.0, 2.0, 3.0, 4.0])
index = jnp.array([0, 0, 1, 1])

# Sum values by index
out = scatter_add(src, index, dim_size=2)
# Result: [3.0, 7.0]

scatter_mean

scatter_mean(src: Array, index: Array, dim_size: int | None = None, dim: int = -2) → Array

Scatter mean operation - averages values from src at indices specified by index.

Empty segments are filled with zero.

Parameters:

• src (Array) – Source tensor to scatter

• index (Array) – Indices where to scatter

• dim_size (int | None, default: None) – Size of the output dimension, inferred from index if None (which requires a concrete index)

• dim (int, default: -2) – Dimension along which to scatter

Returns:

Array – Tensor with scattered values

Scatter mean operation for averaging values by index.

Example:

from jraphx.utils.scatter import scatter_mean
import jax.numpy as jnp

src = jnp.array([1.0, 2.0, 3.0, 4.0])
index = jnp.array([0, 0, 1, 1])

# Average values by index
out = scatter_mean(src, index, dim_size=2)
# Result: [1.5, 3.5]

scatter_max

scatter_max(src: Array, index: Array, dim_size: int | None = None, dim: int = -2, fill_value: float | None = None) → Array

Scatter max operation - takes maximum of values from src at indices specified by index.

Segments that receive no value are filled with fill_value. Emptiness is determined from index alone, so genuine infinities and NaNs present in src are propagated untouched and the output keeps src’s dtype.

Parameters:

• src (Array) – Source tensor to scatter

• index (Array) – Indices where to scatter

• dim_size (int | None, default: None) – Size of the output dimension, inferred from index if None (which requires a concrete index)

• dim (int, default: -2) – Dimension along which to scatter

• fill_value (float | None, default: None) – Value assigned to empty segments, cast to src.dtype. None means zero.

Returns:

Array – Tensor with scattered values

Scatter max operation for finding maximum values by index.

Example:

from jraphx.utils.scatter import scatter_max
import jax.numpy as jnp

src = jnp.array([1.0, 3.0, 2.0, 4.0])
index = jnp.array([0, 0, 1, 1])

# Find max values by index
out = scatter_max(src, index, dim_size=2)
# Result: [3.0, 4.0]

scatter_min

scatter_min(src: Array, index: Array, dim_size: int | None = None, dim: int = -2, fill_value: float | None = None) → Array

Scatter min operation - takes minimum of values from src at indices specified by index.

Segments that receive no value are filled with fill_value. Emptiness is determined from index alone, so genuine infinities and NaNs present in src are propagated untouched and the output keeps src’s dtype.

Parameters:

• src (Array) – Source tensor to scatter

• index (Array) – Indices where to scatter

• dim_size (int | None, default: None) – Size of the output dimension, inferred from index if None (which requires a concrete index)

• dim (int, default: -2) – Dimension along which to scatter

• fill_value (float | None, default: None) – Value assigned to empty segments, cast to src.dtype. None means zero.

Returns:

Array – Tensor with scattered values

Scatter min operation for finding minimum values by index.

scatter_std

scatter_std(src: Array, index: Array, dim_size: int | None = None, dim: int = -2, unbiased: bool = True) → Array

Scatter standard deviation - computes std of values at indices.

The deviations are accumulated around the per-segment mean, which avoids the catastrophic cancellation of the E[X^2] - E[X]^2 form. Segments holding fewer than two values (and empty segments) get a standard deviation of zero.

Parameters:

• src (Array) – Source tensor to scatter

• index (Array) – Indices where to scatter

• dim_size (int | None, default: None) – Size of the output dimension, inferred from index if None (which requires a concrete index)

• dim (int, default: -2) – Dimension along which to scatter

• unbiased (bool, default: True) – Whether to apply Bessel’s correction, i.e. divide by count - 1 instead of count

Returns:

Array – Tensor with scattered standard deviations

Scatter standard deviation operation. Bessel’s correction is applied by default (unbiased=True, dividing by count - 1); pass unbiased=False for the population standard deviation.

scatter

scatter(src: Array, index: Array, dim_size: int | None = None, dim: int = -2, reduce: str = 'add') → Array

Generic scatter operation using JAX’s optimized segment operations.

This function scatters values from src tensor at indices specified by index tensor, applying the specified reduction operation. Uses JAX’s built-in segment operations which are XLA-optimized for better performance on GPU/TPU.

Parameters:

• src (Array) – Source tensor to scatter [N, *]

• index (Array) – One-dimensional indices where to scatter [N], one per row of src

• dim_size (int | None, default: None) – Size of the output dimension, inferred from index if None (which requires a concrete index)

• dim (int, default: -2) – Dimension along which to scatter (default: -2, which maps to 0)

• reduce (str, default: 'add') – Reduction operation - “add” (alias “sum”), “mean”, “max”, “min”

Returns:

Array – Output tensor with scattered values [*, dim_size, *]

Generic scatter operation with configurable reduction.

Parameters:

• src: Source tensor to scatter

• index: Index tensor for scattering

• dim: Dimension to scatter along

• dim_size: Size of the output dimension

• reduce: Reduction operation (‘add’, with ‘sum’ as an alias, ‘mean’, ‘max’, ‘min’). Any other value raises a ValueError.

Graph Utilities

degree

degree(index: Array, num_nodes: int | None = None, dtype: str | type | dtype | None = None) → Array

Computes the (unweighted) degree of a given one-dimensional index tensor.

Parameters:

• index (Array) – Index tensor.

• num_nodes (int | None, default: None) – The number of nodes, i.e. the maximum entry of index plus one. (default: None)

• dtype (str | type | dtype | None, default: None) – The desired data type of the returned tensor; strings such as "int32" or "jnp.float16" are resolved with parse_dtype().

Returns:

Array – Node degrees, one entry per node.

Example

>>> import jax.numpy as jnp
>>> row = jnp.array([0, 1, 0, 2, 0])
>>> degree(row, dtype=jnp.int32)
Array([3, 1, 1], dtype=int32)

Compute the degree of each node in a graph.

Example:

from jraphx.utils import degree
import jax.numpy as jnp

edge_index = jnp.array([[0, 1, 2], [1, 2, 0]])

# Compute in-degree and out-degree
in_deg = degree(edge_index[1], num_nodes=3)
out_deg = degree(edge_index[0], num_nodes=3)

parse_dtype

parse_dtype(dtype: str | type | dtype) → type

Resolve a dtype spec to the matching jax.numpy scalar type (e.g. jnp.float32).

Parameters:

dtype (str | type | dtype) – A string naming a jax.numpy dtype, with or without a jnp. / jax.numpy. / np. / numpy. prefix ("float32", "jnp.bfloat16", "int32", …); an already-resolved scalar type such as jnp.float32; or a jax.numpy.dtype / numpy.dtype object. Non-string dtypes are normalized to the jax.numpy scalar type. The Python builtins float and int are rejected – name the width explicitly.

Returns:

type – The matching jax.numpy scalar type, e.g. jnp.float32.

Raises:

• TypeError – If dtype is neither a string nor a dtype-like value.

• ValueError – If dtype does not name a concrete jax.numpy dtype. Abstract categories such as "floating" or "integer" are rejected here rather than failing later at the first array construction.

Resolve a dtype spec – a plain or prefixed string, a scalar type, or a dtype object – to the matching jax.numpy scalar type. Every jraphx dtype argument accepts these specs, so a dtype can come straight from a configuration file.

Example:

from jraphx.utils import degree, parse_dtype
import jax.numpy as jnp

assert parse_dtype("float32") is jnp.float32
assert parse_dtype("jnp.bfloat16") is jnp.bfloat16

edge_index = jnp.array([[0, 1, 2], [1, 2, 0]])
deg = degree(edge_index[1], num_nodes=3, dtype="int32")

to_undirected

to_undirected(edge_index: Array, edge_attr: Array | None = None, num_nodes: int | None = None, reduce: str = 'add') → tuple[Array, Array | None]

Converts the graph given by edge_index to an undirected graph such that \((j,i) \in \mathcal{E}\) for every edge \((i,j) \in \mathcal{E}\).

The result is coalesced: the edge list is row-wise sorted, duplicated edges appear once, and their features are merged with reduce.

Parameters:

• edge_index (Array) – The edge indices.

• edge_attr (Array | None, default: None) – Edge weights or multi-dimensional edge features. (default: None)

• num_nodes (int | None, default: None) – The number of nodes, i.e. max_val + 1 of edge_index. (default: None)

• reduce (str, default: 'add') – The reduce operation to use for merging edge features ("add" / "sum", "mean", "min", "max"). (default: "add")

Returns:

tuple[Array, Array | None] – Tuple of (undirected edge_index, undirected edge_attr).

Note

The number of unique edges is data-dependent, so this function cannot be traced by jax.jit.

Convert a directed graph to undirected by adding reverse edges. The result is coalesced: the edge list is row-wise sorted, duplicated edges appear once, and their features are merged with reduce. The number of resulting edges is data-dependent, so this function cannot be traced by jax.jit().

Example:

from jraphx.utils import to_undirected
import jax.numpy as jnp

edge_index = jnp.array([[0, 1], [1, 2]])
edge_attr = jnp.array([[1.0], [2.0]])

# Convert to undirected
edge_index_undirected, edge_attr_undirected = to_undirected(
    edge_index, edge_attr
)

add_self_loops

add_self_loops(edge_index: Array, edge_attr: Array | None = None, fill_value: float | str = 1.0, num_nodes: int | None = None) → tuple[Array, Array | None]

Adds a self-loop \((i,i) \in \mathcal{E}\) to every node \(i \in \mathcal{V}\) in the graph given by edge_index. In case the graph is weighted or has multi-dimensional edge features (edge_attr is not None), edge features of self-loops will be added according to fill_value. One self-loop is appended per node unconditionally, so a node that already has a self-loop ends up with two; use add_remaining_self_loops() to add only the missing ones.

Parameters:

• edge_index (Array) – The edge indices.

• edge_attr (Array | None, default: None) – Edge weights or multi-dimensional edge features. (default: None)

• fill_value (Union[float, str], default: 1.0) – The way to generate edge features of self-loops. If float, edge features are set to this value. If str, edge features are computed by aggregating existing edge features that point to each node using the specified reduction (‘mean’, ‘add’ (alias ‘sum’), ‘max’, ‘min’). (default: 1.0)

• num_nodes (int | None, default: None) – The number of nodes, i.e. max_val + 1 of edge_index. (default: None)

Returns:

tuple[Array, Array | None] – Tuple of (edge_index with self-loops, edge_attr with self-loops).

Note

The output shape depends on num_nodes, so it must be given as a static integer under jax.jit. Inferring it from edge_index reads the array on the host.

Add self-loop edges to a graph.

Example:

from jraphx.utils import add_self_loops
import jax.numpy as jnp

edge_index = jnp.array([[0, 1], [1, 2]])

# Add self-loops
edge_index_with_loops, edge_attr = add_self_loops(
    edge_index, num_nodes=3
)

remove_self_loops

remove_self_loops(edge_index: Array, edge_attr: Array | None = None) → tuple[Array, Array | None]

Remove self-loops from edge indices.

Parameters:

• edge_index (Array) – Edge indices [2, num_edges]

• edge_attr (Array | None, default: None) – Optional edge attributes [num_edges, *]

Returns:

tuple[Array, Array | None] – Tuple of (edge_index without self-loops, edge_attr without self-loops)

Remove self-loop edges from a graph.

coalesce

coalesce(edge_index: Array, edge_attr: Array | None = None, num_nodes: int | None = None, reduce: str = 'add') → tuple[Array, Array | None]

Row-wise sorts edge_index and removes its duplicated entries. Duplicate entries in edge_attr are merged by scattering them together according to the given reduce option.

Parameters:

• edge_index (Array) – The edge indices.

• edge_attr (Array | None, default: None) – Edge weights or multi-dimensional edge features. (default: None)

• num_nodes (int | None, default: None) – The number of nodes, i.e. max_val + 1 of edge_index. Used only to validate edge_index. (default: None)

• reduce (str, default: 'add') – The reduce operation to use for merging edge features ("add" / "sum", "mean", "min", "max"). (default: "add")

Returns:

tuple[Array, Array | None] – Tuple of (coalesced edge_index, coalesced edge_attr).

Raises:

ValueError – If edge_index addresses a node outside of num_nodes.

Note

The number of unique edges is data-dependent, so this function cannot be traced by jax.jit.

Remove duplicate edges and optionally sum their attributes.

Example:

from jraphx.utils import coalesce
import jax.numpy as jnp

# Graph with duplicate edges
edge_index = jnp.array([[0, 0, 1], [1, 1, 2]])
edge_attr = jnp.array([[1.0], [2.0], [3.0]])

# Remove duplicates and sum attributes
edge_index_clean, edge_attr_clean = coalesce(
    edge_index, edge_attr, reduce='sum'
)

Conversion Utilities

to_dense_adj

to_dense_adj(edge_index: Array, edge_attr: Array | None = None, max_num_nodes: int | None = None) → Array

Convert edge indices to dense adjacency matrix.

Parallel edges are accumulated, so a duplicated edge contributes the sum of its attributes (or its multiplicity when edge_attr is None).

Parameters:

• edge_index (Array) – Edge indices [2, num_edges]

• edge_attr (Array | None, default: None) – Optional edge attributes [num_edges] or [num_edges, num_features]

• max_num_nodes (int | None, default: None) – Number of nodes of the dense output. Inferred from edge_index when None, which requires a concrete edge_index.

Returns:

Array – Dense adjacency matrix [num_nodes, num_nodes] or [num_nodes, num_nodes, num_features]

Convert edge indices to a dense adjacency matrix.

Example:

from jraphx.utils import to_dense_adj
import jax.numpy as jnp

edge_index = jnp.array([[0, 1, 2], [1, 2, 0]])

# Convert to dense adjacency matrix
adj = to_dense_adj(edge_index, max_num_nodes=3)

to_edge_index

to_edge_index(adj: Array) → tuple[Array, Array]

Convert adjacency matrix to edge indices.

An edge is emitted for every entry that is non-zero (for a feature tensor, for every entry with at least one non-zero feature), and its stored value is always returned as the edge attribute.

Parameters:

adj (Array) – Adjacency matrix [num_nodes, num_nodes] or [num_nodes, num_nodes, num_features]

Returns:

tuple[Array, Array] – Tuple of (edge_index [2, num_edges], edge_attr [num_edges] or [num_edges, num_features])

Note

The number of edges is data-dependent, so this function cannot be traced by jax.jit.

Convert adjacency representation to edge index format. Returns a (edge_index, edge_attr) tuple; the stored value of every non-zero entry is always returned as the edge attribute.