mlx_graphs.utils.transformations.to_edge_index

Contents

mlx_graphs.utils.transformations.to_edge_index#

mlx_graphs.utils.transformations.to_edge_index(adjacency_matrix: mlx.core.array, *, dtype: mlx.core.Dtype = mlx.core.uint32) → mlx.core.array[source]#

Converts an adjacency matrix to an edge index representation.

Parameters:

  • adjacency_matrix (array) – the input adjacency matrix

  • dtype (Dtype) – type of the output edge_index. Default to uint32.

Return type:

array

Returns:

A [2, num_edges] array representing the source and target nodes of each edge

Example:

matrix = mx.array(
    [
        [0, 1, 2],
        [3, 0, 0],
        [5, 1, 2],
    ]
)
edge_index = to_edge_index(matrix)
# mx.array([[0, 0, 1, 2, 2, 2], [1, 2, 0, 0, 1, 2]])
Contents