JaxGraph.from_numpy_graph

classmethod JaxGraph.from_numpy_graph(graph, device=None, dtype='float32')[source]

Convert a classical numpy graph to a jax.numpy format for GNN processing.

This method transforms all array-like attributes of a Graph object into their JAX equivalents, allowing efficient use with JAX transformations and accelerators.

Parameters:

  • graph (Graph) – A graph object containing NumPy arrays to convert.

  • device (Device | None) – Optional JAX device (e.g., CPU, GPU) to place the converted arrays on. If None, JAX uses the default device.

  • dtype (str) – Desired floating-point precision for converted arrays (e.g., “float32”, “float64”).

Returns:

A JAX-compatible version of the graph, ready for use in GNN pipelines.

Return type:

JaxGraph