{"id":28719152,"url":"https://github.com/0xsooki/extending-jax","last_synced_at":"2026-04-25T11:31:59.997Z","repository":{"id":296021433,"uuid":"992068323","full_name":"0xSooki/extending-jax","owner":"0xSooki","description":"JAX Custom Operations with C++ and CUDA (using Pybind11)","archived":false,"fork":false,"pushed_at":"2025-05-28T15:53:26.000Z","size":20,"stargazers_count":0,"open_issues_count":0,"forks_count":0,"subscribers_count":0,"default_branch":"main","last_synced_at":"2025-05-28T16:36:08.098Z","etag":null,"topics":["cuda","jax","pybind11","xla"],"latest_commit_sha":null,"homepage":"","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":null,"status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/0xSooki.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":null,"code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null,"zenodo":null}},"created_at":"2025-05-28T15:09:52.000Z","updated_at":"2025-05-28T15:54:33.000Z","dependencies_parsed_at":"2025-05-28T16:38:43.875Z","dependency_job_id":"32e96a70-b376-461b-9499-0cbc640521b8","html_url":"https://github.com/0xSooki/extending-jax","commit_stats":null,"previous_names":["0xsooki/extending-jax"],"tags_count":0,"template":false,"template_full_name":null,"purl":"pkg:github/0xSooki/extending-jax","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/0xSooki%2Fextending-jax","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/0xSooki%2Fextending-jax/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/0xSooki%2Fextending-jax/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/0xSooki%2Fextending-jax/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/0xSooki","download_url":"https://codeload.github.com/0xSooki/extending-jax/tar.gz/refs/heads/main","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/0xSooki%2Fextending-jax/sbom","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":259929946,"owners_count":22933527,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["cuda","jax","pybind11","xla"],"created_at":"2025-06-15T06:00:18.295Z","updated_at":"2026-04-25T11:31:59.957Z","avatar_url":"https://github.com/0xSooki.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"# JAX Custom Operations with Pybind11, C++ and CUDA\n\nA demonstration project showing how to extend JAX with custom C++ and CUDA operations using the XLA FFI (Foreign Function Interface). This project implements a simple mathematical operation `foo(a, b) = sum(a * (b + 1))` with both CPU and GPU kernels, complete with gradient support.\n\n[![PyPI version](https://img.shields.io/pypi/v/permanentboost)](https://pypi.org/project/permanentboost/)  \n[![Build](https://github.com/0xSooki/permanent-boost/actions/workflows/tests.yml/badge.svg)](https://github.com/0xSooki/permanent-boost/actions)  \n[![License: MIT](https://img.shields.io/badge/license-MIT-blue.svg)](LICENSE)\n\n## Features\n\n- ✅ **Custom JAX operations** with C++ and CUDA implementations\n- ✅ **Automatic differentiation** support via `jax.grad`\n- ✅ **CPU and GPU kernels** with automatic fallback\n- ✅ **JAX transformations** (JIT, vmap, etc.) compatibility\n- ✅ **Comprehensive test suite** with 20+ test cases\n- ✅ **Scalar output** for gradient computation\n- ✅ **Modern CMake build system** with scikit-build-core\n\n## Mathematical Operation\n\nThe implemented operation computes:\n\n```\nf(a, b) = sum(a * (b + 1))\n```\n\nWhere:\n\n- `a`, `b` are input tensors of the same shape\n- The output is a scalar (sum of element-wise products)\n- Gradients: `∇_a f = b + 1`, `∇_b f = a`\n\n## Quick Start\n\n### Installation\n\n```bash\n# Clone the repository\ngit clone \u003crepository-url\u003e\ncd permanent-boost\n\n# Install dependencies\npip install -r requirements.txt\n\n# Build and install the package\npip install -e .\n```\n\n### Basic Usage\n\n```python\nimport jax.numpy as jnp\nimport jax\nfrom sooki import foo\n\n# Create input tensors (must be float32)\na = jnp.array([1.0, 2.0], dtype=jnp.float32)\nb = jnp.array([3.0, 4.0], dtype=jnp.float32)\n\n# Compute the operation\nresult = foo(a, b)\nprint(f\"foo(a, b) = {result}\")  # Output: 14.0\n\n# Compute gradients\ngrad_fn = jax.grad(foo, argnums=0)\ngrad_a = grad_fn(a, b)\nprint(f\"∇_a foo = {grad_a}\")  # Output: [4.0, 5.0] = b + 1\n```\n\n### Advanced Usage\n\n```python\n# Works with JAX transformations\njit_foo = jax.jit(foo)\nresult = jit_foo(a, b)\n\n# Vectorized operations\nbatch_a = jnp.array([[1.0, 2.0], [3.0, 4.0]], dtype=jnp.float32)\nbatch_b = jnp.array([[5.0, 6.0], [7.0, 8.0]], dtype=jnp.float32)\nbatch_result = jax.vmap(foo)(batch_a, batch_b)\n\n# Value and gradient simultaneously\nvalue, grad = jax.value_and_grad(foo, argnums=0)(a, b)\n```\n\n## Project Structure\n\n```\n├── src/\n│   ├── cpu_ops.hpp         # CPU kernel implementations\n│   ├── gpu_ops.cc          # GPU FFI bindings\n│   ├── kernels.cc.cu       # CUDA kernel implementations\n│   ├── kernels.h           # GPU function declarations\n│   ├── main.cpp            # CPU FFI bindings\n│   └── sooki/\n│       ├── __init__.py     # Package initialization\n│       └── ops.py          # Python interface and custom VJP\n├── tests/\n│   └── test_foo.py         # Comprehensive test suite\n├── CMakeLists.txt          # Build configuration\n├── pyproject.toml          # Python package configuration\n└── README.md               # This file\n```\n\n## Implementation Details\n\n### Python Interface\n\n- Custom VJP (Vector-Jacobian Product) implementation\n- Automatic CPU/GPU dispatch based on availability\n- Integration with JAX's transformation system\n\n## Testing\n\nRun the comprehensive test suite:\n\n```bash\n# Run all tests\npytest tests/test_foo.py -v\n\n# Run specific test categories\npytest tests/test_foo.py::TestFooGradients -v\npytest tests/test_foo.py::TestFooJAXTransformations -v\n```\n\nTest coverage includes:\n\n- ✅ Basic functionality\n- ✅ Gradient computation accuracy\n- ✅ JAX transformations (JIT, vmap, value_and_grad)\n- ✅ Error handling and edge cases\n- ✅ Performance and consistency\n- ✅ Mathematical correctness verification\n\n## Requirements\n\n### System Requirements\n\n- Python 3.8+\n- CMake 3.15+\n- C++14 compatible compiler\n- CUDA Toolkit (optional, for GPU support)\n\n### Python Dependencies\n\n- JAX \u003e= 0.4.31\n- JAXlib \u003e= 0.4.31\n- NumPy\n- pybind11\n\n### Build Dependencies\n\n- scikit-build-core\n- ninja (build system)\n\n## GPU Support\n\nGPU kernels are automatically compiled if CUDA is available. The package gracefully falls back to CPU-only mode if:\n\n- CUDA toolkit is not installed\n- No CUDA-capable GPU is detected\n- CUDA compilation fails\n\nCheck GPU support:\n\n```python\nimport sooki\nprint(\"GPU support:\", hasattr(sooki, 'gpu_ops'))\n```\n\n## Contributing\n\n1. Fork the repository\n2. Create a feature branch\n3. Add tests for new functionality\n4. Ensure all tests pass: `pytest tests/`\n5. Submit a pull request\n\n## License\n\nThis project demonstrates JAX extension techniques and is provided for educational purposes.\n\n## Acknowledgments\n\nThis project was inspired by and builds upon the excellent tutorial and examples from:\n\n**[dfm/extending-jax](https://github.com/dfm/extending-jax)** - A comprehensive guide to extending JAX with custom operations\n\nThe original repository by Dan Foreman-Mackey provides foundational examples and best practices for JAX extensions that were instrumental in developing this project.\n\n## References\n\n- [JAX Custom Operations Guide](https://jax.readthedocs.io/en/latest/Custom_Operation_for_GPUs.html)\n- [XLA FFI Documentation](https://github.com/google/jax/tree/main/jaxlib/xla_extension)\n- [JAX Autodiff Cookbook](https://jax.readthedocs.io/en/latest/notebooks/autodiff_cookbook.html)\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2F0xsooki%2Fextending-jax","html_url":"https://awesome.ecosyste.ms/projects/github.com%2F0xsooki%2Fextending-jax","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2F0xsooki%2Fextending-jax/lists"}