Metadata-Version: 2.1
Name: jaxbind
Version: 1.3.1
Summary: Bind any function written in another language to JAX with support for JVP/VJP/batching/jit compilation
Author-Email: Jakob Roth <roth@mpa-garching.mpg.de>, Martin Reinecke <martin@mpa-garching.mpg.de>, Gordian Edenhofer <gordian.edenhofer@gmail.com>
License: BSD 2-Clause License
         
         Copyright (c) 2024, Max-Planck-Society
         
         Redistribution and use in source and binary forms, with or without
         modification, are permitted provided that the following conditions are met:
         
         1. Redistributions of source code must retain the above copyright notice, this
            list of conditions and the following disclaimer.
         
         2. Redistributions in binary form must reproduce the above copyright notice,
            this list of conditions and the following disclaimer in the documentation
            and/or other materials provided with the distribution.
         
         THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
         AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
         IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
         DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
         FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
         DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
         SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
         CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
         OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
         OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Science/Research
Classifier: Topic :: Scientific/Engineering :: Mathematics
Classifier: Topic :: Scientific/Engineering :: Physics
Classifier: License :: OSI Approved :: BSD License
Classifier: Programming Language :: C++
Classifier: Programming Language :: Python
Project-URL: Source Code, https://github.com/NIFTy-PPL/JAXbind
Project-URL: Bug Tracker, https://github.com/NIFTy-PPL/JAXbind/issues
Project-URL: README, https://github.com/NIFTy-PPL/JAXbind/blob/main/README.md
Requires-Python: >=3.10
Requires-Dist: packaging
Requires-Dist: numpy>=1.17.0
Requires-Dist: jax>=0.5
Requires-Dist: jaxlib>=0.5
Provides-Extra: test
Requires-Dist: pytest; extra == "test"
Requires-Dist: scipy; extra == "test"
Requires-Dist: ducc0; extra == "test"
Provides-Extra: doc
Requires-Dist: sphinx; extra == "doc"
Requires-Dist: pydata-sphinx-theme; extra == "doc"
Requires-Dist: myst-parser; extra == "doc"
Requires-Dist: ipykernel; extra == "doc"
Requires-Dist: jupyterlab; extra == "doc"
Requires-Dist: notebook; extra == "doc"
Requires-Dist: jupytext; extra == "doc"
Requires-Dist: nbconvert; extra == "doc"
Requires-Dist: ducc0; extra == "doc"
Requires-Dist: scipy; extra == "doc"
Provides-Extra: all
Requires-Dist: jaxbind[doc,test]; extra == "all"
Description-Content-Type: text/markdown

# JAXbind: Bind any function to JAX

