Convert to KLUJAX to new FFI API

.github/run_tests.py

@@ -0,0 +1,18 @@
+import os
+import platform
+import sys
+from subprocess import call
+
+PROJECT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+TESTS_PATH = os.path.join(PROJECT, "tests.py")
+
+print(f"{PROJECT=}", os.path.exists(PROJECT))
+print(f"{TESTS_PATH=}", os.path.exists(PROJECT))
+
+if sys.platform == "darwin":
+    architecture = platform.machine()
+    print(f"{architecture=}")
+    if architecture != "arm64":
+        exit(print("skipping tests as we only run them on arm64."))
+
+exit(call(["pytest", "-s", TESTS_PATH]))

.github/workflows/main.yml

@@ -17,11 +17,13 @@ jobs:
       - uses: actions/checkout@v3
 
       - name: Clone Suitesparse
-        run: git clone --depth 1 --branch v7.2.0 https://github.com/DrTimothyAldenDavis/SuiteSparse suitesparse
+        run: make suitesparse
 
-      - name: Clone Pybind
-        run: |
-          git clone --depth 1 --branch stable https://github.com/pybind/pybind11 pybind11
+      - name: Clone XLA
+        run: make xla
+
+      - name: Clone Pybind11
+        run: make pybind11
 
       - name: Build wheels
         uses: pypa/[email protected]
@@ -30,6 +32,8 @@ jobs:
           CIBW_ARCHS_LINUX: x86_64
           CIBW_ARCHS_WINDOWS: AMD64
           CIBW_SKIP: '*-musllinux* pp*'
+          CIBW_BEFORE_TEST: 'pip install pytest'
+          CIBW_TEST_COMMAND: 'pytest tests.py'
 
       - uses: actions/upload-artifact@v3
         with:

.github/workflows/pr.yml

@@ -0,0 +1,36 @@
+name: pr
+on:
+  pull_request:
+
+jobs:
+  build_wheels:
+    name: Build wheels on ${{ matrix.os }}
+    runs-on: ${{ matrix.os }}
+    strategy:
+      fail-fast: false
+      matrix:
+        os: [ubuntu-latest, windows-latest, macos-latest]
+
+    # for a PR we just check if it can build for python=3.12
+    steps:
+      - uses: actions/checkout@v3
+
+      - name: Clone Suitesparse
+        run: make suitesparse
+
+      - name: Clone XLA
+        run: make xla
+
+      - name: Clone Pybind11
+        run: make pybind11
+
+      - name: Build wheels
+        uses: pypa/[email protected]
+        env:
+          CIBW_ARCHS_MACOS: x86_64 arm64
+          CIBW_ARCHS_LINUX: x86_64
+          CIBW_ARCHS_WINDOWS: AMD64
+          CIBW_BUILD: '*cp312*'
+          CIBW_SKIP: '*-musllinux* pp*'
+          CIBW_BEFORE_TEST: 'pip install pytest'
+          CIBW_TEST_COMMAND: 'python {project}/.github/run_tests.py'

.gitignore

@@ -3,6 +3,8 @@ __pycache__/
 *.py[cod]
 *$py.class
 suitesparse
+xla
+pybind11
 
 # C extensions
 *.so

.pre-commit-config.yaml

@@ -9,9 +9,15 @@ repos:
     rev: 23.9.1
     hooks:
       - id: black
-        language_version: python3.11
   - repo: https://github.com/pocc/pre-commit-hooks
     rev: '336fdd7'
     hooks:
       - id: clang-format
-        args: [--style=Google]
+  - repo: https://github.com/kynan/nbstripout
+    rev: 0.6.0
+    hooks:
+      - id: nbstripout
+  - repo: https://github.com/psf/black-pre-commit-mirror
+    rev: 23.9.1
+    hooks:
+      - id: black-jupyter

CHANGELOG.md

@@ -0,0 +1,71 @@
+# Changelog
+
+## 0.3.0
+- Implement new FFI API for C++ extension
+- Enhance C++ testing and error handling (proper error throwing instead of segfaulting)
+- Enable JIT and `vmap` for optimized performance
+- Improve shape handling for arrays
+- Retrieve array sizes from C++ buffers
+- Refactor workflows and remove deprecated implementations
+- Remove old notebook files
+- Upgrade dependencies and ensure compatibility with C++17 standard
+
+## 0.2.10
+- Run tests post wheel build
+- Pin exact dependency versions
+- Streamline GitHub workflows and update `setup.py`
+- Fix issues with `vmap` over array `b`
+
+## 0.2.8
+- Refine CI/CD environment variable configuration for `cibuildwheel`
+
+## 0.2.7
+- Prevent memory leaks
+- Introduce pre-commit configuration
+- Update `cibuildwheel` configuration
+- Clone specific SuiteSparse version
+
+## 0.2.5
+- Address deprecations in XLA translations
+
+## 0.2.4
+- Add support for Python 3.12
+- Consolidate GitHub workflow files and update package metadata
+
+## 0.2.0
+- Vendor SuiteSparse library in source distribution
+- Re-enable `PIP_FIND_LINKS`
+- Update build recipes and dependencies
+- Improve README with setup and build instructions
+
+## 0.1.4
+- Add support for Python 3.11
+
+## 0.1.3
+- Enable installation on macOS
+- Fix issues with static linking on macOS (C++11 requirement)
+
+## 0.1.1
+- Publish release on PyPI and include tarball for distribution
+- Add support for multiple Python versions and manylinux2014 wheels
+
+## 0.1.0
+- Enable custom JVP/VJP rules
+- Improve differentiation features with forward-mode JVP and transposition
+- Add `pyproject.toml` for better build configuration
+
+## 0.0.6
+- Add more library/include paths for builds
+- Refine README and setup instructions
+- Initial integration of complex value handling in `vmap`
+
+## 0.0.4
+- Add `bump2version` configuration for automated versioning
+- Bugfix: Correct matrix-vector multiplication in COO format (`mul_coo_vec`)
+
+## 0.0.3
+- Set up core functionality with sparse matrix multiplication (COO format)
+- Integrate `vmap` for float64 and complex128 arrays
+- Initial setup with Makefile, test suites, and Docker configuration
+- Begin development of XLA translations and gradient support
+

