metagraph-numba Documentation

metagraph-numba is a plugin for Metagraph that enables algorithms to be just-in-time (JIT) compiled with Numba. This package currently does not include any graph algorithm implementations in Numba, and is mostly used to support the metagraph-mlir plugin. To learn more about compiler plugins in Metagraph, see the Compiler Plugins section of the Metagraph Plugin Author Guide.

metagraph-numba is licensed under the Apache 2.0 license and the source code can be found on Github.

Installation

metagraph-numba is currently only distributed via conda. To install:

conda install -c metagraph -c conda-forge metagraph-numba

Implementing Algorithms with metagraph-numba

The metagraph-numba compiler assumes that the decorated algorithm function is compilable with Numba directly. For example, suppose an abstract algorithm has already been defined to add two vectors:

@abstract_algorithm("example.add")
def example_add(a: Vector, b: Vector) -> Vector:
    pass

A concrete implementation of this algorithm can be written this way:

@concrete_algorithm("example.add", compiler="numba")
def compiled_add(a: NumpyVectorType, b: NumpyVectorType) -> NumpyVectorType:
   return a + b

Numba has support for the input and output types, which are 1D NumPy arrays, and the body of the function, which does array addition. For more details on what Numba supports, see the Numba Supported Python Features and Numba Supported NumPy Features.

Internally, the compiler wraps each algorithm function in a numba.jit(inline="always") call. This pulls every algorithm body in the task subgraph into the generated wrapper function, allowing for more optimization opportunies. Note that the resulting function is not JIT-compiled until the first call, which allows for further specialization on argument attributes like the dtypes of input arrays.