JAXbind API documentation: [nifty-ppl.github.io/JAXbind/](https://nifty-ppl.github.io/JAXbind/)
| Found a bug? [github.com/NIFTy-PPL/JAXbind/issues](https://github.com/NIFTy-PPL/JAXbind/issues)
| Need help? [github.com/NIFTy-PPL/JAXbind/discussions](https://github.com/NIFTy-PPL/JAXbind/discussions)

## Summary

The existing interface in JAX for connecting fully differentiable custom code requires deep knowledge of JAX and its C++ backend.
The aim of `JAXbind` is to drastically lower the burden of connecting custom functions implemented in other programming languages to JAX.
Specifically, `JAXbind` provides an easy-to-use Python interface for defining custom, so-called JAX primitives.
Via `JAXbind`, any function callable from Python can be exposed as a JAX primitive.
`JAXbind` allows to interface the JAX function transformation engine with custom derivatives and batching rules, enabling all JAX transformations for the custom primitive.
In contrast, the JAX built-in [external callback interface](https://jax.readthedocs.io/en/latest/notebooks/external_callbacks.html) also has a Python endpoint but the external callbacks cannot be fully integrated into the JAX transformation engine, as only the [Jacobian-vector product](https://jax.readthedocs.io/en/latest/_autosummary/jax.custom_jvp.html) or the [vector-Jacobian product](https://jax.readthedocs.io/en/latest/_autosummary/jax.custom_vjp.html) can be added but not both.

### Automatic Differentiation and Code Example

Automatic differentiation is a core feature of JAX and often one of the main reasons for using it.
Thus, it is essential that custom functions registered with JAX support automatic differentiation.
In the following, we will outline which functions our package respectively JAX requires to enable automatic differentiation.
For simplicity, we assume that we want to connect the nonlinear function $f(x_1,x_2) = x_1x_2^2$ to JAX.
The `JAXbind` package expects the Python function for $f$ to take three positional arguments.
The first argument, `out`, is a `tuple` into which the function results are written.
The second argument is also a `tuple` containing the input to the function, in our case, $x_1$ and $x_2$.
Via `kwargs_dump`, potential keyword arguments given to the later registered Jax primitive can be forwarded to `f` in serialized form.

```python
import jaxbind

def f(out, args, kwargs_dump):
    kwargs = jaxbind.load_kwargs(kwargs_dump)
    x1, x2 = args
    out[0][()] = x1 * x2**2
```

JAX's automatic differentiation engine can compute the Jacobian-vector product `jvp` and vector-Jacobian product `vjp` of JAX primitives.
The Jacobian-vector product in JAX is a function applying the Jacobian of $f$ at a position $x$ to a tangent vector.
In mathematical nomenclature this operation is called the pushforward of $f$ and can be denoted as $\partial f(x): T_x X \mapsto T_{f(x)} Y$, with $T_x X$ and $T_{f(x)} Y$ being the tangent spaces of $X$ and $Y$ at the positions $x$ and $f(x)$.
As the implementation of $f$ is not JAX native, JAX cannot automatically compute the `jvp`.
Instead, an implementation of the pushforward has to be provided, which `JAXbind` will register as the `jvp` of the JAX primitive of $f$.
For our example, this Jacobian-vector-product function is given by $\partial f(x_1,x_2)(dx_1,dx_2) = x_2^2dx_1 + 2x_1x_2dx_2$.

```python
def f_jvp(out, args, kwargs_dump):
    kwargs = jaxbind.load_kwargs(kwargs_dump)
    x1, x2, dx1, dx2 = args
    out[0][()] = x2**2 * dx1 + 2 * x1 * x2 * dx2
```

The vector-Jacobian product `vjp` in JAX is the linear transpose of the Jacobian-vector product.
In mathematical nomenclature this is the pullback $(\partial f(x))^{T}: T_{f(x)}Y \mapsto T_x X$ of $f$.
Analogously to the `jvp`, the user has to implement this function as JAX cannot automatically construct it.
For our example function, the vector-Jacobian product is $(\partial f(x_1,x_2))^{T}(dy) = (x_2^2dy, 2x_1x_2dy)$.

```python
def f_vjp(out, args, kwargs_dump):
    kwargs = jaxbind.load_kwargs(kwargs_dump)
    x1, x2, dy = args
    out[0][()] = x2**2 * dy
    out[1][()] = 2 * x1 * x2 * dy
```

To just-in-time compile the function, JAX needs to abstractly evaluate the code, i.e. it needs to be able to know the shape and dtype of the output of the custom function given only the shape and dtype of the input.
We have to provide these abstract evaluation functions returning the output shape and dtype given an input shape and dtype for `f` as well as for the `vjp` application.
The output shape of the `jvp` is identical to the output shape of `f` itself and does not need to be specified again.
Due to the internals of JAX the abstract evaluation functions take normal keyword arguments and not serialized keyword arguments.

```python
def f_abstract(*args, **kwargs):
    assert args[0].shape == args[1].shape
    return ((args[0].shape, args[0].dtype),)

def f_abstract_T(*args, **kwargs):
    return (
        (args[0].shape, args[0].dtype),
        (args[0].shape, args[0].dtype),
    )
```

We have now defined all ingredients necessary to register a JAX primitive for our function $f$ using the `JAXbind` package.

```python
f_jax = jaxbind.get_nonlinear_call(
    f, (f_jvp, f_vjp), f_abstract, f_abstract_T
)
```

`f_jax` is a JAX primitive registered via the `JAXbind` package supporting all JAX transformations.
We can now compute the `jvp` and `vjp` of the new JAX primitive and even jit-compile and batch it.

```python
import jax
import jax.numpy as jnp

inp = (jnp.full((4,3), 4.), jnp.full((4,3), 2.))
tan = (jnp.full((4,3), 1.), jnp.full((4,3), 1.))
res, res_tan = jax.jvp(f_jax, inp, tan)

cotan = [jnp.full((4,3), 6.)]
res, f_vjp = jax.vjp(f_jax, *inp)
res_cotan = f_vjp(cotan)

f_jax_jit = jax.jit(f_jax)
res = f_jax_jit(*inp)
```

### Higher Order Derivatives and Linear Functions

JAX supports higher order derivatives and can differentiate a `jvp` or `vjp` with respect to the position at which the Jacobian was taken.
Similar to first derivatives, JAX can not automatically compute higher derivatives of a general function $f$ that is not natively implemented in JAX.
Higher order derivatives would again need to be provided by the user.
For many algorithms, first derivatives are sufficient, and higher order derivatives are often not implemented by the high-performance codes.
Therefore, the current interface of `JAXbind` is, for simplicity, restricted to first derivatives.
In the future, the interface could be easily expanded if specific use cases require higher order derivatives.

In scientific computing, linear functions such as, e.g., spherical harmonic transforms are widespread.
If the function $f$ is linear, differentiation becomes trivial.
Specifically for a linear function $f$, the pushforward respectively the `jvp` of $f$ is identical to $f$ itself and independent of the position at which it is computed.
Expressed in formulas, $\partial f(x)(dx) = f(dx)$ if $f$ is linear in $x$.
Analogously, the pullback respectively the `vjp` becomes independent of the initial position and is given by the linear transpose of $f$, thus $(\partial f(x))^{T}(dy) = f^T(dy)$.
Also, all higher order derivatives can be expressed in terms of $f$ and its transpose.
To make use of these simplifications, `JAXbind` provides a special interface for linear functions, supporting higher order derivatives, only requiring an implementation of the function and its transpose.

### Demos and Documentation

Additional demos can be found in the demos folder.
Specifically, there is a basic demo [01_linear_function.py](https://github.com/NIFTy-PPL/JAXbind/blob/main/demos/01_linear_function.py) showcasing the interface for linear functions and custom batching rules.
[02_multilinear_function.py](https://github.com/NIFTy-PPL/JAXbind/blob/main/demos/02_multilinear_function.py) binds a multi-linear function as a JAX primitive.
Finally, [03_nonlinear_function.py](https://github.com/NIFTy-PPL/JAXbind/blob/main/demos/03_nonlinear_function.py) demonstrates the interface for non-linear functions and shows how to deal with fixed arguments, which cannot be differentiated.
JAXbind provides bindings to parts of the functionality of the [DUCC](https://gitlab.mpcdf.mpg.de/mtr/ducc) package.
The DUCC bindings are also exposed as a [webpage](https://nifty-ppl.github.io/JAXbind/contrib) to showcase a real-world example of the usage of JAXbind.
The documentation of the JAXbind API is available [here](https://nifty-ppl.github.io/JAXbind/).

## Platforms

Currently, `JAXbind` only has CPU but no GPU support.
With some expertise on Python bindings for GPU kernels adding GPU support should be fairly simple.
The interfacing with the JAX automatic differentiation engine is identical for CPU and GPU.
Contributions are welcome!

## Installation

Binary wheels for JAXbind can be obtained and installed from PyPI via:

```
pip install jaxbind
```

To install JAXbind from source, clone the repository and install the package via pip.

```
git clone https://github.com/NIFTy-PPL/jaxbind.git
cd jaxbind
pip install .
```

## Contributing

Contributions are highly appreciated!
Please open an issue first if you think your PR changes current code substantially.
Please format your code using black.
PRs affecting the public API, including adding new features, should update the public documentation.
If possible, add appropriate tests to your PR.
Feel free to open a PR early on in the development process, we are happy to help in the development process and provide feedback along the way.

## Licensing terms

All source code in this package is released under the 2-clause BSD license.
All of `JAXbind` is distributed *without any warranty*.

## Citing JAXbind

To cite `JAXbind`, please use the citation provided below.

```
@article{jaxbind,
    title = {JAXbind: Bind any function to JAX},
    author = {Jakob Roth and Martin Reinecke and Gordian Edenhofer},
    year = {2024},
    journal = {Journal of Open Source Software},
    publisher = {The Open Journal},
    volume = {9},
    number = {98},
    pages = {6532},
    doi = {10.21105/joss.06532},
    url = {https://doi.org/10.21105/joss.06532},
}
```
