Metadata-Version: 2.3
Name: numba4jax
Version: 0.0.12.post3
Summary: Usa numba in jax-compiled kernels.
Project-URL: homepage, http://github.com/PhilipVinc/numba4jax
Project-URL: repository, http://github.com/PhilipVinc/numba4jax
Author-email: Filippo Vicentini <filippovicentini@gmail.com>
License: MIT
License-File: LICENSE.txt
Keywords: Jax,Numba,compile,jit,kernel
Classifier: Development Status :: 5 - Production/Stable
Classifier: Intended Audience :: Science/Research
Classifier: Operating System :: MacOS :: MacOS X
Classifier: Operating System :: POSIX :: Linux
Classifier: Operating System :: Unix
Classifier: Programming Language :: Python :: 3
Requires-Python: >=3.9
Requires-Dist: cffi>=1.14.4
Requires-Dist: jax<=0.4.30,>=0.4.16
Requires-Dist: numba<0.61,>=0.53
Requires-Dist: numpy<3.0,>=1.22
Provides-Extra: dev
Requires-Dist: black==24.8.0; extra == 'dev'
Requires-Dist: build; extra == 'dev'
Requires-Dist: coverage>=5; extra == 'dev'
Requires-Dist: flaky>=3.7; extra == 'dev'
Requires-Dist: pre-commit; extra == 'dev'
Requires-Dist: pytest-cov>=2.10.1; extra == 'dev'
Requires-Dist: pytest>=6; extra == 'dev'
Requires-Dist: ruff==0.5.6; extra == 'dev'
Requires-Dist: wheel; extra == 'dev'
Description-Content-Type: text/markdown

# numba4jax



A small experimental python package allowing you to use numba-jitted functions from within jax 
with no overhead.

This package uses the CFFI of Numba to expose the C Function pointer of your compiled
function to XLA. It works both for CPU and GPU functions.

This package exports a single decorator `@njit4jax`, which takes an argument, a function
or Tuple describing the output shape of the function itself.
See the brief example below.

```python

import jax
import jax.numpy as jnp

from numba4jax import ShapedArray, njit4jax


def compute_type(*x):
    return x[0]


@njit4jax(compute_type)
def test(args):
    y, x, x2 = args
    y[:] = x[:] + 1


z = jnp.ones((1, 2), dtype=float)

jax.make_jaxpr(test)(z, z)

print("output: ", test(z, z))
print("output: ", jax.jit(test)(z, z))

z = jnp.ones((2, 3), dtype=float)
print("output: ", jax.jit(test)(z, z))

z = jnp.ones((1, 3, 1), dtype=float)
print("output: ", jax.jit(test)(z, z))

```

## Backend support

This package supports both the CPU and GPU backends of jax.
The GPU backend is only supported on linux, and is highly experimental.
It requires CUDA to be installed in a standard path.
CUDA is found through `numba.cuda`, so you should first check that `numba.cuda`
works.
