jax-mlx-plugin 0.0.1__tar.gz

jax_mlx_plugin-0.0.1/CMakeLists.txt

@@ -0,0 +1,94 @@
+cmake_minimum_required(VERSION 3.15)
+project(jax_mlx)
+
+set(CMAKE_CXX_STANDARD 17)
+set(CMAKE_OSX_DEPLOYMENT_TARGET "26.0")
+
+# --- Dependencies ---
+
+# 1. Find Python and pybind11
+find_package(Python3 REQUIRED COMPONENTS Interpreter Development)
+
+# Get pybind11 cmake path from pip install
+execute_process(
+    COMMAND ${Python3_EXECUTABLE} -c "import pybind11; print(pybind11.get_cmake_dir())"
+    OUTPUT_VARIABLE PYBIND11_CMAKE_DIR
+    OUTPUT_STRIP_TRAILING_WHITESPACE
+    RESULT_VARIABLE PYBIND11_FIND_RESULT
+)
+if(NOT PYBIND11_FIND_RESULT EQUAL 0)
+    message(FATAL_ERROR "Could not find pybind11. Install with: pip install pybind11")
+endif()
+list(APPEND CMAKE_PREFIX_PATH "${PYBIND11_CMAKE_DIR}")
+
+find_package(pybind11 REQUIRED)
+
+# 2. Find MLX from pip installation
+# The mlx package installs headers in site-packages/mlx/include
+# and libmlx.dylib in site-packages/mlx/lib
+execute_process(
+    COMMAND ${Python3_EXECUTABLE} -c "import mlx.core; import os; print(os.path.dirname(mlx.core.__file__))"
+    OUTPUT_VARIABLE MLX_PYTHON_PATH
+    OUTPUT_STRIP_TRAILING_WHITESPACE
+    RESULT_VARIABLE MLX_FIND_RESULT
+)
+
+if(NOT MLX_FIND_RESULT EQUAL 0)
+    message(FATAL_ERROR "Could not find MLX Python package. Install with: pip install mlx")
+endif()
+
+set(MLX_INCLUDE_DIR "${MLX_PYTHON_PATH}/include")
+set(MLX_LIB_DIR "${MLX_PYTHON_PATH}/lib")
+
+message(STATUS "Found MLX at: ${MLX_PYTHON_PATH}")
+message(STATUS "  Include: ${MLX_INCLUDE_DIR}")
+message(STATUS "  Lib: ${MLX_LIB_DIR}")
+
+# Verify paths exist
+if(NOT EXISTS "${MLX_INCLUDE_DIR}")
+    message(FATAL_ERROR "MLX include directory not found: ${MLX_INCLUDE_DIR}")
+endif()
+if(NOT EXISTS "${MLX_LIB_DIR}/libmlx.dylib")
+    message(FATAL_ERROR "MLX library not found: ${MLX_LIB_DIR}/libmlx.dylib")
+endif()
+
+# 3. OpenXLA Headers (for PJRT C API)
+include_directories(${CMAKE_CURRENT_SOURCE_DIR}/third_party)
+
+# --- The Plugin Library ---
+add_library(mlx_pjrt_plugin SHARED  
+    src/jax_mlx_pjrt.cpp
+)
+
+# Include directories
+target_include_directories(mlx_pjrt_plugin PRIVATE
+    ${MLX_INCLUDE_DIR}
+    ${Python3_INCLUDE_DIRS}
+)
+
+# Link against MLX and pybind11
+target_link_libraries(mlx_pjrt_plugin   
+    PRIVATE   
+    ${MLX_LIB_DIR}/libmlx.dylib
+    pybind11::pybind11
+)
+
+# Set RPATH to find libmlx.dylib at runtime
+# Plugin installs to site-packages/jax_mlx/
+# MLX is at site-packages/mlx/lib/
+# So relative path is @loader_path/../mlx/lib
+set_target_properties(mlx_pjrt_plugin PROPERTIES
+    INSTALL_RPATH "@loader_path;@loader_path/../mlx/lib"
+    BUILD_WITH_INSTALL_RPATH TRUE
+)
+
+# macOS specific: Allow undefined symbols (Python provides them at runtime)
+if (APPLE)  
+    target_link_options(mlx_pjrt_plugin PRIVATE "-undefined" "dynamic_lookup")  
+endif()
+
+# Install target
+install(TARGETS mlx_pjrt_plugin
+    LIBRARY DESTINATION jax_mlx
+    RUNTIME DESTINATION jax_mlx
+)

