Extending JAX with custom Rust code

This repository demonstrate how to add custom XLA ops to JAX with custom Rust code. The codes are based on demonstration using c++ and CUDA.

Running the demonstration

This repository uses Rye to manage python environment and Pytest to manage tests. First of all, compile Rust codes and set up python environment:

# rye sync

Then, run tests:

# rye run pytest

Name		Name	Last commit message	Last commit date
Latest commit History 62 Commits
.github		.github
lib		lib
python		python
src		src
tests		tests
.clang-format		.clang-format
.gitattributes		.gitattributes
.gitignore		.gitignore
CMakeLists.txt		CMakeLists.txt
Cargo.lock		Cargo.lock
Cargo.toml		Cargo.toml
LICENSE		LICENSE
MANIFEST.in		MANIFEST.in
README.md		README.md
demo.ipynb		demo.ipynb
pyproject.toml		pyproject.toml
requirements-dev.lock		requirements-dev.lock
requirements.lock		requirements.lock