mlx_graphs.algorithms.Node2Vec

class mlx_graphs.algorithms.Node2Vec(edge_index: mlx.core.array, embedding_dim: int, walk_length: int, context_size: int, num_nodes: int, walks_per_node: int = 1, p: float = 1.0, q: float = 1.0, num_negative_samples: int = 1, use_gpu: bool = True)

Bases: Module

The Node2Vec model from the “node2vec: Scalable Feature Learning for Networks” paper where random walks of length walk_length are sampled in a given graph, and node embeddings are learned via negative sampling optimization.

Note

This feature is experimental and requires mlx_cluster to be installed which currently requires mlx 0.18.

Parameters:

• edge_index (array) – The edge indices.

• embedding_dim (int) – The size of each embedding vector.

• walk_length (int) – The walk length.

• context_size (int) – The actual context size which is considered for positive samples. This parameter increases the effective sampling rate by reusing samples across different source nodes.

• num_nodes (int) – Number of nodes in a graph

• walks_per_node (int) – The number of walks to sample for each node. (default: 1)

• p (float) – Likelihood of immediately revisiting a node in the walk. (default: 1)

• q (float) – Control parameter to interpolate between breadth-first strategy and depth-first strategy (default: 1)

• num_negative_samples (int, optional) – The number of negative samples to use for each positive sample. (default: 1)

__call__(batch)

Returns the embeddings for the nodes in batch.

Methods

dataloader(batch_size)	Dataloader for nodes
loss(pos_array, neg_array)	Computes the loss given positive and negative random walks.
neg_sample(batch)
pos_sample(batch)

dataloader(batch_size)

Dataloader for nodes

loss(pos_array, neg_array)

Computes the loss given positive and negative random walks.