jax_mlx_plugin-0.0.1/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Thomas Summe
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

jax_mlx_plugin-0.0.1/MANIFEST.in

@@ -0,0 +1,15 @@
+# Include CMake and C++ source files for building the extension
+include CMakeLists.txt
+recursive-include src *.cpp *.h
+recursive-include third_party *.h
+
+# Include documentation
+include README.md
+include LICENSE
+
+# Exclude build artifacts
+global-exclude __pycache__
+global-exclude *.py[cod]
+global-exclude *.so
+global-exclude *.dylib
+prune benchmarks/build

jax_mlx_plugin-0.0.1/PKG-INFO

@@ -0,0 +1,119 @@
+Metadata-Version: 2.4
+Name: jax-mlx-plugin
+Version: 0.0.1
+Summary: JAX PJRT plugin for Apple Silicon using MLX
+Home-page: https://github.com/tsumme1/jax-mlx
+Author: Thomas Summe
+License: MIT
+Project-URL: Homepage, https://github.com/tsumme1/jax-mlx
+Project-URL: Repository, https://github.com/tsumme1/jax-mlx
+Keywords: jax,mlx,apple-silicon,machine-learning,pjrt
+Classifier: Development Status :: 3 - Alpha
+Classifier: Environment :: GPU
+Classifier: Environment :: MacOS X
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: MacOS :: MacOS X
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: jax>=0.5.0
+Requires-Dist: jaxlib>=0.5.0
+Requires-Dist: mlx
+Dynamic: home-page
+Dynamic: license-file
+Dynamic: requires-python
+
+# JAX MLX Plugin
+
+A PJRT plugin enabling JAX to use Apple's MLX framework as a backend on Apple Silicon Macs.
+
+## Status
+
+✅ **362 ops tested and passing**
+
+
+## Requirements
+
+- **Apple Silicon Mac** (M1/M2/M3/M4)
+- **Python:** 3.11+
+- **Dependencies:** jax, jaxlib, mlx
+
+## Installation
+
+```bash
+# Install build dependencies
+pip install mlx jaxlib jax
+
+# Install the plugin
+pip install .
+```
+
+## Usage
+
+```python
+import jax
+import jax.numpy as jnp
+
+# List available devices
+print(jax.devices())  # [mlx:0]
+
+# Use MLX as default device
+mlx = jax.devices('mlx')[0]
+with jax.default_device(mlx):
+    x = jnp.array([1.0, 2.0, 3.0])
+    y = jnp.sin(x) + jnp.cos(x)
+    print(y)
+```
+
+## Features
+
+- ✅ All core JAX operations (362 tested)
+- ✅ Full autodiff support (`jax.grad`, `jax.value_and_grad`)
+- ✅ JIT compilation with `mx.compile()` kernel fusion
+- ✅ Vectorization (`jax.vmap`)
+- ✅ Control flow (`lax.cond`, `lax.while_loop`, `lax.scan`)
+- ✅ Linear algebra, FFT, convolutions
+- ✅ Neural network training (Flax, Optax)
+
+## Environment Variables
+
+| Variable | Description |
+|----------|-------------|
+| `MLX_PJRT_DEBUG=1` | Enable verbose debug logging |
+| `MLX_NO_COMPILE=1` | Disable mx.compile() kernel fusion |
+| `MLX_TIMING=1` | Enable timing output |
+
+## Development
+
+```bash
+# Run exhaustive tests (362 ops)
+python tests/test_exhaustive.py
+
+# Run CNN benchmarks
+python benchmarks/benchmark_cnn.py           # JAX/Flax
+python benchmarks/benchmark_mlx_native.py    # Native MLX Python
+
+# Build and run C++ benchmark
+cd benchmarks && cmake -B build && cmake --build build
+./build/mlx_cpp_benchmark
+```
+
+## Architecture
+
+The plugin implements the PJRT (Portable JAX Runtime) C API. StableHLO operations from JAX are parsed and converted to MLX operations at runtime using a lightweight MLIR parser. The plugin uses `mx.compile()` for GPU kernel fusion.
+
+## Known Limitations
+
+- **Float64:** Not supported on Metal GPU (use Float32)
+- **While loops:** Block kernel fusion (require runtime eval)
+
+## License
+
+MIT License - see [LICENSE](LICENSE) for details.