MANIFEST.in

@@ -1,12 +1,19 @@
+include tests.py
 include klujax.cpp
-include suitesparse/SuiteSparse_config/*.h
-include suitesparse/SuiteSparse_config/*.c
+include pybind11/include/pybind11/*.h
+include pybind11/include/pybind11/detail/*.h
+include pybind11/include/pybind11/eigen/*.h
+include pybind11/include/pybind11/stl/*.h
 include suitesparse/AMD/Include/*.h
-include suitesparse/COLAMD/Include/*.h
-include suitesparse/BTF/Include/*.h
-include suitesparse/KLU/Include/*.h
 include suitesparse/AMD/Source/*.c
-include suitesparse/COLAMD/Source/*.c
+include suitesparse/BTF/Include/*.h
 include suitesparse/BTF/Source/*.c
+include suitesparse/COLAMD/Include/*.h
+include suitesparse/COLAMD/Source/*.c
+include suitesparse/KLU/Include/*.h
 include suitesparse/KLU/Source/*.c
 include suitesparse/LICENSE.txt
+include suitesparse/SuiteSparse_config/*.c
+include suitesparse/SuiteSparse_config/*.h
+include xla/xla/ffi/api/*.cc
+include xla/xla/ffi/api/*.h

Makefile

@@ -1,24 +1,40 @@
 dist:
 	python setup.py build sdist bdist_wheel
 
+test:
+	pytest tests.py
+
 .PHONY: build
 build:
 	python setup.py build_ext --inplace
 
-test:
-	pytest tests.py
-
+.PHONY: suitesparse
 suitesparse:
+	rm -rf suitesparse
 	git clone --depth 1 --branch v7.2.0 https://github.com/DrTimothyAldenDavis/SuiteSparse suitesparse || true
 	cd suitesparse && rm -rf .git
 
+.PHONY: xla
+xla:
+	rm -rf xla
+	git clone --depth 1 --branch main https://github.com/openxla/xla xla
+	cd xla && rm -rf .git
+
+.PHONY: pybind11
+pybind11:
+	rm -rf pybind11
+	git clone --depth 1 --branch stable https://github.com/pybind/pybind11 pybind11
+	cd pybind11 && rm -rf .git
+
 clean:
-	find . -not -path "./suitesparse*" -name "dist" | xargs rm -rf
-	find . -not -path "./suitesparse*" -name "build" | xargs rm -rf
-	find . -not -path "./suitesparse*" -name "builds" | xargs rm -rf
-	find . -not -path "./suitesparse*" -name "__pycache__" | xargs rm -rf
-	find . -not -path "./suitesparse*" -name "*.so" | xargs rm -rf
-	find . -not -path "./suitesparse*" -name "*.egg-info" | xargs rm -rf
-	find . -not -path "./suitesparse*" -name ".ipynb_checkpoints" | xargs rm -rf
-	find . -not -path "./suitesparse*" -name ".pytest_cache" | xargs rm -rf
+	find . -name "dist" | xargs rm -rf
+	find . -name "build" | xargs rm -rf
+	find . -name "builds" | xargs rm -rf
+	find . -name "__pycache__" | xargs rm -rf
+	find . -name "*.so" | xargs rm -rf
+	find . -name "*.egg-info" | xargs rm -rf
+	find . -name ".ipynb_checkpoints" | xargs rm -rf
+	find . -name ".pytest_cache" | xargs rm -rf
 
+env:
+	@echo export CPLUS_INCLUDE_PATH='/home/flaport/Projects/klujax/xla:/home/flaport/.anaconda/include/python3.12:/home/flaport/.anaconda/lib/python3.12/site-packages/pybind11/include:/home/flaport/Projects/klujax/suitesparse/SuiteSparse_config:/home/flaport/Projects/klujax/suitesparse/AMD/Include:/home/flaport/Projects/klujax/suitesparse/COLAMD/Include:/home/flaport/Projects/klujax/suitesparse/BTF/Include:/home/flaport/Projects/klujax/suitesparse/KLU/Include'