haiku_geometric.transforms#

add_self_loops

Adds self edges.

remove_self_loops

Removes self loops from a graph.

normalize_features

Normalizes node features to sum to 1 on the last axis.
haiku_geometric.transforms.add_self_loops(senders, receivers, edge_attr=None, fill_value=None, num_nodes=None)[source]#

Adds self edges. Assumes self edges are not in the graph yet.

Parameters

  • senders (jnp.ndarray) – Array of sender node indices.

  • receivers (jnp.ndarray) – Array of receiver node indices.

  • edge_attr (jnp.ndarray) – Array of edge attributes or edge weights. (default: None)

  • fill_value (float) – Value to fill the edge features with. (default: None)

  • num_nodes (int) – Total number of nodes in the graph. (default: None)

Return type

Tuple[ndarray, ndarray]

haiku_geometric.transforms.normalize_features(x)[source]#

Normalizes node features to sum to 1 on the last axis.

Parameters

x (jnp.ndarray) – Array of node features.

Return type

ndarray

haiku_geometric.transforms.remove_self_loops(senders, receivers, edge_weight)[source]#

Removes self loops from a graph.