jax_mlx_plugin-0.0.1/README.md

@@ -0,0 +1,87 @@
+# JAX MLX Plugin
+
+A PJRT plugin enabling JAX to use Apple's MLX framework as a backend on Apple Silicon Macs.
+
+## Status
+
+✅ **362 ops tested and passing**
+
+
+## Requirements
+
+- **Apple Silicon Mac** (M1/M2/M3/M4)
+- **Python:** 3.11+
+- **Dependencies:** jax, jaxlib, mlx
+
+## Installation
+
+```bash
+# Install build dependencies
+pip install mlx jaxlib jax
+
+# Install the plugin
+pip install .
+```
+
+## Usage
+
+```python
+import jax
+import jax.numpy as jnp
+
+# List available devices
+print(jax.devices())  # [mlx:0]
+
+# Use MLX as default device
+mlx = jax.devices('mlx')[0]
+with jax.default_device(mlx):
+    x = jnp.array([1.0, 2.0, 3.0])
+    y = jnp.sin(x) + jnp.cos(x)
+    print(y)
+```
+
+## Features
+
+- ✅ All core JAX operations (362 tested)
+- ✅ Full autodiff support (`jax.grad`, `jax.value_and_grad`)
+- ✅ JIT compilation with `mx.compile()` kernel fusion
+- ✅ Vectorization (`jax.vmap`)
+- ✅ Control flow (`lax.cond`, `lax.while_loop`, `lax.scan`)
+- ✅ Linear algebra, FFT, convolutions
+- ✅ Neural network training (Flax, Optax)
+
+## Environment Variables
+
+| Variable | Description |
+|----------|-------------|
+| `MLX_PJRT_DEBUG=1` | Enable verbose debug logging |
+| `MLX_NO_COMPILE=1` | Disable mx.compile() kernel fusion |
+| `MLX_TIMING=1` | Enable timing output |
+
+## Development
+
+```bash
+# Run exhaustive tests (362 ops)
+python tests/test_exhaustive.py
+
+# Run CNN benchmarks
+python benchmarks/benchmark_cnn.py           # JAX/Flax
+python benchmarks/benchmark_mlx_native.py    # Native MLX Python
+
+# Build and run C++ benchmark
+cd benchmarks && cmake -B build && cmake --build build
+./build/mlx_cpp_benchmark
+```
+
+## Architecture
+
+The plugin implements the PJRT (Portable JAX Runtime) C API. StableHLO operations from JAX are parsed and converted to MLX operations at runtime using a lightweight MLIR parser. The plugin uses `mx.compile()` for GPU kernel fusion.
+
+## Known Limitations
+
+- **Float64:** Not supported on Metal GPU (use Float32)
+- **While loops:** Block kernel fusion (require runtime eval)
+
+## License
+
+MIT License - see [LICENSE](LICENSE) for details.

jax_mlx_plugin-0.0.1/pyproject.toml

@@ -0,0 +1,41 @@
+[build-system]
+requires = ["setuptools>=42", "wheel", "pybind11", "cmake", "mlx"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "jax-mlx-plugin"
+version = "0.0.1"
+description = "JAX PJRT plugin for Apple Silicon using MLX"
+readme = "README.md"
+license = {text = "MIT"}
+requires-python = ">=3.11"
+authors = [
+    {name = "Thomas Summe"}
+]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Environment :: GPU",
+    "Environment :: MacOS X",
+    "Intended Audience :: Developers",
+    "Intended Audience :: Science/Research",
+    "License :: OSI Approved :: MIT License",
+    "Operating System :: MacOS :: MacOS X",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+keywords = ["jax", "mlx", "apple-silicon", "machine-learning", "pjrt"]
+dependencies = [
+    "jax>=0.5.0",
+    "jaxlib>=0.5.0",
+    "mlx",
+]
+
+[project.urls]
+Homepage = "https://github.com/tsumme1/jax-mlx"
+Repository = "https://github.com/tsumme1/jax-mlx"
+
+[project.entry-points."jax_plugins"]
+mlx_plugin = "jax_mlx.plugin"

jax_mlx_plugin-0.0.1/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

jax_mlx_plugin-0.0.1/setup.py

@@ -0,0 +1,83 @@
+import os
+import sys
+import subprocess
+import shutil
+from setuptools import setup, Extension, find_packages
+from setuptools.command.build_ext import build_ext
+
+class CMakeExtension(Extension):
+    def __init__(self, name, sourcedir=''):
+        Extension.__init__(self, name, sources=[])
+        self.sourcedir = os.path.abspath(sourcedir)
+
+class CMakeBuild(build_ext):
+    def run(self):
+        for ext in self.extensions:
+            self.build_extension(ext)
+
+    def build_extension(self, ext):
+        extdir = os.path.abspath(os.path.dirname(self.get_ext_fullpath(ext.name)))
+        if not extdir.endswith(os.path.sep):
+            extdir += os.path.sep
+
+        cfg = 'Debug' if self.debug else 'Release'
+        
+        cmake_args = [
+            f'-DCMAKE_LIBRARY_OUTPUT_DIRECTORY={extdir}',
+            f'-DPython3_EXECUTABLE={sys.executable}',
+            f'-DCMAKE_BUILD_TYPE={cfg}',
+        ]
+
+        build_args = ['--config', cfg]
+
+        if not os.path.exists(self.build_temp):
+            os.makedirs(self.build_temp)
+
+        subprocess.check_call(['cmake', ext.sourcedir] + cmake_args, cwd=self.build_temp)
+        subprocess.check_call(['cmake', '--build', '.', '-j'] + build_args, cwd=self.build_temp)
+        
+        # Copy the built library to the package source directory as well
+        src_package_dir = os.path.join(ext.sourcedir, 'src', 'jax_mlx')
+        built_lib = os.path.join(extdir, 'libmlx_pjrt_plugin.dylib')
+        if os.path.exists(built_lib) and os.path.isdir(src_package_dir):
+            shutil.copy2(built_lib, src_package_dir)
+
+setup(
+    name='jax-mlx-plugin',
+    version='0.0.1',
+    author='Thomas Summe',
+    description='JAX PJRT plugin for Apple Silicon using MLX',
+    long_description=open('README.md').read(),
+    long_description_content_type='text/markdown',
+    url='https://github.com/tsumme1/jax-mlx',
+    packages=find_packages(where='src'),
+    package_dir={'': 'src'},
+    ext_modules=[CMakeExtension('jax_mlx.mlx_pjrt_plugin')],
+    cmdclass=dict(build_ext=CMakeBuild),
+    package_data={'jax_mlx': ['*.dylib', '*.so']},
+    include_package_data=True,
+    zip_safe=False,
+    python_requires='>=3.11',
+    install_requires=[
+        'jax>=0.5.0',
+        'jaxlib>=0.5.0',
+        'mlx',
+    ],
+    entry_points={
+        "jax_plugins": [
+            "mlx_plugin = jax_mlx.plugin",
+        ],
+    },
+    classifiers=[
+        'Development Status :: 3 - Alpha',
+        'Intended Audience :: Developers',
+        'Intended Audience :: Science/Research',
+        'License :: OSI Approved :: MIT License',
+        'Operating System :: MacOS :: MacOS X',
+        'Programming Language :: Python :: 3',
+        'Programming Language :: Python :: 3.11',
+        'Programming Language :: Python :: 3.12',
+        'Programming Language :: Python :: 3.13',
+        'Topic :: Scientific/Engineering :: Artificial Intelligence',
+    ],
+)

jax_mlx_plugin-0.0.1/src/jax_mlx/__init__.py

@@ -0,0 +1,6 @@
+# jax-mlx: JAX PJRT plugin for Apple Silicon using MLX
+#
+# This package provides MLX acceleration for JAX via the PJRT plugin interface.
+# Registration is handled automatically by plugin.py via JAX's plugin mechanism.
+
+__version__ = "0.0.1"