Repository: NIFTy-PPL/JAXbind (language: python)

Issue:
Merge paper sources into multi_arg
Just for the purpose of general branch cleanup


Discussion / hints:

Fix patch:
diff --git a/.github/actions/setup_package/action.yml b/.github/actions/setup_package/action.yml
new file mode 100644
index 0000000..ae20407
--- /dev/null
+++ b/.github/actions/setup_package/action.yml
@@ -0,0 +1,20 @@
+name: Reusable setup for the python environment and the package installation
+runs:
+  using: composite
+  steps:
+  - uses: actions/checkout@v4
+  - name: Set up Python 3.10
+    uses: actions/setup-python@v5
+    with:
+      python-version: "3.10"
+      cache: "pip"
+  - name: Install dependencies
+    shell: bash
+    run: |
+      python -m pip install --upgrade pip
+      pip install setuptools pybind11 pytest numpy scipy ducc0
+      pip install "jax[cpu]"
+  - name: Install package
+    shell: bash
+    run: |
+      pip install --user .
diff --git a/.github/workflows/build_wheels.yml b/.github/workflows/build_wheels.yml
new file mode 100644
index 0000000..f65f111
--- /dev/null
+++ b/.github/workflows/build_wheels.yml
@@ -0,0 +1,47 @@
+name: Build all wheels
+on: [workflow_dispatch]
+jobs:
+  build_wheels:
+    name: Build wheels on ${{ matrix.os }}
+    strategy:
+      matrix:
+        os: [ubuntu-22.04, ubuntu-latest, windows-2022, windows-latest, macos-13, macos-14]
+    runs-on: ${{ matrix.os }}
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: '3.11'
+      - uses: pypa/cibuildwheel@v2.17.0
+        env:
+          CIBW_ARCHS: auto64
+          CIBW_SKIP: cp36-* cp37-* cp38-* pp* *-musllinux*
+          CIBW_TEST_REQUIRES: pytest
+          CIBW_TEST_COMMAND: python -m pytest -x {project}/test
+          CIBW_TEST_SKIP: cp36-* cp37-* cp38-* cp39-* cp310-* cp311-*
+          CIBW_MANYLINUX_X86_64_IMAGE: manylinux_2_28
+          CIBW_MANYLINUX_I686_IMAGE: manylinux_2_28
+          CIBW_MANYLINUX_PYPY_X86_64_IMAGE: manylinux_2_28
+          CIBW_MANYLINUX_AARCH64_IMAGE: manylinux_2_28
+          CIBW_MANYLINUX_PPC64LE_IMAGE: manylinux_2_28
+          CIBW_MANYLINUX_S390X_IMAGE: manylinux_2_28
+          CIBW_MANYLINUX_PYPY_AARCH64_IMAGE: manylinux_2_28
+          CIBW_MANYLINUX_PYPY_I686_IMAGE: manylinux_2_28
+        with:
+          package-dir: .
+          output-dir: wheelhouse
+      - uses: actions/upload-artifact@v4
+        with:
+          name: cibw-wheels-${{ matrix.os }}-${{ strategy.job-index }}
+          path: ./wheelhouse/*.whl
+  build_sdist:
+    name: Build source distribution
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - name: Build sdist
+        run: pipx run build --sdist
+      - uses: actions/upload-artifact@v4
+        with:
+          name: cibw-sdist
+          path: dist/*.tar.gz
diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..bf9f3c2
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,107 @@
+# from https://github.com/github/gitignore/blob/master/Python.gitignore
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.c
+*.o
+*.so
+
+# Distribution / packaging
+.Python
+env/
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+*.egg-info/
+.installed.cfg
+*.egg
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*,cover
+.hypothesis/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/build/
+docs/source/mod
+
+# PyBuilder
+target/
+
+# IPython Notebook
+.ipynb_checkpoints
+*.ipynb
+
+# pyenv
+.python-version
+
+# celery beat schedule file
+celerybeat-schedule
+
+# dotenv
+.env
+
+# virtualenv
+venv/
+ENV/
+
+# Spyder project settings
+.spyderproject
+
+# Rope project settings
+.ropeproject
+
+# latex
+*.aux
+*.bbl
+*.bcf
+*.blg
+*.fdb_latexmk
+*.fls
+*.out
+*.pdf
+*.xml
+*.gz
diff --git a/LICENSE b/LICENSE
new file mode 100644
index 0000000..1fd4aea
--- /dev/null
+++ b/LICENSE
@@ -0,0 +1,24 @@
+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.
\ No newline at end of file
diff --git a/README.md b/README.md
index 8384caa..6e5eaed 100644
--- a/README.md
+++ b/README.md
@@ -1,10 +1,138 @@
-Simple linear operators for Jax
-===============================
+# JAXbind: Easy bindings to JAX
 
-This package provides convenience functionality which allows easy integration
-of arbitrary linear operations into Jax.
+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)
 
-### Requirements
+## Summary
+
+The existing interface in JAX for connecting 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 supporting any JAX transformations.
+
+### 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
+
+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/multi_arg/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/multi_arg/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/multi_arg/demos/03_nonlinear_function.py) demonstrates the interface for non-linear functions and shows how to deal with fixed arguments, which cannot be differentiated.
+
+
+## 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!
+
+## Requirements
 
 - [Python >= 3.8](https://www.python.org/)
 - only when compiling from source: [pybind11](https://github.com/pybind/pybind11)
@@ -15,12 +143,30 @@ of arbitrary linear operations into Jax.
   - Intel `icpx` (oneAPI compiler series). (Note that the older `icpc` compilers
     are not supported.)
 
-### Sources
+## Installation
+
+Binary wheels for JAXbind can be obtained and installed from PyPI via:
+
+```
+pip install --user jaxbind
+```
+
+To install JAXbind from source clone the repository and install JAXbind via pip.
+
+```
+git clone https://github.com/NIFTy-PPL/jaxbind.git
+cd jaxbind
+pip install --user .
+```
 
-The latest version of jax_linop can be obtained by cloning the repository via
+## TODOs
 
-    git clone https://gitlab.mpcdf.mpg.de/mtr/jax_linop.git
+* Paper
+  * final editing
+* PiPy release
+* README
 
-### Licensing terms
+## Licensing terms
 
-All source code in this package is released under the 3-clause BSD license.
+All source code in this package is released under the 2-clause BSD license.
+All of JAXbind is distributed *without any warranty*.
diff --git a/demos/01_linear_function.py b/demos/01_linear_function.py
new file mode 100644
index 0000000..28c4b8d
--- /dev/null
+++ b/demos/01_linear_function.py
@@ -0,0 +1,193 @@
+# SPDX-License-Identifier: BSD-2-Clause
+# Authors: Martin Reinecke, Jakob Roth, Gordian Edenhofer
+
+# Copyright(C) 2024 Max-Planck-Society
+
+# %%
+import jax
+import jax.numpy as jnp
+import scipy
+from jax import random
+
+import jaxbind
+
+jax.config.update("jax_enable_x64", True)
+
+# %% [markdown]
+
+# # Binding Scipy's FFT to JAX
+
+# This demo showcases the use of the interface of JAXbind for binding linear
+# functions to jax. Specifically, it wraps scipy's FFT as a JAX primitive.
+
+# The code wraps the scipy.fft.fftn function as a JAX primitive using JAXbind.
+# It defines the fftn function with the required signature for JAXbind. The
+# fftn function takes three arguments: out, args, and kwargs_dump. It extracts
+# the input array from args, deserializes the keyword arguments from
+# kwargs_dump, and computes the FFT using scipy.fft.fftn. Similarly, it defines
+# the fftn_transposed function for the linear transposed operation. The code
+# also includes abstract evaluation functions for fftn and fftn_transposed.
+# Finally, it explains the purpose of wrapping scipy.fft.fftn as a JAX
+# primitive.
+
+# All native JAX primitives can be batched via 'jax.vmap'. Primitives
+# registered via JAXbind are no exception here. By default, JAXbind performs
+# the batching by sequential applying the function along the batching axis.
+# However, it also exposes an interface to allow the user to perform the
+# batching themself. As the FFT natively support mapping over input axis and
+# yields significant speedups compared to a sequential computation, we
+# demonstrate in the following how JAXbind can be used to register a custom
+# batching.
+
+# %%
+
+
+def fftn(out, args, kwargs_dump):
+    # extract the input for the fft form the input tuple
+    (x,) = args
+
+    # deserialize keyword arguments
+    kwargs = jaxbind.load_kwargs(kwargs_dump)
+    # extract keyword argument which can be given to the JAX primitive
+    workers = kwargs.pop("workers", None)
+    # extract the axes over which the FFT should be batched. This is only
+    # necessary when supporting custom batching.
+    batch_axes = kwargs.pop("batch_axes", None)
+
+    # translate the batch_axes into the axes argument for scipy.fft.fftn
+    axes = list(range(len(x.shape)))
+    if batch_axes:
+        axes = [i for i in range(len(x.shape)) if not i in batch_axes[0]]
+
+    # compute the fft and write the result in the out tuple
+    out[0][()] = scipy.fft.fftn(x, axes=axes, norm="forward", workers=workers)
+
+
+# %% [markdown]
+
+# In addition to applying the FFT, JAXbind also requires the implementation
+# of the linear transposed function. The syntax for computing the linear
+# transposed is identical to fftn.
+
+# %%
+
+
+def fftn_transposed(out, args, kwargs_dump):
+    (x,) = args
+    kwargs = jaxbind.load_kwargs(kwargs_dump)
+
+    workers = kwargs.pop("workers", None)
+    batch_axes = kwargs.pop("batch_axes", None)
+
+    axes = list(range(len(x.shape)))
+    if batch_axes:
+        axes = [i for i in range(len(x.shape)) if not i in batch_axes[0]]
+
+    # JAX need the transposed function and not he adjoint (transposed+complex
+    # conjugated). Since scipy.fft.ifftn is the adjoint of scipy.fft.fftn we
+    # have to undo the complex conjugation.
+    out[0][()] = scipy.fft.ifftn(
+        x.conj(), axes=axes, norm="backward", workers=workers
+    ).conj()
+
+
+# %% [markdown]
+
+# JAX needs to abstractly evaluate the code, thus needs to be able to
+# evaluate the shape and dtype of the output of a function given the shape
+# and dtype of the input. For this we have to provide the abstract eval
+# functions for fftn and fftn_transposed.
+
+# The abstract eval functions take normal arguments and keyword arguments
+# and return a tuple containing the output information for each output
+# argument of the function. Since fftn has only one output argument the
+# output tuple of the abstract eval function has length one.
+
+# The output description of each output argument is also a tuple. The
+# first entry in the tuple contains the shape of the output array. The
+# second entry is the dtype of this array. The third entry in the tuple
+# is only required for functions supporting custom batching and indicates
+# the batching axis of the output of the function (thus fftn). In our case
+# the batching axis of the output is identical to the batching axis of the
+# input.
+
+
+# %%
+
+
+def fftn_abstract_eval(*args, **kwargs):
+    # extract the input
+    (x,) = args
+
+    # extract potential batching axis
+    batch_axes = kwargs.pop("batch_axes", None)
+
+    # indicate the output batching axis if fftn is batched
+    out_ax = ()
+    if batch_axes:
+        if len(batch_axes[0]) > 0:  # check if function is batched
+            # check along which axis the fftn is batched. If fftn was batched
+            # multiple times take the last batching axis
+            out_ax = batch_axes[0][-1]
+
+    # return shape dtype and potential batching axis of output
+    return ((x.shape, x.dtype, out_ax),)
+
+
+# %% [markdown]
+
+# JAX also needs to abstractly evaluate the transposed function. For that we
+# have to provide the same information as for fftn. Since a fft is not changing
+# the shape or dtype this function is identical to the fftn_abstract_eval. For
+# general linear functions this might be different.
+
+# %%
+
+
+def fftn_transposed_abstract_eval(*args, **kwargs):
+    (a,) = args
+    batch_axes = kwargs.pop("batch_axes", None)
+    out_ax = ()
+    if batch_axes:
+        if len(batch_axes[0]) > 0:
+            out_ax = batch_axes[0][-1]
+    return ((a.shape, a.dtype, out_ax),)
+
+
+# %% [markdown]
+
+# Now we register our function as a custom JAX primitive using the interface for
+# linear functions of JAXbind. JAXbind returns the resulting JAX primitive.
+
+# %%
+fftn_jax = jaxbind.get_linear_call(
+    fftn,
+    fftn_transposed,
+    fftn_abstract_eval,
+    fftn_transposed_abstract_eval,
+    func_can_batch=True,  # indicate that our function supports custom batching
+)
+
+
+# generate some random input to showcase the use of the newly register JAX primitive
+key = random.PRNGKey(42)
+key, subkey = random.split(key)
+inp = jax.random.uniform(subkey, shape=(10, 10), dtype=jnp.float64)
+inp = inp + 1j * jax.random.uniform(subkey, shape=(10, 10), dtype=jnp.float64)
+
+
+# apply the new primitive
+res = fftn_jax(inp)
+
+# apply the new primitive and pass the keyword argument "workers=2" to the scipy fft
+res2 = fftn_jax(inp, workers=2)
+
+# jit compile the new primitive
+fftn_jax_jit = jax.jit(fftn_jax)
+res_jit = fftn_jax_jit(inp)
+
+# vmap fft_jax over the fist axis of the input
+res_vmap = jax.vmap(fftn_jax, in_axes=0)
+
+# compute the jvp of fftn
+res_jvp = jax.jvp(fftn_jax, (inp,), (inp,))
diff --git a/demos/02_multilinear_function.py b/demos/02_multilinear_function.py
new file mode 100644
index 0000000..0890ae6
--- /dev/null
+++ b/demos/02_multilinear_function.py
@@ -0,0 +1,123 @@
+# SPDX-License-Identifier: BSD-2-Clause
+# Authors: Martin Reinecke, Jakob Roth, Gordian Edenhofer
+
+# Copyright(C) 2024 Max-Planck-Society
+
+# %%
+import jax
+import jax.numpy as jnp
+from jax import random
+
+
+import jaxbind
+
+jax.config.update("jax_enable_x64", True)
+
+# %% [markdown]
+
+# # Binding a multi-linear function to JAX
+
+# This demo showcases the use of JAXbind for binding multi-linear functions to
+# JAX. As an example we bind the python function mlin computing
+# (x,y) -> (x*y, x*y) to a JAX primitive. Note: multi-linear functions can also
+# be regarded als general non-linear functions. For the JAXbind interface for
+# non-linear functions see the 'demo_nonlin.py' and the docstring of the
+# jaxbind.get_nonlinear_cal'. Additional information for linear functions can
+# also be found in 'demo_scipy_fft.py'.
+
+
+# %%
+
+
+def mlin(out, args, kwargs_dump):
+    # extract the input from the input tuple
+    x, y = args[0], args[1]
+
+    # do the computation and write result in the out tuple
+    out[0][()] = x * y
+    out[1][()] = x * y
+
+
+# %% [markdown]
+
+# Besides the application of the function ('mlin') itself, JAXbind required the
+# linear transposed of the partial derivatives of 'mlin'.
+# %%
+
+
+# linear transpose of the partial derivative of 'mlin' with respect to the fist
+# variable x.
+def mlin_T1(out, args, kwargs_dump):
+    y, da, db = args[0], args[1], args[2]
+    out[0][()] = y * da + y * db
+
+
+# linear transpose of the partial derivative of 'mlin' with respect to the second
+# variable y.
+def mlin_T2(out, args, kwargs_dump):
+    x, da, db = args[0], args[1], args[2]
+    out[0][()] = x * da + x * db
+
+
+# %% [markdown]
+
+# JAX needs to abstractly evaluate the code, thus needs to be abel to evaluate
+# the shape and dtype of the output of a function given the shape and dtype of
+# the input. For this we have to provide the abstract eval functions for mlin,
+# mlin_T1, and mlin_T2. The abstract evaluations functions return for each
+# output argument a tuple containing the shape and dtype of this output. More
+# details are in the 'demo_scipy_fft.py'
+
+
+# %%
+def mlin_abstract(*args, **kwargs):
+    assert args[0].shape == args[1].shape
+    return (
+        (args[0].shape, args[0].dtype),
+        (args[0].shape, args[0].dtype),
+    )
+
+
+def mlin_abstract_T1(*args, **kwargs):
+    assert args[0].shape == args[1].shape
+    return ((args[0].shape, args[0].dtype),)
+
+
+def mlin_abstract_T2(*args, **kwargs):
+    assert args[0].shape == args[1].shape
+    return ((args[0].shape, args[0].dtype),)
+
+
+# %% [markdown]
+
+# Now we can register the JAX primitive corresponding to the python function
+# mlin.
+
+# %%
+
+func_T = (mlin_T1, mlin_T2)
+func_abstract_T = (mlin_abstract_T1, mlin_abstract_T2)
+mlin_jax = jaxbind.get_linear_call(
+    mlin,
+    func_T,
+    mlin_abstract,
+    func_abstract_T,
+)
+
+# generate some random input to showcase the use of the newly register JAX primitive
+key = random.PRNGKey(42)
+key, subkey = random.split(key)
+inp0 = jax.random.uniform(subkey, shape=(10, 10), dtype=jnp.float64)
+key, subkey = random.split(key)
+inp1 = jax.random.uniform(subkey, shape=(10, 10), dtype=jnp.float64)
+inp = (inp0, inp1)
+
+# apply the new primitive
+res = mlin_jax(*inp)
+
+# jit compile the new primitive
+mlin_jit = jax.jit(mlin_jax)
+res_jit = mlin_jit(*inp)
+
+# compute the jvp
+res_jvp = jax.jvp(mlin_jit, inp, inp)
diff --git a/demos/03_nonlinear_function.py b/demos/03_nonlinear_function.py
new file mode 100644
index 0000000..2576f8a
--- /dev/null
+++ b/demos/03_nonlinear_function.py
@@ -0,0 +1,97 @@
+# SPDX-License-Identifier: BSD-2-Clause
+# Authors: Martin Reinecke, Jakob Roth, Gordian Edenhofer
+
+# Copyright(C) 2024 Max-Planck-Society
+
+# %%
+from functools import partial
+
+import jax
+from jax import numpy as jnp
+
+import jaxbind
+
+jax.config.update("jax_enable_x64", True)
+
+
+# %% [markdown]
+
+# # Binding non-linear functions to JAX with fixed arguments
+
+# Bind a nonlinear function to JAX using JAXbind. Some input parts are fixed
+# (nondifferentiable) to showcase JAXbind's handling of fixed arguments.
+
+# The script begins by defining the nonlinear function using the JAXbind
+# compatible signature of out, args, and kwargs_dump. The derivative
+# function, which is registered with JAX, has the same signature as the
+# nonlinear function but includes an additional variable for the tangent of
+# the variable being differentiated. Additionally, we define the transpose of
+# the derivative function, as JAX may transpose it for operations like
+# retrieving the VJP.
+
+# %%
+
+
+def nonlin(out, args, kwargs_dump):
+    # (x,y) -> (xy, y**2)
+    x, y = args
+    out[0][()] = x * y
+    out[1][()] = y * y
+
+
+def nonlin_deriv(out, args, kwargs_dump):
+    # (x,y,dy) -> (ydx + xdy, 2 * y * dy)
+    x, y, dy = args
+    out[0][()] = x * dy
+    out[1][()] = 2 * y * dy
+
+
+def nonlin_deriv_T(out, args, kwargs_dump):
+    # (x, y, da, db) -> (yda, xda + 2ydb)
+    x, y, da, db = args
+    out[0][()] = x * da + 2 * y * db
+
+
+# %% [markdown]
+
+# Define the abstract evaluation functions for JAX that translate input
+# shape and dtypes to output shape and dtypes.
+
+# %%
+
+
+def nonlin_abstract(*args, **kwargs):
+    # Returns `shape` and `dtype` of output
+    x, y, *_ = args
+    assert x.shape == x.shape and x.dtype is y.dtype
+    return ((x.shape, x.dtype), (x.shape, x.dtype))
+
+
+def nonlin_abstract_T(*args, **kwargs):
+    # Returns `shape` and `dtype` of output
+    a, b, da, db = args
+    assert a.shape == b.shape == da.shape == db.shape
+    assert a.dtype is b.dtype and a.dtype is da.dtype and a.dtype is db.dtype
+    return ((a.shape, a.dtype),)
+
+
+# %%
+nonlin_jax = jaxbind.get_nonlinear_call(
+    nonlin,
+    (nonlin_deriv, nonlin_deriv_T),
+    nonlin_abstract,
+    nonlin_abstract_T,
+    # Tell JAXbind that the first parameter to `nonlin` (and derived functions)
+    # is not to be differentiated.
+    first_n_args_fixed=1,
+    func_can_batch=True,
+)
+
+inp_f = 4 + jnp.zeros((1,))
+inp = 6 + jnp.zeros((1,))
+nonlin_jax_pt = partial(nonlin_jax, inp_f)
+
+_ = jax.jvp(nonlin_jax_pt, (inp,), (inp,))
+_, nonlin_jax_pt_vjp = jax.vjp(nonlin_jax_pt, inp)
+nonlin_jax_pt_vjp = jax.jit(nonlin_jax_pt_vjp)
+_ = nonlin_jax_pt_vjp([inp, inp])
diff --git a/demos/demo.py b/demos/demo.py
deleted file mode 100644
index 1e33d04..0000000
--- a/demos/demo.py
+++ /dev/null
@@ -1,39 +0,0 @@
-import jax
-import numpy as np
-import scipy.fft
-from jax.test_util import check_grads
-
-import jax_linop
-
-jax.config.update("jax_enable_x64", True)
-
-
-def fft_operator(axes):
-    def fft(inp, out, state):
-        # we don't want normalization, just transpose, hence the strange norms
-        out[()] = scipy.fft.fftn(inp, axes=state["axes"], norm="backward")
-
-    def fft_T(inp, out, state):
-        out[()
-           ] = scipy.fft.ifftn(inp.conj(), axes=state["axes"],
-                               norm="forward").conj()
-
-    def fft_abstract(shape, dtype, state):
-        return shape, dtype
-
-    return jax_linop.get_linear_call(
-        fft, fft_T, fft_abstract, fft_abstract, axes=tuple(axes)
-    )
-
-
-# create an FFT operator that transforms the first axis of its input array,
-# as well as its adjoint operator
-fft_op = fft_operator(axes=(0, ))
-
-rng = np.random.default_rng(42)
-a = rng.random((100, 20, 30)) - 0.5 + 1j * (rng.random((100, 20, 30)) - 0.5)
-
-fft_op_T = jax.linear_transpose(fft_op, a)
-
-check_grads(fft_op, (a, ), order=2, modes=["fwd", "rev"], eps=1.)
-check_grads(fft_op_T, ([a], ), order=2, modes=["fwd", "rev"], eps=1.)
diff --git a/demos/s2fft_demo.py b/demos/s2fft_demo.py
deleted file mode 100644
index c124afb..0000000
--- a/demos/s2fft_demo.py
+++ /dev/null
@@ -1,251 +0,0 @@
-import jax
-
-from jax import config
-config.update("jax_enable_x64", True)
-
-import ducc0
-import numpy as np
-import s2fft
-import jax_linop
-from time import time
-
-r2cdict = {np.dtype(np.float32) : np.dtype(np.complex64),
-           np.dtype(np.float64) : np.dtype(np.complex128)}
-
-c2rdict = {np.dtype(np.complex64) : np.dtype(np.float32),
-           np.dtype(np.complex128) : np.dtype(np.float64)}
-
-def isrealtype(dtype):
-    return np.dtype(dtype) in r2cdict
-
-def iscomplextype(dtype):
-    return np.dtype(dtype) in c2rdict
-
-def isvalidtype(dtype):
-    return isrealtype(dtype) or iscomplextype(dtype)
-
-def realtype(dtype):
-    return c2rdict[np.dtype(dtype)]
-
-def complextype(dtype):
-    return r2cdict[np.dtype(dtype)]
-
-def ssht_operator(L, sampling, spin, reality, nthreads, nside=0):
-    def get_config(alm, map, state):
-        if not iscomplextype(alm.dtype):
-            raise RuntimeError("invalid alm data type")
-        if not isvalidtype(map.dtype):
-            raise RuntimeError("invalid map data type")
-        reality = state["reality"]
-        L = state["L"]
-        spin = state["spin"]
-        nthreads = state["nthreads"]
-        sampling = state["sampling"]
-        hpxparam={}
-        if sampling == "mw":
-            geometry = "MW"
-            mapshape = (L, 2*L-1)
-        elif sampling == "mwss":
-            geometry = "CC"
-            mapshape = (L+1, 2*L)
-        elif sampling == "dh":
-            geometry = "F1"
-            mapshape = (2*L, 2*L-1)
-        elif sampling == "healpix":
-            geometry = "healpix"
-            mapshape = (12*state["nside"]**2,)
-            base = ducc0.healpix.Healpix_Base(state["nside"], "RING")
-            hpxparam = base.sht_info()
-        if map.shape!=mapshape:
-            raise RuntimeError("bad map shape")
-        if alm.shape!=(L,2*L-1):
-            raise RuntimeError("bad alm shape")
-        return L, spin, nthreads, reality, geometry, mapshape, hpxparam
-
-    def ssht_T(inp, out, state):
-        L, spin, nthreads, reality, geometry, mapshape, hpxparam = get_config(out, inp, state)
-
-        if spin==0:
-            inp2 = inp.reshape((1, *mapshape))
-            inp2 = inp2.conj()
-            tmp = np.empty((1 if reality else 2, (L*(L+1))//2), out.dtype)
-            if sampling == "healpix":
-                ducc0.sht.adjoint_synthesis(lmax=L-1, spin=np.abs(spin), map=inp2.real, alm=tmp[0:1], nthreads=nthreads, **hpxparam)
-                if not reality:
-                    ducc0.sht.adjoint_synthesis(lmax=L-1, spin=np.abs(spin), map=inp2.imag, alm=tmp[1:2], nthreads=nthreads, **hpxparam)
-            else:
-                ducc0.sht.adjoint_synthesis_2d(
-                    lmax=L-1, spin=np.abs(spin),
-                    map=inp2.real, alm=tmp[0:1], nthreads=nthreads,
-                    geometry=geometry)
-                if not reality:
-                    ducc0.sht.adjoint_synthesis_2d(
-                        lmax=L-1, spin=np.abs(spin),
-                        map=inp2.imag, alm=tmp[1:2], nthreads=nthreads,
-                        geometry=geometry)
-        else:
-            if spin > 0:
-                tmap = inp.view(realtype(out.dtype)).reshape((*mapshape,2))
-                tmap = np.moveaxis(tmap,-1,0)
-            else:
-                tmap = np.empty((2, inp.shape[0], inp.shape[1]), dtype=realtype(inp.dtype))
-                tmap[0] = -inp.real
-                tmap[1] = inp.imag
-            tmap = tmap.copy()
-            tmap[1] *= -1
-            if sampling == "healpix":
-                tmp = ducc0.sht.adjoint_synthesis(lmax=L-1, spin=np.abs(spin), map=tmap, nthreads=nthreads, **hpxparam)
-            else:
-                tmp = ducc0.sht.adjoint_synthesis_2d(
-                    lmax=L-1, spin=np.abs(spin),
-                    map=tmap, nthreads=nthreads,
-                    geometry=geometry)
-            if spin > 0:
-                tmp *= -1
-        ducc0.sht.experimental.alm2flm(tmp, spin, out)
-        out.imag *= -1
-
-    def ssht(inp, out, state):
-        L, spin, nthreads, reality, geometry, mapshape, hpxparam = get_config(inp, out, state)
-
-        tmp = ducc0.sht.experimental.flm2alm(inp, spin, real=reality)
-        out2 = out.view(realtype(inp.dtype))
-        out2 = out2.reshape((*mapshape, 1 if reality else 2))
-        out2 = np.moveaxis(out2, -1, 0)
-        if spin==0:
-            if sampling == "healpix":
-                ducc0.sht.synthesis(lmax=L-1, spin=np.abs(spin), alm=tmp[0:1], map=out2[0:1], nthreads=nthreads, **hpxparam)
-                if not reality:
-                    ducc0.sht.synthesis(lmax=L-1, spin=np.abs(spin), alm=tmp[1:2], map=out2[1:2], nthreads=nthreads, **hpxparam)
-            else:
-                ducc0.sht.synthesis_2d(
-                    lmax=L-1, spin=np.abs(spin),
-                    alm=tmp[0:1], map=out2[0:1], nthreads=nthreads, geometry=geometry)
-                if not reality:
-                    ducc0.sht.synthesis_2d(
-                        lmax=L-1, spin=np.abs(spin),
-                        alm=tmp[1:2], map = out2[1:2], nthreads=nthreads, geometry=geometry)
-        else:
-            if sampling == "healpix":
-                ducc0.sht.synthesis(lmax=L-1, spin=np.abs(spin), alm=tmp, map=out2, nthreads=nthreads, **hpxparam)
-            else:
-                ducc0.sht.synthesis_2d(
-                    lmax=L-1, spin=np.abs(spin),
-                    alm=tmp, map=out2, nthreads=nthreads,
-                    geometry=geometry)
-            if spin > 0:
-                out *= -1
-            else:
-                out.real *= -1
-
-    def get_config2(shape, dtype, synthesis, state):
-        if synthesis and not iscomplextype(dtype):
-            raise RuntimeError("invalid alm data type")
-        if not isvalidtype(dtype):
-            raise RuntimeError("invalid map data type")
-        reality = state["reality"]
-        L = state["L"]
-        spin = state["spin"]
-        sampling = state["sampling"]
-        if sampling == "mw":
-            mapshape = (L, 2*L-1)
-        elif sampling == "mwss":
-            mapshape = (L+1, 2*L)
-        elif sampling == "dh":
-            mapshape = (2*L, 2*L-1)
-        elif sampling == "healpix":
-            mapshape = (12*state["nside"]**2,)
-        if synthesis:
-            if shape != (L,2*L-1):
-                raise RuntimeError("bad alm shape")
-        else:
-            if shape != mapshape:
-                raise RuntimeError("bad map shape")
-        return L, spin, reality, mapshape
-
-    def ssht_abstract(shape_in, dtype_in, state):
-        L, spin, reality, mapshape = get_config2(shape_in, dtype_in, True, state)
-        
-        shape_out = mapshape
-        dtype_out = dtype_in if not reality else realtype(dtype_in)
-        return shape_out, dtype_out
-
-    def ssht_abstract_T(shape_in, dtype_in, state):
-        L, spin, reality, mapshape = get_config2(shape_in, dtype_in, False, state)
-
-        if isrealtype(dtype_in):
-            if (not reality):
-                raise RuntimeError("real map provided, but complex map expected")
-            dtype_out = complextype(dtype_in)
-        else:
-            dtype_out = dtype_in
-        shape_out = (L, 2*L-1)
-        return shape_out, dtype_out
-
-    L = int(L)
-    if (L<=0):
-        raise ValueError("L must be positive")
-    spin = int(spin)
-    reality = bool(reality)
-    if spin != 0 and reality:
-        raise ValueError("if spin is nonzero, reality must be False")
-    sampling = str(sampling)
-    if sampling!="mw" and sampling!="dh" and sampling!="mwss" and sampling!="healpix":
-        raise ValueError("unsupported sampling type")
-    return jax_linop.get_linear_call(
-        ssht, ssht_T, ssht_abstract, ssht_abstract_T,
-        L=L, spin=spin,
-        sampling=sampling,
-        reality=reality,
-        nthreads=int(nthreads),
-        nside=int(nside))
-
-
-from jax import config
-config.update("jax_enable_x64", True)
-import jax.numpy as jnp
-from jax.test_util import check_grads
-
-
-L=1024
-lmax = L-1
-spin=0
-nthreads=1
-reality=False
-rng = np.random.default_rng(42)
-sampling="mwss"
-nside=L//2
-
-flm = s2fft.utils.signal_generator.generate_flm(rng, L, 0, spin, reality)
-
-op = ssht_operator(L=L, sampling=sampling, spin=spin, nthreads=nthreads, reality=reality, nside=nside)
-op = jax.jit(op)
-#ducc0.misc.preallocate_memory(8)
-t0 = time()
-for i in range(100):
-    op(flm)
-map0 = np.array(op(flm)[0])
-print("ducc time:", time()-t0)
-#exit()
-def func2(flm):
-    f = op(flm)[0]
-    return jnp.sum(jnp.abs(f)**2)
-t0 = time()
-check_grads(func2, (flm,), order=3, modes=('fwd','rev'))
-print("check_grad ducc time (fwd+rev):", time()-t0)
-
-t0 = time()
-precomps = s2fft.generate_precomputes_jax(L, forward=False, spin=spin, sampling=sampling, nside=nside)
-print("precomp time:", time()-t0)
-
-t0 = time()
-map1 = np.array(s2fft.inverse(flm, L, spin=spin, reality=reality, sampling=sampling, precomps=precomps, method="jax", nside=nside))
-print("s2fft time:", time()-t0)
-print(ducc0.misc.l2error(map0, map1))
-
-def func1(flm):
-    f = s2fft.inverse(flm, L, reality=reality, precomps=precomps, spin=spin, sampling=sampling,nside=nside, method="jax")
-    return jnp.sum(jnp.abs(f)**2)
-t0 = time()
-check_grads(func1, (flm,), order=3, modes=('rev'))
-print("check_grad s2fft time (rev):", time()-t0)
diff --git a/jax_linop/__init__.py b/jax_linop/__init__.py
deleted file mode 100644
index 369981b..0000000
--- a/jax_linop/__init__.py
+++ /dev/null
@@ -1,1 +0,0 @@
-from .jax_linop import get_linear_call
diff --git a/jax_linop/jax_linop.py b/jax_linop/jax_linop.py
deleted file mode 100644
index 901709d..0000000
--- a/jax_linop/jax_linop.py
+++ /dev/null
@@ -1,169 +0,0 @@
-# SPDX-License-Identifier: BSD-2-Clause
-# Authors: Martin Reinecke, Jakob Roth, Gordian Edenhofer
-
-# Copyright(C) 2023, 2024 Max-Planck-Society
-
-from functools import partial
-
-import _jax_linop
-import jax
-import numpy as np
-from jax.interpreters import ad, mlir, batching
-from jaxlib.hlo_helpers import custom_call
-
-# incremented for every registered operator, strictly for uniqueness purposes
-_global_opcounter = 0
-
-for _name, _value in _jax_linop.registrations().items():
-    jax.lib.xla_client.register_custom_call_target(
-        _name, _value, platform="cpu"
-    )
-
-
-def _from_id(objectid):
-    import ctypes
-
-    return ctypes.cast(objectid, ctypes.py_object).value
-
-
-def _exec_abstract(x, *, stateid, stateTid):
-    state = _from_id(stateid)
-    shp, tp = state["_func_abstract"](x.shape, x.dtype, state)
-    return (jax.core.ShapedArray(shp, tp), )
-
-
-# the values are explained in src/duc0/bindings/typecode.h
-_dtype_dict = {
-    np.dtype(np.float32): 3,
-    np.dtype(np.float64): 7,
-    np.dtype(np.complex64): 67,
-    np.dtype(np.complex128): 71,
-}
-
-
-def _lowering(ctx, x, *, platform="cpu", stateid, stateTid):
-    state = _from_id(stateid)
-    if len(ctx.avals_in) != 1:
-        raise RuntimeError("need exactly one input object")
-    shape_in = ctx.avals_in[0].shape
-    dtype_in = ctx.avals_in[0].dtype
-    if len(ctx.avals_out) != 1:
-        raise RuntimeError("need exactly one output object")
-    shape_out, dtype_out = state["_func_abstract"](shape_in, dtype_in, state)
-
-    dtype_out_mlir = mlir.dtype_to_ir_type(dtype_out)
-    jaxtype_out = mlir.ir.RankedTensorType.get(shape_out, dtype_out_mlir)
-    layout_in = tuple(range(len(shape_in) - 1, -1, -1))
-    layout_out = tuple(range(len(shape_out) - 1, -1, -1))
-
-    irc = mlir.ir_constant
-    operands = [
-        x, irc(state["_opid"]), irc(stateid), irc(_dtype_dict[dtype_in]),
-        irc(len(shape_in))] + [irc (i) for i in shape_in] + \
-        [irc(_dtype_dict[dtype_out]), irc(len(shape_out))] + \
-        [irc (i) for i in shape_out]
-    operand_layouts = [layout_in] + [()] * (6 + len(shape_in) + len(shape_out))
-
-    if platform == "cpu":
-        return custom_call(
-            platform + "_pycall",
-            result_types=[jaxtype_out],
-            result_layouts=[layout_out],
-            operands=operands,
-            operand_layouts=operand_layouts,
-        ).results
-    elif platform == "gpu":
-        raise ValueError("No GPU support")
-    raise ValueError("Unsupported platform; this must be either 'cpu' or 'gpu'")
-
-
-def _jvp(args, tangents, *, stateid, stateTid):
-    res = _prim.bind(args[0], stateid=stateid, stateTid=stateTid)
-    return (
-        res,
-        jax.lax.zeros_like_array(res) if type(tangents[0]) is ad.Zero else
-        _prim.bind(tangents[0], stateid=stateid, stateTid=stateTid),
-    )
-
-
-def _transpose(cotangents, args, *, stateid, stateTid):
-    return _prim.bind(cotangents[0], stateid=stateTid, stateTid=stateid)
-
-
-def _batch(args, axes, *, stateid, stateTid):
-    from .custom_map import smap
-
-    return smap(
-        partial(_prim.bind, stateid=stateid, stateTid=stateTid),
-        in_axes=axes[0],
-        out_axes=axes[0]
-    )(*args), axes
-
-
-_prim = jax.core.Primitive("jax_linop_prim")
-_prim.multiple_results = True
-_prim.def_impl(partial(jax.interpreters.xla.apply_primitive, _prim))
-_prim.def_abstract_eval(_exec_abstract)
-
-for platform in ["cpu", "gpu"]:
-    mlir.register_lowering(
-        _prim, partial(_lowering, platform=platform), platform=platform
-    )
-    ad.primitive_jvps[_prim] = _jvp
-    ad.primitive_transposes[_prim] = _transpose
-    batching.primitive_batchers[_prim] = _batch
-
-
-def _call(x, state, stateT):
-    return _prim.bind(x, stateid=id(state), stateTid=id(stateT))
-
-
-def get_linear_call(func, func_T, /, func_abstract, func_abstract_T, **kwargs):
-    """Create Jax functions for the provided linear operator
-
-    Parameters
-    ----------
-    func, func_T : linear function respectively its transpose
-        The function signature must be (inp, out, state), where
-        inp and out are numpy.ndarrays of float[32/64] or complex[64/128] type,
-        and state is a dictionary containing additional information that the
-        operator might need.
-    func_abstract, func_abstract_T : function respectively its tranpose
-        computing the shape and dtype of the operator's output from shape and dtype of its input
-        Its signature must be (shape, dtype, state), where `state` is analogous
-        to the one from `func`, `shape` is a tuple of integers, and dtype is a
-        numpy data type (float[32/64] or complex[64/128]). The function must
-        return the tuple (shape_out, dtype_out).
-    **kwargs : optional arguments that will be provided in `state` when calling
-        `func` and `func_abstract`.
-
-    Returns
-    -------
-    op : function of the operator for use in Jax computations
-
-    Notes
-    -----
-    - `func` must not return anything; the result of the computation must be
-      written into `out`.
-    - the contents of `inp` must not be modified.
-    - no references to `inp` or `out` may be stored beyond the execution
-      time of `func`
-    - `state` must not be modified
-    - when calling `func`* or `func_abstract`*, `state` may contain some entries
-      that the user did not supply in `**kwargs`; these will have names starting
-      with an underscore.
-    """
-    import copy
-
-    # somehow make sure that kwargs_clean only contains deep copies of
-    # everything in kwargs that are not accessible from anywhere else.
-    state = copy.deepcopy(kwargs)  # FIXME TODO
-    stateT = copy.copy(state)
-    global _global_opcounter
-    state["_opid"] = stateT["_opid"] = _global_opcounter
-    _global_opcounter += 1
-    state["_func"] = stateT["_func_T"] = func
-    state["_func_T"] = stateT["_func"] = func_T
-    state["_func_abstract"] = stateT["_func_abstract_T"] = func_abstract
-    state["_func_abstract_T"] = stateT["_func_abstract"] = func_abstract_T
-    return partial(_call, state=state, stateT=stateT)
diff --git a/jaxbind/__init__.py b/jaxbind/__init__.py
new file mode 100644
index 0000000..c65dec6
--- /dev/null
+++ b/jaxbind/__init__.py
@@ -0,0 +1,4 @@
+__version__ = "0.1.0"
+
+from .jaxbind import get_linear_call, get_nonlinear_call
+from .misc import load_kwargs
diff --git a/jax_linop/custom_map.py b/jaxbind/custom_map.py
similarity index 96%
rename from jax_linop/custom_map.py
rename to jaxbind/custom_map.py
index 2f6a7a4..82d75a1 100644
--- a/jax_linop/custom_map.py
+++ b/jaxbind/custom_map.py
@@ -1,4 +1,4 @@
-# Copyright(C) 2022-2023 Gordian Edenhofer
+# Copyright(C) 2022-2024 Gordian Edenhofer
 # SPDX-License-Identifier: GPL-2.0+ OR BSD-2-Clause
 
 from functools import partial
@@ -67,7 +67,7 @@ def _generic_smap(fun, in_axes, out_axes, unroll, *x, _scan=lax.scan, **k):
         fun=fun,
         unmapped=unmapped,
         unflatten=partial(tree_unflatten, x_td),
-        in_axes=in_axes
+        in_axes=in_axes,
     )
     _, y = _scan(fun_reord, None, mapped, unroll=unroll)
 
@@ -98,8 +98,7 @@ def _generic_smap(fun, in_axes, out_axes, unroll, *x, _scan=lax.scan, **k):
 # unnecessary recompiles. Ensure scan is compiled only once by compiling the
 # whole data dependence.
 _smap = jax.jit(
-    _generic_smap,
-    static_argnames=("fun", "in_axes", "out_axes", "unroll", "_scan")
+    _generic_smap, static_argnames=("fun", "in_axes", "out_axes", "unroll", "_scan")
 )
 
 
diff --git a/jaxbind/jaxbind.py b/jaxbind/jaxbind.py
new file mode 100644
index 0000000..d83a82a
--- /dev/null
+++ b/jaxbind/jaxbind.py
@@ -0,0 +1,433 @@
+# SPDX-License-Identifier: BSD-2-Clause
+# Authors: Martin Reinecke, Jakob Roth, Gordian Edenhofer
+
+# Copyright(C) 2023, 2024 Max-Planck-Society
+
+import pickle
+from collections import namedtuple
+from functools import partial
+from typing import Union
+
+import jax
+import jaxlib.mlir.dialects.stablehlo as hlo
+import jaxlib.mlir.ir as ir
+import numpy as np
+from jax.interpreters import ad, batching, mlir
+from jax.interpreters.mlir import ir_constant as irc
+from jaxlib.hlo_helpers import custom_call
+
+
+__all__ = ['get_linear_call', 'get_nonlinear_call']
+
+import _jaxbind
+
+for _name, _value in _jaxbind.registrations().items():
+    jax.lib.xla_client.register_custom_call_target(_name, _value, platform="cpu")
+
+
+# Hack to avoid classes and having to register a PyTree
+_shared_args_names = (
+    "abstract",
+    "abstract_T",
+    "first_n_args_fixed",
+    "can_batch",
+    "batch_axes",
+)
+LinearFunction = namedtuple("LinearFunction", ("f", "T") + _shared_args_names)
+MultiLinearFunction = namedtuple("MultiLinearFunction", ("f", "T") + _shared_args_names)
+NonLinearFunction = namedtuple(
+    "NonLinearFunction", ("f", "derivatives") + _shared_args_names
+)
+
+FunctionType = Union[LinearFunction, MultiLinearFunction, NonLinearFunction]
+
+
+def _exec_abstract(*args, _func: FunctionType, **kwargs):
+    if _func.can_batch:
+        assert "batch_axes" not in kwargs
+        kwargs["batch_axes"] = _func.batch_axes
+    ae = _func.abstract(*args, **kwargs)
+    # NOTE, do not attempt to unpack the batch axes b/c it might not be there
+    return tuple(jax.core.ShapedArray(*sdb[:2]) for sdb in ae)
+
+
+# the values are explained in src/ducc0/bindings/typecode.h
+_dtype_dict = {
+    np.dtype(np.float32): 3,
+    np.dtype(np.float64): 7,
+    np.dtype(np.uint8): 32,
+    np.dtype(np.uint64): 39,
+    np.dtype(np.complex64): 67,
+    np.dtype(np.complex128): 71,
+}
+
+
+def _lowering(ctx, *args, _func: FunctionType, _platform="cpu", **kwargs):
+    operands = [irc(id(_func.f))]  # Pass the ID of the callable to C++
+    operand_layouts = [()]
+
+    operands += [irc(len(args))]
+    operand_layouts += [()]
+    assert len(args) == len(ctx.avals_in)
+    # All `args` are assumed to be JAX arrays
+    for a, ca in zip(args, ctx.avals_in):
+        operands += (
+            [irc(_dtype_dict[ca.dtype]), irc(ca.ndim)]
+            + [irc(i) for i in ca.shape]
+            + [a]
+        )
+        lyt_a = tuple(range(ca.ndim - 1, -1, -1))
+        operand_layouts += [()] * (2 + ca.ndim) + [lyt_a]
+
+    operands += [irc(len(ctx.avals_out))]
+    operand_layouts += [()]
+    result_layouts = []
+    result_types = []
+    for co in ctx.avals_out:
+        operands += [irc(_dtype_dict[co.dtype]), irc(co.ndim)] + [
+            irc(i) for i in co.shape
+        ]
+        operand_layouts += [()] * (2 + co.ndim)
+        result_layouts += [tuple(range(co.ndim - 1, -1, -1))]
+        rs_typ = mlir.ir.RankedTensorType.get(co.shape, mlir.dtype_to_ir_type(co.dtype))
+        result_types += [rs_typ]
+
+    if _func.can_batch:
+        assert "batch_axes" not in kwargs
+        kwargs["batch_axes"] = _func.batch_axes
+    kwargs = np.frombuffer(pickle.dumps(kwargs), dtype=np.uint8)
+    kwargs_ir = hlo.constant(
+        ir.DenseElementsAttr.get(kwargs, type=ir.IntegerType.get_unsigned(8))
+    )
+    operands += [irc(_dtype_dict[kwargs.dtype]), irc(kwargs.size), kwargs_ir]
+    operand_layouts += [(), (), [0]]
+
+    assert len(operand_layouts) == len(operands)
+    if _platform == "cpu":
+        return custom_call(
+            _platform + "_pycall",
+            result_types=result_types,
+            result_layouts=result_layouts,
+            operands=operands,
+            operand_layouts=operand_layouts,
+        ).results
+    elif _platform == "gpu":
+        raise ValueError("No GPU support")
+    raise ValueError("Unsupported platform; this must be either 'cpu' or 'gpu'")
+
+
+def _explicify_zeros(x):
+    if isinstance(x, (tuple, list)):
+        return [ad.instantiate_zeros(t) if isinstance(t, ad.Zero) else t for t in x]
+    return ad.instantiate_zeros(x) if isinstance(x, ad.Zero) else x
+
+
+def _jvp(args, tangents, *, _func: FunctionType, **kwargs):
+    res = _prim.bind(*args, **kwargs, _func=_func)
+
+    def zero_tans(tans):
+        if isinstance(tans, (tuple, list)):
+            return [isinstance(t, ad.Zero) for t in tans]
+        return [isinstance(tans, ad.Zero)]
+
+    n_args = len(args)
+    n_f_args = _func.first_n_args_fixed
+    assert n_args > n_f_args
+    args_fixed = n_f_args * (True,) + (n_args - n_f_args) * (False,)
+    assert len(args) == len(tangents) == len(args_fixed)
+
+    tan_is_zero = zero_tans(tangents)
+    assert len(args) == len(tan_is_zero)
+
+    for i, (a, t) in enumerate(zip(args_fixed, tan_is_zero)):
+        if a and not t:
+            raise RuntimeError(f"{i}th positional argument not differentiable")
+
+    if all(type(t) is ad.Zero for t in tangents):
+        tans = list((jax.lax.zeros_like_array(a) for a in res))
+        return (res, tans)
+
+    tans = None
+    if isinstance(_func, MultiLinearFunction):
+        for i, t in enumerate(tangents):
+            if not args_fixed[i]:
+                t = _explicify_zeros(t)
+                tn = _prim.bind(*args[:i], t, *args[i + 1 :], **kwargs, _func=_func)
+                tans = (
+                    tn if tans is None else tuple(t + tn_i for t, tn_i in zip(tans, tn))
+                )
+    elif isinstance(_func, LinearFunction):
+        inp = []
+        for a, f, t in zip(args, args_fixed, tangents):
+            inp.append(a if f else _explicify_zeros(t))
+        tans = _prim.bind(*inp, **kwargs, _func=_func)
+    elif isinstance(_func, NonLinearFunction):
+        f, f_T = _func.derivatives
+        tan_in = [t for f, t in zip(args_fixed, tangents) if not f]
+        tan_in = _explicify_zeros(tan_in)
+        _func = LinearFunction(
+            f=f,
+            T=f_T,
+            abstract=_func.abstract,
+            abstract_T=_func.abstract_T,
+            first_n_args_fixed=len(args),
+            can_batch=_func.can_batch,
+            batch_axes=_func.batch_axes,
+        )
+        tans = _prim.bind(*args, *tan_in, **kwargs, _func=_func)
+    else:
+        raise TypeError(f"JVP for {type(_func)} not implemented")
+
+    assert tans is not None
+    return (res, tans)
+
+
+# NOTE: for whatever reason JAX will pass each arg separately to _transpose
+# and not as a tuple as for _jvp. Thus we need *args since we don't know
+# the number of arguments.
+def _transpose(cotangents, *args, _func: FunctionType, **kwargs):
+    assert isinstance(_func, (LinearFunction, MultiLinearFunction))
+    if _func.T is None:
+        raise NotImplementedError(f"transpose of {_func} not implemented")
+    n_args = len(args)
+    n_f_args = _func.first_n_args_fixed
+    assert n_args > n_f_args
+    args_fixed = n_f_args * (True,) + (n_args - n_f_args) * (False,)
+    arg_is_lin = [ad.is_undefined_primal(a) for a in args]
+    assert len(args_fixed) >= len(arg_is_lin)
+
+    for i, (a, is_lin) in enumerate(zip(args_fixed, arg_is_lin)):
+        if a and is_lin:
+            raise RuntimeError(
+                f"Cannot transpose with respect to positional argument number {i}"
+            )
+
+    if isinstance(_func, MultiLinearFunction):
+        assert sum(arg_is_lin) == 1
+        lin_arg = arg_is_lin.index(True)
+        c_in = cotangents
+        a_in = args[:lin_arg] + args[lin_arg + 1 :]
+
+        assert isinstance(_func.T, tuple)
+        assert isinstance(_func.abstract_T, tuple)
+        # TODO(edh): I think this should be a LinearFunction
+        _func = _func._replace(
+            f=_func.T[lin_arg],
+            T=None,
+            abstract=_func.abstract_T[lin_arg],
+            abstract_T=None,
+        )
+        res = _prim.bind(*a_in, *c_in, **kwargs, _func=_func)
+        res = [None] * lin_arg + res + [None] * (len(arg_is_lin) - (lin_arg + 1))
+    elif isinstance(_func, LinearFunction):
+        inp = []
+        for a, f in zip(args, args_fixed):
+            if f:
+                assert not ad.is_undefined_primal(a)
+                inp.append(a)
+        cot = _explicify_zeros(cotangents)
+        _func = _func._replace(
+            f=_func.T,
+            T=_func.f,
+            abstract=_func.abstract_T,
+            abstract_T=_func.abstract,
+            first_n_args_fixed=len(inp),
+        )
+        res = _prim.bind(*inp, *cot, **kwargs, _func=_func)
+        res = n_f_args * [None] + res
+    else:
+        raise TypeError(f"transpose for {type(_func)} not implemented")
+    return res
+
+
+def _batch(args, in_axes, *, _func: FunctionType, **kwargs):
+    from .custom_map import smap
+
+    if not _func.can_batch:
+        y = smap(partial(_prim.bind, _func=_func), in_axes=in_axes)(*args)
+        out_axes = [0] * len(y)
+    else:
+        batch_axes = _func.batch_axes
+        batch_axes = ((),) * len(in_axes) if batch_axes is None else batch_axes
+        new_batch_axes = []
+        assert len(in_axes) == len(batch_axes)
+        for ia, baxes in zip(in_axes, batch_axes):
+            baxes_new = []
+            if ia is not None:
+                assert isinstance(ia, int)
+                for b in baxes:
+                    if b >= ia:
+                        b += 1
+                    baxes_new.append(b)
+                baxes = tuple(tuple(baxes_new) + (ia,))
+            new_batch_axes.append(baxes)
+        new_batch_axes = tuple(new_batch_axes)
+        _func = _func._replace(batch_axes=new_batch_axes)
+
+        args_w = [jax.ShapeDtypeStruct(el.shape, el.dtype) for el in args]
+        out_w = _func.abstract(*args_w, batch_axes=new_batch_axes, **kwargs)
+        out_axes = [ba_wb for _, _, ba_wb in out_w]
+        y = _call(*args, _func=_func, **kwargs)
+    return y, out_axes
+
+
+_prim = jax.core.Primitive("jaxbind_prim")
+_prim.multiple_results = True
+_prim.def_impl(partial(jax.interpreters.xla.apply_primitive, _prim))
+_prim.def_abstract_eval(_exec_abstract)
+
+for platform in ["cpu", "gpu"]:
+    mlir.register_lowering(
+        _prim, partial(_lowering, _platform=platform), platform=platform
+    )
+    ad.primitive_jvps[_prim] = _jvp
+    ad.primitive_transposes[_prim] = _transpose
+    batching.primitive_batchers[_prim] = _batch
+
+
+def _call(*args, _func: FunctionType, **kwargs):
+    return _prim.bind(*args, **kwargs, _func=_func)
+
+
+def get_linear_call(
+    f,
+    f_T,
+    /,
+    abstract,
+    abstract_T,
+    *,
+    first_n_args_fixed=0,
+    func_can_batch=False,
+) -> partial:
+    """Create a JAX primitive for the provided linear function
+
+    Parameters
+    ----------
+    f, f_T : linear function respectively its transpose
+        The function signature must be (out, args, kwargs_dump), where
+        out and args are tuples. The results of the functions should be written as
+        numpy.ndarrays of float[32/64] or complex[64/128] type into the out tuple.
+        The args tuple contains the input for the function. In kwargs_dump,
+        potential keyword arguments are contained in serialized form. The
+        keyword arguments can be deserialized via
+        `jaxbind.load_kwargs(kwargs_dump)`.
+    abstract, abstract_T : function respectively its transpose
+        Computing the shape and dtype of the operator's output from shape and
+        dtype of its input. Its signature must be `(*args, **kwargs)`. `args`
+        will be a tuple containing abstract tracer arrays with shape and dtype
+        for each input argument of `f` respectively `f_T`. Via `**kwargs`,
+        potential keyword arguments are passed to the function. The function
+        must return a tuple containing tuples of (shape_out, dtype_out) for each
+        output argument of `f` respectively `f_T`.
+    first_n_args_fixed : int
+        If the function cannot be differentiated with respect to some of the
+        arguments, these can be passed as the first arguments to the function.
+        fist_n_args_fixed indicates the number of non-differential arguments.
+        Note: The function does not need to be linear with respect to these
+        arguments. Default 0 (all arguments are differentiable).
+    func_can_batch : bool
+        Indicator whether the function natively supports batching. If true, the
+        function will receive one additional argument called `batch_axes`. The
+        parameter will be a tuple of tuples, or None if no batching is currently
+        performed. The tuple will be of length of the input and for each input
+        will contain a tuple of integer indices along which the computation
+        shall be batched.
+
+    Returns
+    -------
+    op : Jax primitive corresponding to the function f.
+
+    Notes
+    -----
+    - `f` and 'f_T' must not return anything; the result of the computation must be
+      written into the member arrays of `out`.
+    - the contents of `args` must not be modified.
+    - no reference to the contents of `args` or `out` may be stored beyond
+      the execution time of `f` or `f_T`.
+    """
+    kw = dict(
+        f=f,
+        T=f_T,
+        abstract=abstract,
+        abstract_T=abstract_T,
+        first_n_args_fixed=first_n_args_fixed,
+        batch_axes=None,
+        can_batch=func_can_batch,
+    )
+    if isinstance(f_T, (tuple, list)):
+        _func = MultiLinearFunction(**kw)
+    else:
+        _func = LinearFunction(**kw)
+    return partial(_call, _func=_func)
+
+
+def get_nonlinear_call(
+    f,
+    f_derivative,
+    /,
+    abstract,
+    abstract_reverse,
+    *,
+    first_n_args_fixed=0,
+    func_can_batch=False,
+) -> partial:
+    """Create a JAX primitive for the provided (nonlinear) function
+
+    Parameters
+    ----------
+    f : function
+        The function signature must be (out, args, kwargs_dump), where
+        out and args are tuples. The results of the functions should be written as
+        numpy.ndarrays of float[32/64] or complex[64/128] type into the out tuple.
+        The args tuple contains the input for the function. In kwargs_dump,
+        potential keyword arguments are contained in serialized form. The
+        keyword arguments can be deserialized via
+        `jaxbind.load_kwargs(kwargs_dump)`.
+    f_derivative: tuple of functions
+        Tuple containing functions for evaluating jvp and vjp of f. The fist entry
+        in the function should evaluate jvp, the second vjp. The signature of the
+        jvp and vjp functions should be (out, args, kwargs_dump) analogous to f.
+    abstract, abstract_reverse : functions
+        Computing the shape and dtype of the operator's output from shape and
+        dtype of its input. Its signature must be (*args, **kwargs). *args will
+        be a tuple containing abstract tracer arrays with shape and dtype for
+        each input argument of f. Via **kwargs, potential keyword arguments are
+        passed to the function. The function must return the tuple containing a
+        tuples of (shape_out, dtype_out). abstract should compute the output
+        shapes of f and jvp. abstract_reverse should compute the output shape of
+        vjp.
+    first_n_args_fixed : int
+        If the function cannot be differentiated with respect to some of the
+        arguments, these can be passed as the first arguments to the function.
+        fist_n_args_fixed indicates the number of non-differential arguments.
+        Default 0 (all arguments are differentiable).
+    func_can_batch : bool
+        Indicator whether the function natively supports batching. If true, the
+        function will receive one additional argument called `batch_axes`. The
+        parameter will be a tuple of tuples, or None if no batching is currently
+        performed. The tuple will be of length of the input and for each input
+        will contain a tuple of integer indices along which the computation
+        shall be batched.
+
+    Returns
+    -------
+    op : Jax primitive corresponding to the function f.
+
+    Notes
+    -----
+    - `f` and members of 'f_derivative' must not return anything; the result
+      of the computation must be written into the member arrays `out`.
+    - the contents of `args` must not be modified.
+    - no references to the contents of `args` or `out` may be stored beyond
+      the execution time of `f` or `f_derivative`.
+    """
+    _func = NonLinearFunction(
+        f=f,
+        abstract=abstract,
+        abstract_T=abstract_reverse,
+        first_n_args_fixed=first_n_args_fixed,
+        batch_axes=None,
+        can_batch=func_can_batch,
+        derivatives=f_derivative,
+    )
+    return partial(_call, _func=_func)
diff --git a/jaxbind/misc.py b/jaxbind/misc.py
new file mode 100644
index 0000000..7bee06b
--- /dev/null
+++ b/jaxbind/misc.py
@@ -0,0 +1,11 @@
+# SPDX-License-Identifier: BSD-2-Clause
+# Authors: Martin Reinecke, Jakob Roth, Gordian Edenhofer
+
+# Copyright(C) 2023, 2024 Max-Planck-Society
+
+
+def load_kwargs(kwargs_dump, /):
+    import pickle
+    import numpy as np
+
+    return pickle.loads(np.ndarray.tobytes(kwargs_dump))
diff --git a/paper/paper.bib b/paper/paper.bib
new file mode 100644
index 0000000..885fc89
--- /dev/null
+++ b/paper/paper.bib
@@ -0,0 +1,76 @@
+@software{Jax2018,
+  author  = {James Bradbury and Roy Frostig and Peter Hawkins and Matthew James Johnson and Chris Leary and Dougal Maclaurin and George Necula and Adam Paszke and Jake Vander{P}las and Skye Wanderman-{M}ilne and Qiao Zhang},
+  title   = {{JAX}: composable transformations of {P}ython+{N}um{P}y programs},
+  url     = {http://github.com/google/jax},
+  version = {0.3.13},
+  year    = {2018}
+}
+@software{ducc0,
+  author  = {Martin Reinecke},
+  title   = {{DUCC}: Distinctly Useful Code Collection},
+  url     = {https://gitlab.mpcdf.mpg.de/mtr/ducc},
+  version = {0.33.0},
+  year    = {2024}
+}
+@software{Moses2024,
+  author  = {EnzymeAD},
+  title   = {{Enzyme-JAX}},
+  url     = {https://github.com/EnzymeAD/Enzyme-JAX},
+  version = {0.0.6},
+  year    = {2024}
+}
+@inproceedings{Moses2020,
+  author    = {Moses, William and Churavy, Valentin},
+  booktitle = {Advances in Neural Information Processing Systems},
+  editor    = {H. Larochelle and M. Ranzato and R. Hadsell and M. F. Balcan and H. Lin},
+  pages     = {12472--12485},
+  publisher = {Curran Associates, Inc.},
+  title     = {Instead of Rewriting Foreign Code for Machine Learning, Automatically Synthesize Fast Gradients},
+  url       = {https://proceedings.neurips.cc/paper/2020/file/9332c513ef44b682e9347822c2e457ac-Paper.pdf},
+  volume    = {33},
+  year      = {2020}
+}
+@inproceedings{Moses2021,
+  author    = {Moses, William S. and Churavy, Valentin and Paehler, Ludger and H\"{u}ckelheim, Jan and Narayanan, Sri Hari Krishna and Schanen, Michel and Doerfert, Johannes},
+  title     = {Reverse-Mode Automatic Differentiation and Optimization of GPU Kernels via Enzyme},
+  year      = {2021},
+  isbn      = {9781450384421},
+  publisher = {Association for Computing Machinery},
+  address   = {New York, NY, USA},
+  url       = {https://doi.org/10.1145/3458817.3476165},
+  doi       = {10.1145/3458817.3476165},
+  booktitle = {Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis},
+  articleno = {61},
+  numpages  = {16},
+  keywords  = {CUDA, LLVM, ROCm, HPC, AD, GPU, automatic differentiation},
+  location  = {St. Louis, Missouri},
+  series    = {SC '21}
+}
+@inproceedings{Moses2022,
+  author    = {Moses, William S. and Narayanan, Sri Hari Krishna and Paehler, Ludger and Churavy, Valentin and Schanen, Michel and H\"{u}ckelheim, Jan and Doerfert, Johannes and Hovland, Paul},
+  title     = {Scalable Automatic Differentiation of Multiple Parallel Paradigms through Compiler Augmentation},
+  year      = {2022},
+  isbn      = {9784665454445},
+  publisher = {IEEE Press},
+  booktitle = {Proceedings of the International Conference on High Performance Computing, Networking, Storage and Analysis},
+  articleno = {60},
+  numpages  = {18},
+  keywords  = {automatic differentiation, tasks, OpenMP, compiler, Julia, parallel, Enzyme, C++, RAJA, hybrid parallelization, MPI, distributed, LLVM},
+  location  = {Dallas, Texas},
+  series    = {SC '22}
+}
+@software{Resolve2024,
+  author = {Arras, Philipp and Roth, Jakob and Ding, Simon and Reinecke, Martin and Fuchs, Richard and Johnson, Vishal},
+  title  = {RESOLVE},
+  url    = {https://gitlab.mpcdf.mpg.de/ift/resolve},
+  year   = {2024}
+}
+@misc{Edenhofer2023NIFTyRE,
+  author        = {{Edenhofer}, Gordian and {Frank}, Philipp and {Roth}, Jakob and {Leike}, Reimar H. and {Guerdi}, Massin and {Scheel-Platz}, Lukas I. and {Guardiani}, Matteo and {Eberle}, Vincent and {Westerkamp}, Margret and {Enßlin}, Torsten A.},
+  title         = {{Re-Envisioning Numerical Information Field Theory (NIFTy.re): A Library for Gaussian Processes and Variational Inference}},
+  keywords      = {Software},
+  year          = 2024,
+  eprint        = {2402.16683},
+  archiveprefix = {arXiv},
+  primaryclass  = {astro-ph.IM}
+}
diff --git a/paper/paper.md b/paper/paper.md
new file mode 100644
index 0000000..d444177
--- /dev/null
+++ b/paper/paper.md
@@ -0,0 +1,206 @@
+---
+title: '`JAXbind`: Easy bindings to JAX'
+tags:
+  - Python
+  - Machine Learning
+  - High Performance Computing
+authors:
+  - name: Jakob Roth
+    orcid: 0000-0002-8873-8215
+    equal-contrib: true
+    affiliation: "1, 2, 3"
+  - name: Martin Reinecke
+    orcid:
+    equal-contrib: true
+    affiliation: "1"
+  - name: Gordian Edenhofer
+    equal-contrib: true
+    orcid: 0000-0003-3122-4894
+    affiliation: "1, 2"  # Multiple affiliations must be quoted
+affiliations:
+  - name: Max Planck Institute for Astrophysics, Karl-Schwarzschild-Str. 1, 85748 Garching, Germany
+    index: 1
+  - name: Ludwig Maximilian University of Munich, Geschwister-Scholl-Platz 1, 80539 Munich, Germany
+    index: 2
+  - name: Technical University of Munich, Boltzmannstr. 3, 85748 Garching, Germany
+    index: 3
+date: 22 February 2024
+bibliography: paper.bib
+---
+
+# Summary
+
+JAX is widely used in machine learning and scientific computing.
+Scientific computing relies on existing high-performance code which we would ideally like to use in JAX.
+Reimplementing the existing code in JAX is often impractical and the existing interface in JAX for connecting 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 supporting any JAX transformations.
+
+
+# Statement of Need
+
+The use of JAX [@Jax2018] is widespread in the natural sciences.
+JAX's powerful transformation system is of especially high interest.
+It enables retrieving arbitrary derivatives of functions, batch computations, and just-in-time code compilation for additional performance.
+Its transformation system relies on all constituents of the computation being written in JAX.
+
+A plethora of high-performance code is not written in JAX and thus not accessible from within JAX.
+Rewriting these is often infeasible and/or inefficient.
+Ideally, we would like to intermix existing high-performance code with JAX code.
+However, connecting code to JAX requires knowledge of the internals of JAX and its C++ backend.
+
+In this paper, we present `JAXbind`, a package for bridging any function to JAX without in-depth knowledge of JAX's transformation system.
+The interface is accessible from Python without requiring any development in C++.
+The package is able to register any function, its partial derivatives and their transpose functions as a JAX native call, a so-called primitive.
+Derivatives, compilation rules, and batching rules are automatically registered with JAX.
+
+<!-- Mention (if applicable) a representative set of past or ongoing research projects using the software and recent scholarly publications enabled by it. -->
+We believe `JAXbind` to be highly useful in scientific computing.
+<!-- There are a lot of well-developed packages in JAX for, e.g., optimization and sampling that could be used once existing code is able to interface with JAX. -->
+We intend to use this package to connect the Hartley transform and the spherical harmonic transform from ducc [@ducc0] to NIFTy [@Edenhofer2023NIFTyRE] as well as the radio interferometry response from ducc with \texttt{resolve} [@Resolve2024] for radio astronomy.
+Furthermore, we intend to connect the non-uniform FFT from ducc with JAX for applications in strong-lensing astrophysics.
+We envision many further applications within and outside of astrophysics.
+
+<!-- A list of key references, including to other software addressing related needs. Note that the references should include full names of venues, e.g., journals and conferences, not abbreviations only understood in the context of a specific discipline. -->
+To the best of our knowledge no other code currently exists for connecting generic functions to JAX.
+The package that comes the closest is Enzyme-JAX [@Moses2024].
+Enzyme-JAX allows one to differentiate a C++ function with Enzyme [@Moses2020; @Moses2021; @Moses2022] and connect it together with its derivative to JAX.
+However, it enforces the use of Enzyme for deriving derivatives and does not allow for connecting arbitrary code to JAX.
+
+# 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.
+<!-- Should we point out specifically that the abstract functions take "traditional" args and kwargs? -->
+
+```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.
+
+# 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. We especially want to highlight that the interfacing with the JAX automatic differentiation engine is identical for CPU and GPU.
+
+# Acknowledgements
+
+We would like to thank Dan Foreman-Mackey for his detailed guide (https://dfm.io/posts/extending-jax/) on connecting C++ code to JAX.
+Jakob Roth acknowledges financial support from the German Federal Ministry of Education and Research (BMBF) under grant 05A23WO1 (Verbundprojekt D-MeerKAT III).
+Gordian Edenhofer acknowledges support from the German Academic Scholarship Foundation in the form of a PhD scholarship ("Promotionsstipendium der Studienstiftung des Deutschen Volkes").
+
+# References
+
+<!-- Citations to entries in paper.bib should be in
+[rMarkdown](http://rmarkdown.rstudio.com/authoring_bibliographies_and_citations.html)
+format.
+
+For a quick reference, the following citation commands can be used:
+- `@author:2001`  ->  "Author et al. (2001)"
+- `[@author:2001]` -> "(Author et al., 2001)"
+- `[@author1:2001; @author2:2001]` -> "(Author1 et al., 2001; Author2 et al., 2002)"
+# Figures
+
+Figures can be included like this:
+![Caption for example figure.\label{fig:example}](figure.png)
+and referenced from text using \autoref{fig:example}.
+
+Figure sizes can be customized by adding an optional second parameter:
+![Caption for example figure.](figure.png){ width=20% }
+-->
diff --git a/paper/paper.tex b/paper/paper.tex
new file mode 100644
index 0000000..ffbbd8c
--- /dev/null
+++ b/paper/paper.tex
@@ -0,0 +1,626 @@
+\documentclass[10pt,a4paper,onecolumn]{article}
+\usepackage{marginnote}
+\usepackage{graphicx}
+\usepackage{xcolor}
+\usepackage{authblk,etoolbox}
+\usepackage{titlesec}
+\usepackage{calc}
+\usepackage{tikz}
+\usepackage{hyperref}
+\hypersetup{colorlinks,breaklinks=true,
+            urlcolor=[rgb]{0.0, 0.5, 1.0},
+            linkcolor=[rgb]{0.0, 0.5, 1.0}}
+\usepackage{caption}
+\usepackage{tcolorbox}
+\usepackage{amssymb,amsmath}
+\usepackage{ifxetex,ifluatex}
+\usepackage{seqsplit}
+\usepackage{xstring}
+
+\usepackage{float}
+\let\origfigure\figure
+\let\endorigfigure\endfigure
+\renewenvironment{figure}[1][2] {
+    \expandafter\origfigure\expandafter[H]
+} {
+    \endorigfigure
+}
+
+
+\usepackage{fixltx2e} % provides \textsubscript
+\usepackage[
+  backend=biber,
+%  style=alphabetic,
+%  citestyle=numeric
+]{biblatex}
+\bibliography{paper.bib}
+
+% --- Splitting \texttt --------------------------------------------------
+
+\let\textttOrig=\texttt
+\def\texttt#1{\expandafter\textttOrig{\seqsplit{#1}}}
+\renewcommand{\seqinsert}{\ifmmode
+  \allowbreak
+  \else\penalty6000\hspace{0pt plus 0.02em}\fi}
+
+
+% --- Pandoc does not distinguish between links like [foo](bar) and
+% --- [foo](foo) -- a simplistic Markdown model.  However, this is
+% --- wrong:  in links like [foo](foo) the text is the url, and must
+% --- be split correspondingly.
+% --- Here we detect links \href{foo}{foo}, and also links starting
+% --- with https://doi.org, and use path-like splitting (but not
+% --- escaping!) with these links.
+% --- Another vile thing pandoc does is the different escaping of
+% --- foo and bar.  This may confound our detection.
+% --- This problem we do not try to solve at present, with the exception
+% --- of doi-like urls, which we detect correctly.
+
+
+\makeatletter
+\let\href@Orig=\href
+\def\href@Urllike#1#2{\href@Orig{#1}{\begingroup
+    \def\Url@String{#2}\Url@FormatString
+    \endgroup}}
+\def\href@Notdoi#1#2{\def\tempa{#1}\def\tempb{#2}%
+  \ifx\tempa\tempb\relax\href@Urllike{#1}{#2}\else
+  \href@Orig{#1}{#2}\fi}
+\def\href#1#2{%
+  \IfBeginWith{#1}{https://doi.org}%
+  {\href@Urllike{#1}{#2}}{\href@Notdoi{#1}{#2}}}
+\makeatother
+
+\newlength{\cslhangindent}
+\setlength{\cslhangindent}{1.5em}
+\newlength{\csllabelwidth}
+\setlength{\csllabelwidth}{3em}
+\newenvironment{CSLReferences}[3] % #1 hanging-ident, #2 entry spacing
+ {% don't indent paragraphs
+  \setlength{\parindent}{0pt}
+  % turn on hanging indent if param 1 is 1
+  \ifodd #1 \everypar{\setlength{\hangindent}{\cslhangindent}}\ignorespaces\fi
+  % set entry spacing
+  \ifnum #2 > 0
+  \setlength{\parskip}{#2\baselineskip}
+  \fi
+ }%
+ {}
+\usepackage{calc}
+\newcommand{\CSLBlock}[1]{#1\hfill\break}
+\newcommand{\CSLLeftMargin}[1]{\parbox[t]{\csllabelwidth}{#1}}
+\newcommand{\CSLRightInline}[1]{\parbox[t]{\linewidth - \csllabelwidth}{#1}}
+\newcommand{\CSLIndent}[1]{\hspace{\cslhangindent}#1}
+
+% --- Page layout -------------------------------------------------------------
+\usepackage[top=3.5cm, bottom=3cm, right=2.54cm, left=2.54cm,
+            headheight=2.2cm]{geometry}
+
+% --- Default font ------------------------------------------------------------
+\renewcommand\familydefault{\sfdefault}
+
+% --- Style -------------------------------------------------------------------
+\renewcommand{\bibfont}{\small \sffamily}
+\renewcommand{\captionfont}{\small\sffamily}
+\renewcommand{\captionlabelfont}{\bfseries}
+
+% --- Section/SubSection/SubSubSection ----------------------------------------
+\titleformat{\section}
+  {\normalfont\sffamily\Large\bfseries}
+  {}{0pt}{}
+\titleformat{\subsection}
+  {\normalfont\sffamily\large\bfseries}
+  {}{0pt}{}
+\titleformat{\subsubsection}
+  {\normalfont\sffamily\bfseries}
+  {}{0pt}{}
+\titleformat*{\paragraph}
+  {\sffamily\normalsize}
+
+
+% --- Header / Footer ---------------------------------------------------------
+\usepackage{fancyhdr}
+\pagestyle{fancy}
+\fancyhf{}
+%\renewcommand{\headrulewidth}{0.50pt}
+\renewcommand{\headrulewidth}{0pt}
+\fancyhead[L]{}
+\fancyhead[C]{}
+\fancyhead[R]{}
+\renewcommand{\footrulewidth}{0.0pt}
+
+\fancyfoot[L]{}
+
+
+\fancyfoot[R]{\sffamily \thepage}
+\makeatletter
+\let\ps@plain\ps@fancy
+\fancyheadoffset[L]{4.5cm}
+\fancyfootoffset[L]{4.5cm}
+
+% --- Macros ---------
+
+\definecolor{linky}{rgb}{0.0, 0.5, 1.0}
+
+\newtcolorbox{repobox}
+   {colback=red, colframe=red!75!black,
+     boxrule=0.5pt, arc=2pt, left=6pt, right=6pt, top=3pt, bottom=3pt}
+
+\newcommand{\ExternalLink}{%
+   \tikz[x=1.2ex, y=1.2ex, baseline=-0.05ex]{%
+       \begin{scope}[x=1ex, y=1ex]
+           \clip (-0.1,-0.1)
+               --++ (-0, 1.2)
+               --++ (0.6, 0)
+               --++ (0, -0.6)
+               --++ (0.6, 0)
+               --++ (0, -1);
+           \path[draw,
+               line width = 0.5,
+               rounded corners=0.5]
+               (0,0) rectangle (1,1);
+       \end{scope}
+       \path[draw, line width = 0.5] (0.5, 0.5)
+           -- (1, 1);
+       \path[draw, line width = 0.5] (0.6, 1)
+           -- (1, 1) -- (1, 0.6);
+       }
+   }
+
+% --- Title / Authors ---------------------------------------------------------
+% patch \maketitle so that it doesn't center
+\patchcmd{\@maketitle}{center}{flushleft}{}{}
+\patchcmd{\@maketitle}{center}{flushleft}{}{}
+% patch \maketitle so that the font size for the title is normal
+\patchcmd{\@maketitle}{\LARGE}{\LARGE\sffamily}{}{}
+% patch the patch by authblk so that the author block is flush left
+\def\maketitle{{%
+  \renewenvironment{tabular}[2][]
+    {\begin{flushleft}}
+    {\end{flushleft}}
+  \AB@maketitle}}
+\makeatletter
+\renewcommand\AB@affilsepx{ \protect\Affilfont}
+%\renewcommand\AB@affilnote[1]{{\bfseries #1}\hspace{2pt}}
+\renewcommand\AB@affilnote[1]{{\bfseries #1}\hspace{3pt}}
+\renewcommand{\affil}[2][]%
+   {\newaffiltrue\let\AB@blk@and\AB@pand
+      \if\relax#1\relax\def\AB@note{\AB@thenote}\else\def\AB@note{#1}%
+        \setcounter{Maxaffil}{0}\fi
+        \begingroup
+        \let\href=\href@Orig
+        \let\texttt=\textttOrig
+        \let\protect\@unexpandable@protect
+        \def\thanks{\protect\thanks}\def\footnote{\protect\footnote}%
+        \@temptokena=\expandafter{\AB@authors}%
+        {\def\\{\protect\\\protect\Affilfont}\xdef\AB@temp{#2}}%
+         \xdef\AB@authors{\the\@temptokena\AB@las\AB@au@str
+         \protect\\[\affilsep]\protect\Affilfont\AB@temp}%
+         \gdef\AB@las{}\gdef\AB@au@str{}%
+        {\def\\{, \ignorespaces}\xdef\AB@temp{#2}}%
+        \@temptokena=\expandafter{\AB@affillist}%
+        \xdef\AB@affillist{\the\@temptokena \AB@affilsep
+          \AB@affilnote{\AB@note}\protect\Affilfont\AB@temp}%
+      \endgroup
+       \let\AB@affilsep\AB@affilsepx
+}
+\makeatother
+\renewcommand\Authfont{\sffamily\bfseries}
+\renewcommand\Affilfont{\sffamily\small\mdseries}
+\setlength{\affilsep}{1em}
+
+
+\ifnum 0\ifxetex 1\fi\ifluatex 1\fi=0 % if pdftex
+  \usepackage[T1]{fontenc}
+  \usepackage[utf8]{inputenc}
+
+\else % if luatex or xelatex
+  \ifxetex
+    \usepackage{mathspec}
+    \usepackage{fontspec}
+
+  \else
+    \usepackage{fontspec}
+  \fi
+  \defaultfontfeatures{Ligatures=TeX,Scale=MatchLowercase}
+
+\fi
+% use upquote if available, for straight quotes in verbatim environments
+\IfFileExists{upquote.sty}{\usepackage{upquote}}{}
+% use microtype if available
+\IfFileExists{microtype.sty}{%
+\usepackage{microtype}
+\UseMicrotypeSet[protrusion]{basicmath} % disable protrusion for tt fonts
+}{}
+
+\usepackage{hyperref}
+\hypersetup{unicode=true,
+            pdftitle={JAXbind: Easy bindings to JAX},
+            pdfborder={0 0 0},
+            breaklinks=true}
+\urlstyle{same}  % don't use monospace font for urls
+\usepackage{color}
+\usepackage{fancyvrb}
+\newcommand{\VerbBar}{|}
+\newcommand{\VERB}{\Verb[commandchars=\\\{\}]}
+\DefineVerbatimEnvironment{Highlighting}{Verbatim}{commandchars=\\\{\}}
+% Add ',fontsize=\small' for more characters per line
+\newenvironment{Shaded}{}{}
+\newcommand{\AlertTok}[1]{\textcolor[rgb]{1.00,0.00,0.00}{\textbf{#1}}}
+\newcommand{\AnnotationTok}[1]{\textcolor[rgb]{0.38,0.63,0.69}{\textbf{\textit{#1}}}}
+\newcommand{\AttributeTok}[1]{\textcolor[rgb]{0.49,0.56,0.16}{#1}}
+\newcommand{\BaseNTok}[1]{\textcolor[rgb]{0.25,0.63,0.44}{#1}}
+\newcommand{\BuiltInTok}[1]{#1}
+\newcommand{\CharTok}[1]{\textcolor[rgb]{0.25,0.44,0.63}{#1}}
+\newcommand{\CommentTok}[1]{\textcolor[rgb]{0.38,0.63,0.69}{\textit{#1}}}
+\newcommand{\CommentVarTok}[1]{\textcolor[rgb]{0.38,0.63,0.69}{\textbf{\textit{#1}}}}
+\newcommand{\ConstantTok}[1]{\textcolor[rgb]{0.53,0.00,0.00}{#1}}
+\newcommand{\ControlFlowTok}[1]{\textcolor[rgb]{0.00,0.44,0.13}{\textbf{#1}}}
+\newcommand{\DataTypeTok}[1]{\textcolor[rgb]{0.56,0.13,0.00}{#1}}
+\newcommand{\DecValTok}[1]{\textcolor[rgb]{0.25,0.63,0.44}{#1}}
+\newcommand{\DocumentationTok}[1]{\textcolor[rgb]{0.73,0.13,0.13}{\textit{#1}}}
+\newcommand{\ErrorTok}[1]{\textcolor[rgb]{1.00,0.00,0.00}{\textbf{#1}}}
+\newcommand{\ExtensionTok}[1]{#1}
+\newcommand{\FloatTok}[1]{\textcolor[rgb]{0.25,0.63,0.44}{#1}}
+\newcommand{\FunctionTok}[1]{\textcolor[rgb]{0.02,0.16,0.49}{#1}}
+\newcommand{\ImportTok}[1]{#1}
+\newcommand{\InformationTok}[1]{\textcolor[rgb]{0.38,0.63,0.69}{\textbf{\textit{#1}}}}
+\newcommand{\KeywordTok}[1]{\textcolor[rgb]{0.00,0.44,0.13}{\textbf{#1}}}
+\newcommand{\NormalTok}[1]{#1}
+\newcommand{\OperatorTok}[1]{\textcolor[rgb]{0.40,0.40,0.40}{#1}}
+\newcommand{\OtherTok}[1]{\textcolor[rgb]{0.00,0.44,0.13}{#1}}
+\newcommand{\PreprocessorTok}[1]{\textcolor[rgb]{0.74,0.48,0.00}{#1}}
+\newcommand{\RegionMarkerTok}[1]{#1}
+\newcommand{\SpecialCharTok}[1]{\textcolor[rgb]{0.25,0.44,0.63}{#1}}
+\newcommand{\SpecialStringTok}[1]{\textcolor[rgb]{0.73,0.40,0.53}{#1}}
+\newcommand{\StringTok}[1]{\textcolor[rgb]{0.25,0.44,0.63}{#1}}
+\newcommand{\VariableTok}[1]{\textcolor[rgb]{0.10,0.09,0.49}{#1}}
+\newcommand{\VerbatimStringTok}[1]{\textcolor[rgb]{0.25,0.44,0.63}{#1}}
+\newcommand{\WarningTok}[1]{\textcolor[rgb]{0.38,0.63,0.69}{\textbf{\textit{#1}}}}
+
+% --- We redefined \texttt, but in sections and captions we want the
+% --- old definition
+\let\addcontentslineOrig=\addcontentsline
+\def\addcontentsline#1#2#3{\bgroup
+  \let\texttt=\textttOrig\addcontentslineOrig{#1}{#2}{#3}\egroup}
+\let\markbothOrig\markboth
+\def\markboth#1#2{\bgroup
+  \let\texttt=\textttOrig\markbothOrig{#1}{#2}\egroup}
+\let\markrightOrig\markright
+\def\markright#1{\bgroup
+  \let\texttt=\textttOrig\markrightOrig{#1}\egroup}
+
+
+\IfFileExists{parskip.sty}{%
+\usepackage{parskip}
+}{% else
+\setlength{\parindent}{0pt}
+\setlength{\parskip}{6pt plus 2pt minus 1pt}
+}
+\setlength{\emergencystretch}{3em}  % prevent overfull lines
+\providecommand{\tightlist}{%
+  \setlength{\itemsep}{0pt}\setlength{\parskip}{0pt}}
+\setcounter{secnumdepth}{0}
+% Redefines (sub)paragraphs to behave more like sections
+\ifx\paragraph\undefined\else
+\let\oldparagraph\paragraph
+\renewcommand{\paragraph}[1]{\oldparagraph{#1}\mbox{}}
+\fi
+\ifx\subparagraph\undefined\else
+\let\oldsubparagraph\subparagraph
+\renewcommand{\subparagraph}[1]{\oldsubparagraph{#1}\mbox{}}
+\fi
+
+\title{\texttt{JAXbind}: Easy bindings to JAX}
+
+        \author[1, 2, 3*]{Jakob Roth}
+          \author[1*]{Martin Reinecke}
+          \author[1, 2*]{Gordian Edenhofer}
+    
+      \affil[1]{Max Planck Institute for Astrophysics,
+Karl-Schwarzschild-Str. 1, 85748 Garching, Germany}
+      \affil[2]{Ludwig Maximilian University of Munich,
+Geschwister-Scholl-Platz 1, 80539 Munich, Germany}
+      \affil[3]{Technical University of Munich, Boltzmannstr. 3, 85748
+Garching, Germany}
+      \affil[*]{These authors contributed equally.}
+  \date{\vspace{-7ex}}
+
+\begin{document}
+\maketitle
+
+\vspace{1em}
+
+\hypertarget{summary}{%
+\section{Summary}\label{summary}}
+
+JAX is widely used in machine learning and scientific computing.
+Scientific computing relies on existing high-performance code which we
+would ideally like to use in JAX. Reimplementing the existing code in
+JAX is often impractical and the existing interface in JAX for
+connecting custom code requires deep knowledge of JAX and its C++
+backend. The aim of \texttt{JAXbind} is to drastically lower the burden
+of connecting custom functions implemented in other programming
+languages to JAX. Specifically, \texttt{JAXbind} provides an easy-to-use
+Python interface for defining custom, so-called JAX primitives
+supporting any JAX transformations.
+
+\hypertarget{statement-of-need}{%
+\section{Statement of Need}\label{statement-of-need}}
+
+The use of JAX (Bradbury et al., 2018) is widespread in the natural
+sciences. JAX's powerful transformation system is of especially high
+interest. It enables retrieving arbitrary derivatives of functions,
+batch computations, and just-in-time code compilation for additional
+performance. Its transformation system relies on all constituents of the
+computation being written in JAX.
+
+A plethora of high-performance code is not written in JAX and thus not
+accessible from within JAX. Rewriting these is often infeasible and/or
+inefficient. Ideally, we would like to intermix existing
+high-performance code with JAX code. However, connecting code to JAX
+requires knowledge of the internals of JAX and its C++ backend.
+
+In this paper, we present \texttt{JAXbind}, a package for bridging any
+function to JAX without in-depth knowledge of JAX's transformation
+system. The interface is accessible from Python without requiring any
+development in C++. The package is able to register any function, its
+partial derivatives and their transpose functions as a JAX native call,
+a so-called primitive. Derivatives, compilation rules, and batching
+rules are automatically registered with JAX.
+
+We believe \texttt{JAXbind} to be highly useful in scientific computing.
+We intend to use this package to connect the Hartley transform and the
+spherical harmonic transform from ducc (Reinecke, 2024) to NIFTy
+(Edenhofer et al., 2024) as well as the radio interferometry response
+from ducc with \texttt{resolve} (Arras et al., 2024) for radio
+astronomy. Furthermore, we intend to connect the non-uniform FFT from
+ducc with JAX for applications in strong-lensing astrophysics. We
+envision many further applications within and outside of astrophysics.
+
+To the best of our knowledge no other code currently exists for
+connecting generic functions to JAX. The package that comes the closest
+is Enzyme-JAX (EnzymeAD, 2024). Enzyme-JAX allows one to differentiate a
+C++ function with Enzyme (W. Moses \& Churavy, 2020; W. S. Moses et al.,
+2021, 2022) and connect it together with its derivative to JAX. However,
+it enforces the use of Enzyme for deriving derivatives and does not
+allow for connecting arbitrary code to JAX.
+
+\hypertarget{automatic-differentiation-and-code-example}{%
+\section{Automatic Differentiation and Code
+Example}\label{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
+\texttt{JAXbind} package expects the Python function for \(f\) to take
+three positional arguments. The first argument, \texttt{out}, is a
+\texttt{tuple} into which the function results are written. The second
+argument is also a \texttt{tuple} containing the input to the function,
+in our case, \(x_1\) and \(x_2\). Via \texttt{kwargs\_dump}, potential
+keyword arguments given to the later registered Jax primitive can be
+forwarded to \texttt{f} in serialized form.
+
+\begin{Shaded}
+\begin{Highlighting}[]
+\ImportTok{import}\NormalTok{ jaxbind }
+
+\KeywordTok{def}\NormalTok{ f(out, args, kwargs\_dump):}
+\NormalTok{    kwargs }\OperatorTok{=}\NormalTok{ jaxbind.load\_kwargs(kwargs\_dump)}
+\NormalTok{    x1, x2 }\OperatorTok{=}\NormalTok{ args}
+\NormalTok{    out[}\DecValTok{0}\NormalTok{][()] }\OperatorTok{=}\NormalTok{ x1 }\OperatorTok{*}\NormalTok{ x2}\OperatorTok{**}\DecValTok{2}
+\end{Highlighting}
+\end{Shaded}
+
+JAX's automatic differentiation engine can compute the Jacobian-vector
+product \texttt{jvp} and vector-Jacobian product \texttt{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
+\texttt{jvp}. Instead, an implementation of the pushforward has to be
+provided, which \texttt{JAXbind} will register as the \texttt{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\).
+
+\begin{Shaded}
+\begin{Highlighting}[]
+\KeywordTok{def}\NormalTok{ f\_jvp(out, args, kwargs\_dump):}
+\NormalTok{    kwargs }\OperatorTok{=}\NormalTok{ jaxbind.load\_kwargs(kwargs\_dump)}
+\NormalTok{    x1, x2, dx1, dx2 }\OperatorTok{=}\NormalTok{ args}
+\NormalTok{    out[}\DecValTok{0}\NormalTok{][()] }\OperatorTok{=}\NormalTok{ x2}\OperatorTok{**}\DecValTok{2} \OperatorTok{*}\NormalTok{ dx1 }\OperatorTok{+} \DecValTok{2} \OperatorTok{*}\NormalTok{ x1 }\OperatorTok{*}\NormalTok{ x2 }\OperatorTok{*}\NormalTok{ dx2}
+\end{Highlighting}
+\end{Shaded}
+
+The vector-Jacobian product \texttt{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 \texttt{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)\).
+
+\begin{Shaded}
+\begin{Highlighting}[]
+\KeywordTok{def}\NormalTok{ f\_vjp(out, args, kwargs\_dump):}
+\NormalTok{    kwargs }\OperatorTok{=}\NormalTok{ jaxbind.load\_kwargs(kwargs\_dump)}
+\NormalTok{    x1, x2, dy }\OperatorTok{=}\NormalTok{ args}
+\NormalTok{    out[}\DecValTok{0}\NormalTok{][()] }\OperatorTok{=}\NormalTok{ x2}\OperatorTok{**}\DecValTok{2} \OperatorTok{*}\NormalTok{ dy}
+\NormalTok{    out[}\DecValTok{1}\NormalTok{][()] }\OperatorTok{=} \DecValTok{2} \OperatorTok{*}\NormalTok{ x1 }\OperatorTok{*}\NormalTok{ x2 }\OperatorTok{*}\NormalTok{ dy}
+\end{Highlighting}
+\end{Shaded}
+
+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 \texttt{f}
+as well as for the \texttt{vjp} application. The output shape of the
+\texttt{jvp} is identical to the output shape of \texttt{f} itself and
+does not need to be specified again.
+
+\begin{Shaded}
+\begin{Highlighting}[]
+\KeywordTok{def}\NormalTok{ f\_abstract(}\OperatorTok{*}\NormalTok{args, }\OperatorTok{**}\NormalTok{kwargs):}
+    \ControlFlowTok{assert}\NormalTok{ args[}\DecValTok{0}\NormalTok{].shape }\OperatorTok{==}\NormalTok{ args[}\DecValTok{1}\NormalTok{].shape}
+    \ControlFlowTok{return}\NormalTok{ ((args[}\DecValTok{0}\NormalTok{].shape, args[}\DecValTok{0}\NormalTok{].dtype),)}
+
+\KeywordTok{def}\NormalTok{ f\_abstract\_T(}\OperatorTok{*}\NormalTok{args, }\OperatorTok{**}\NormalTok{kwargs):}
+    \ControlFlowTok{return}\NormalTok{ (}
+\NormalTok{        (args[}\DecValTok{0}\NormalTok{].shape, args[}\DecValTok{0}\NormalTok{].dtype),}
+\NormalTok{        (args[}\DecValTok{0}\NormalTok{].shape, args[}\DecValTok{0}\NormalTok{].dtype),}
+\NormalTok{    )}
+\end{Highlighting}
+\end{Shaded}
+
+We have now defined all ingredients necessary to register a JAX
+primitive for our function \(f\) using the \texttt{JAXbind} package.
+
+\begin{Shaded}
+\begin{Highlighting}[]
+\NormalTok{f\_jax }\OperatorTok{=}\NormalTok{ jaxbind.get\_nonlinear\_call(}
+\NormalTok{    f, (f\_jvp, f\_vjp), f\_abstract, f\_abstract\_T}
+\NormalTok{)}
+\end{Highlighting}
+\end{Shaded}
+
+\texttt{f\_jax} is a JAX primitive registered via the \texttt{JAXbind}
+package supporting all JAX transformations. We can now compute the
+\texttt{jvp} and \texttt{vjp} of the new JAX primitive and even
+jit-compile and batch it.
+
+\begin{Shaded}
+\begin{Highlighting}[]
+\ImportTok{import}\NormalTok{ jax}
+\ImportTok{import}\NormalTok{ jax.numpy }\ImportTok{as}\NormalTok{ jnp}
+
+\NormalTok{inp }\OperatorTok{=}\NormalTok{ (jnp.full((}\DecValTok{4}\NormalTok{,}\DecValTok{3}\NormalTok{), }\FloatTok{4.}\NormalTok{), jnp.full((}\DecValTok{4}\NormalTok{,}\DecValTok{3}\NormalTok{), }\FloatTok{2.}\NormalTok{))}
+\NormalTok{tan }\OperatorTok{=}\NormalTok{ (jnp.full((}\DecValTok{4}\NormalTok{,}\DecValTok{3}\NormalTok{), }\FloatTok{1.}\NormalTok{), jnp.full((}\DecValTok{4}\NormalTok{,}\DecValTok{3}\NormalTok{), }\FloatTok{1.}\NormalTok{))}
+\NormalTok{res, res\_tan }\OperatorTok{=}\NormalTok{ jax.jvp(f\_jax, inp, tan)}
+
+\NormalTok{cotan }\OperatorTok{=}\NormalTok{ (jnp.full((}\DecValTok{4}\NormalTok{,}\DecValTok{3}\NormalTok{), }\FloatTok{6.}\NormalTok{),)}
+\NormalTok{res, f\_vjp }\OperatorTok{=}\NormalTok{ jax.vjp(f\_jax, }\OperatorTok{*}\NormalTok{inp)}
+\NormalTok{res\_cotan }\OperatorTok{=}\NormalTok{ f\_vjp(cotan)}
+
+\NormalTok{f\_jax\_jit }\OperatorTok{=}\NormalTok{ jax.jit(f\_jax)}
+\NormalTok{res }\OperatorTok{=}\NormalTok{ f\_jax\_jit(}\OperatorTok{*}\NormalTok{inp)}
+\end{Highlighting}
+\end{Shaded}
+
+\hypertarget{higher-order-derivatives-and-linear-functions}{%
+\section{Higher Order Derivatives and Linear
+Functions}\label{higher-order-derivatives-and-linear-functions}}
+
+JAX supports higher order derivatives and can differentiate a
+\texttt{jvp} or \texttt{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 \texttt{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 \texttt{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
+\texttt{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, \texttt{JAXbind} provides a special
+interface for linear functions, supporting higher order derivatives,
+only requiring an implementation of the function and its transpose.
+
+\hypertarget{platforms}{%
+\section{Platforms}\label{platforms}}
+
+Currently, \texttt{JAXbind} only has CPU but no GPU support. With some
+expertise on Python bindings for GPU kernels adding GPU support should
+be fairly simple. We especially want to highlight that the interfacing
+with the JAX automatic differentiation engine is identical for CPU and
+GPU.
+
+\hypertarget{acknowledgements}{%
+\section{Acknowledgements}\label{acknowledgements}}
+
+We would like to thank Dan Foreman-Mackey for his detailed guide
+(https://dfm.io/posts/extending-jax/) on connecting C++ code to JAX.
+Jakob Roth acknowledges financial support from the German Federal
+Ministry of Education and Research (BMBF) under grant 05A23WO1
+(Verbundprojekt D-MeerKAT III). Gordian Edenhofer acknowledges support
+from the German Academic Scholarship Foundation in the form of a PhD
+scholarship (``Promotionsstipendium der Studienstiftung des Deutschen
+Volkes'').
+
+\hypertarget{references}{%
+\section{References}\label{references}}
+
+\hypertarget{refs}{}
+\begin{CSLReferences}{1}{0}
+\leavevmode\hypertarget{ref-Resolve2024}{}%
+Arras, P., Roth, J., Ding, S., Reinecke, M., Fuchs, R., \& Johnson, V.
+(2024). \emph{RESOLVE}. \url{https://gitlab.mpcdf.mpg.de/ift/resolve}
+
+\leavevmode\hypertarget{ref-Jax2018}{}%
+Bradbury, J., Frostig, R., Hawkins, P., Johnson, M. J., Leary, C.,
+Maclaurin, D., Necula, G., Paszke, A., VanderPlas, J., Wanderman-Milne,
+S., \& Zhang, Q. (2018). \emph{{JAX}: Composable transformations of
+{P}ython+{N}um{P}y programs} (Version 0.3.13) {[}Computer software{]}.
+\url{http://github.com/google/jax}
+
+\leavevmode\hypertarget{ref-Edenhofer2023NIFTyRE}{}%
+Edenhofer, G., Frank, P., Roth, J., Leike, R. H., Guerdi, M.,
+Scheel-Platz, L. I., Guardiani, M., Eberle, V., Westerkamp, M., \&
+Enßlin, T. A. (2024). \emph{{Re-Envisioning Numerical Information Field
+Theory (NIFTy.re): A Library for Gaussian Processes and Variational
+Inference}}. \url{http://arxiv.org/abs/2402.16683}
+
+\leavevmode\hypertarget{ref-Moses2024}{}%
+EnzymeAD. (2024). \emph{{Enzyme-JAX}} (Version 0.0.6) {[}Computer
+software{]}. \url{https://github.com/EnzymeAD/Enzyme-JAX}
+
+\leavevmode\hypertarget{ref-Moses2020}{}%
+Moses, W., \& Churavy, V. (2020). Instead of rewriting foreign code for
+machine learning, automatically synthesize fast gradients. In H.
+Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, \& H. Lin (Eds.),
+\emph{Advances in neural information processing systems} (Vol. 33, pp.
+12472--12485). Curran Associates, Inc.
+\url{https://proceedings.neurips.cc/paper/2020/file/9332c513ef44b682e9347822c2e457ac-Paper.pdf}
+
+\leavevmode\hypertarget{ref-Moses2021}{}%
+Moses, W. S., Churavy, V., Paehler, L., Hückelheim, J., Narayanan, S. H.
+K., Schanen, M., \& Doerfert, J. (2021). Reverse-mode automatic
+differentiation and optimization of GPU kernels via enzyme.
+\emph{Proceedings of the International Conference for High Performance
+Computing, Networking, Storage and Analysis}.
+\url{https://doi.org/10.1145/3458817.3476165}
+
+\leavevmode\hypertarget{ref-Moses2022}{}%
+Moses, W. S., Narayanan, S. H. K., Paehler, L., Churavy, V., Schanen,
+M., Hückelheim, J., Doerfert, J., \& Hovland, P. (2022). Scalable
+automatic differentiation of multiple parallel paradigms through
+compiler augmentation. \emph{Proceedings of the International Conference
+on High Performance Computing, Networking, Storage and Analysis}.
+ISBN:~\href{https://worldcat.org/isbn/9784665454445}{9784665454445}
+
+\leavevmode\hypertarget{ref-ducc0}{}%
+Reinecke, M. (2024). \emph{{DUCC}: Distinctly useful code collection}
+(Version 0.33.0) {[}Computer software{]}.
+\url{https://gitlab.mpcdf.mpg.de/mtr/ducc}
+
+\end{CSLReferences}
+
+\end{document}
diff --git a/pyproject.toml b/pyproject.toml
index f3dd715..df232f1 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [project]
 requires-python = ">=3.8"
-name = "jax_linop"
+name = "jaxbind"
 readme = "README.md"
 authors = [
   {name = "Martin Reinecke", email = "martin@mpa-garching.mpg.de"},
@@ -8,5 +8,5 @@ authors = [
 dynamic = ["license", "description", "version", "urls", "classifiers", "dependencies"]
 
 [build-system]
-requires = ["setuptools >= 40.6.0", "pybind11 >= 2.6.0", "numpy >= 1.17.0"]
+requires = ["setuptools >= 40.6.0", "pybind11 >= 2.6.0", "numpy >= 1.17.0", "jax >= 0.4", "jaxlib"]
 build-backend = "setuptools.build_meta"
diff --git a/setup.cfg b/setup.cfg
index 4476ba7..6ab0771 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -1,10 +1,10 @@
 [metadata]
-name = jax_linop
+name = jaxbind
 description = Convenience functions for adding arbitrary linear operations to JAX.
 long_description = file: README.md
 long_description_content_type = text/markdown
-url = https://gitlab.mpcdf.mpg.de/mtr/jax_linop
-author = Martin Reinecke
+url = https://https://github.com/NIFTy-PPL/JAXbind
+author = Gordian Edenhofer, Martin Reinecke, Jakob Roth
 author_email = martin@mpa-garching.mpg.de
 python_requires = >=3.8
 license = BSD3
@@ -18,11 +18,11 @@ classifiers =
   Programming Language :: Python
   Intended Audience :: Science/Research
 project_urls =
-  Bug Tracker = https://gitlab.mpcdf.mpg.de/mtr/jax_linop/issues
-  Source Code = https://gitlab.mpcdf.mpg.de/mtr/jax_linop
+  Bug Tracker = https://gitlab.mpcdf.mpg.de/mtr/jaxbind/issues
+  Source Code = https://gitlab.mpcdf.mpg.de/mtr/jaxbind
   README = https://gitlab.mpcdf.mpg.de/mtr/ducc/-/blob/main/README.md
 
 [options]
 include_package_data = True
 packages = 
-install_requires = numpy>=1.17.0; jax>=0.4
+install_requires = numpy>=1.17.0; jax>=0.4; jaxlib
diff --git a/setup.py b/setup.py
index b661210..13e918a 100644
--- a/setup.py
+++ b/setup.py
@@ -1,32 +1,54 @@
 import sys
+import os
+import codecs
 import pybind11
 from setuptools import find_packages, setup, Extension
 
-pkgname = 'jax_linop'
-version = '0.1.0'
+pkgname = "jaxbind"
 
-include_dirs = [pybind11.get_include(True), pybind11.get_include(False)]
 
-extra_compile_args = ['-std=c++17', '-fvisibility=hidden']
+def read(rel_path):
+    here = os.path.abspath(os.path.dirname(__file__))
+    with codecs.open(os.path.join(here, rel_path), "r") as fp:
+        return fp.read()
+
+
+def get_version(rel_path):
+    for line in read(rel_path).splitlines():
+        if line.startswith("__version__"):
+            delim = '"' if '"' in line else "'"
+            return line.split(delim)[1]
+    else:
+        raise RuntimeError("Unable to find version string.")
+
+
+include_dirs = [pybind11.get_include()]
+
+extra_compile_args = ["-std=c++17", "-fvisibility=hidden"]
 
 python_module_link_args = []
 
-if sys.platform == 'darwin':
-    extra_compile_args += ['-mmacosx-version-min=10.14']
-    python_module_link_args += ['-mmacosx-version-min=10.14']
-elif sys.platform == 'win32':
-    extra_compile_args = ['/EHsc', '/std:c++17']
-
-extensions = [Extension("_jax_linop",
-                        language='c++',
-                        sources=['src/_jax_linop.cc'],
-                        depends=['src/_jax_linop.cc', 'setup.py'],
-                        include_dirs=include_dirs,
-                        extra_compile_args=extra_compile_args,
-                        extra_link_args=python_module_link_args)]
-
-setup(name=pkgname,
-      version=version,
-      packages=find_packages(include=["jax_linop"]),
-      ext_modules=extensions
-      )
+if sys.platform == "darwin":
+    extra_compile_args += ["-mmacosx-version-min=10.14"]
+    python_module_link_args += ["-mmacosx-version-min=10.14"]
+elif sys.platform == "win32":
+    extra_compile_args = ["/EHsc", "/std:c++17"]
+
+extensions = [
+    Extension(
+        "_jaxbind",
+        language="c++",
+        sources=["src/_jaxbind.cc"],
+        depends=["src/_jaxbind.cc", "setup.py"],
+        include_dirs=include_dirs,
+        extra_compile_args=extra_compile_args,
+        extra_link_args=python_module_link_args,
+    )
+]
+
+setup(
+    name=pkgname,
+    version=get_version(f"{pkgname}/__init__.py"),
+    packages=find_packages(include=["jaxbind"]),
+    ext_modules=extensions,
+)
diff --git a/src/_jax_linop.cc b/src/_jax_linop.cc
deleted file mode 100644
index daa9808..0000000
--- a/src/_jax_linop.cc
+++ /dev/null
@@ -1,104 +0,0 @@
-/* SPDX-License-Identifier: BSD-3-Clause */
-
-/*
- *  Jax_linop is being developed at the Max-Planck-Institut fuer Astrophysik
- */
-
-/*
- *  Copyright (C) 2023, 2024 Max-Planck-Society
- *  Authors: Martin Reinecke, Jakob Roth, Gordian Edenhofer 
- */
-
-#include <pybind11/pybind11.h>
-#include <pybind11/numpy.h>
-
-#include <vector>
-#include <map>
-
-namespace detail_pymodule_jax {
-
-namespace py = pybind11;
-using namespace std;
-
-// https://en.cppreference.com/w/cpp/numeric/bit_cast
-template <class To, class From>
-typename std::enable_if<sizeof(To) == sizeof(From) && std::is_trivially_copyable<From>::value &&
-                            std::is_trivially_copyable<To>::value,
-                        To>::type
-bit_cast(const From& src) noexcept
-  {
-  static_assert(
-      std::is_trivially_constructible<To>::value,
-      "This implementation additionally requires destination type to be trivially constructible");
-
-  To dst;
-  memcpy(&dst, &src, sizeof(To));
-  return dst;
-  }
-
-template <typename T>
-pybind11::capsule EncapsulateFunction(T* fn)
-  { return pybind11::capsule(bit_cast<void*>(fn), "xla._CUSTOM_CALL_TARGET"); }
-
-void pycall(void *out, void **in)
-  {
-  py::gil_scoped_acquire get_GIL;
-
-  static const map<uint8_t, py::object> tcdict = {
-    { 3, py::dtype::of<float>()},
-    { 7, py::dtype::of<double>()},
-    {67, py::dtype::of<complex<float>>()},
-    {71, py::dtype::of<complex<double>>()}};
-
-  // the "opid" in in[1] is not used; it is only passed to guarantee uniqueness
-  // of the passed parameters for every distinvt operator, so that JAX knows
-  // when and when not to recompile.
-
-  // Getting the "state" dictionary from the passed ID
-  py::handle hnd(*reinterpret_cast<PyObject **>(in[2]));
-  auto obj = py::reinterpret_borrow<py::object>(hnd);
-  const py::dict state(obj);
-
-  size_t idx = 3;
-  // Getting type, rank, and shape of the input
-  auto dtin = tcdict.at(uint8_t(*reinterpret_cast<int64_t *>(in[idx++])));
-  size_t ndim_in = *reinterpret_cast<uint64_t *>(in[idx++]);
-  vector<size_t> shape_in;
-  for (size_t i=0; i<ndim_in; ++i)
-    shape_in.push_back(*reinterpret_cast<uint64_t *>(in[idx++]));
-  // Getting type, rank, and shape of the output
-  auto dtout = tcdict.at(uint8_t(*reinterpret_cast<int64_t *>(in[idx++])));
-  size_t ndim_out = *reinterpret_cast<uint64_t *>(in[idx++]);
-  vector<size_t> shape_out;
-  for (size_t i=0; i<ndim_out; ++i)
-    shape_out.push_back(*reinterpret_cast<uint64_t *>(in[idx++]));
-
-  // Building "pseudo" numpy.ndarays on top of the provided memory regions.
-  // This should be completely fine, as long as the called function does not
-  // keep any references to them.
-  py::str dummy;
-  py::array pyin (dtin, shape_in, in[0], dummy);
-//  MR_assert(!pyin.owndata(), "owndata should be false");
-  pyin.attr("flags").attr("writeable")=false;
-//  MR_assert(!pyin.writeable(), "input array should not be writeable");
-  py::array pyout (dtout, shape_out, out, dummy);
-//  MR_assert(!pyout.owndata(), "owndata should be false");
-//  MR_assert(pyout.writeable(), "output data must be writable");
-
-  // Execute the Python function implementing the linear operation
-  state["_func"](pyin, pyout, state);
-  }
-
-pybind11::dict Registrations()
-  {
-  pybind11::dict dict;
-  dict["cpu_pycall"] = EncapsulateFunction(pycall);
-  return dict;
-  }
-
-}
-
-PYBIND11_MODULE(_jax_linop, m) {
-  m.def("registrations", detail_pymodule_jax::Registrations);
-}
-
diff --git a/src/_jaxbind.cc b/src/_jaxbind.cc
new file mode 100644
index 0000000..2e0fcb6
--- /dev/null
+++ b/src/_jaxbind.cc
@@ -0,0 +1,118 @@
+/* SPDX-License-Identifier: BSD-2-Clause */
+
+/*
+ *  JAXbind is being developed at the Max-Planck-Institut fuer Astrophysik
+ */
+
+/*
+ *  Copyright (C) 2023, 2024 Max-Planck-Society
+ *  Authors: Martin Reinecke, Jakob Roth, Gordian Edenhofer
+ */
+
+#include <pybind11/pybind11.h>
+#include <pybind11/numpy.h>
+#include <iostream>
+
+#include <vector>
+#include <map>
+
+namespace detail_pymodule_jax {
+
+namespace py = pybind11;
+using namespace std;
+
+// https://en.cppreference.com/w/cpp/numeric/bit_cast
+template <class To, class From>
+typename std::enable_if<sizeof(To) == sizeof(From) && std::is_trivially_copyable<From>::value &&
+                            std::is_trivially_copyable<To>::value,
+                        To>::type
+bit_cast(const From& src) noexcept
+  {
+  static_assert(
+      std::is_trivially_constructible<To>::value,
+      "This implementation additionally requires destination type to be trivially constructible");
+
+  To dst;
+  memcpy(&dst, &src, sizeof(To));
+  return dst;
+  }
+
+template <typename T>
+pybind11::capsule EncapsulateFunction(T* fn)
+  { return pybind11::capsule(bit_cast<void*>(fn), "xla._CUSTOM_CALL_TARGET"); }
+
+void pycall(void *out_raw, void **in)
+  {
+  py::gil_scoped_acquire get_GIL;
+
+  static const map<uint8_t, py::object> tcdict = {
+    { 3, py::dtype::of<float>()},
+    { 7, py::dtype::of<double>()},
+    {32, py::dtype::of<uint8_t>()},
+    {39, py::dtype::of<uint64_t>()},
+    {67, py::dtype::of<complex<float>>()},
+    {71, py::dtype::of<complex<double>>()}};
+
+  py::str dummy;
+
+  py::handle hnd(*reinterpret_cast<PyObject **>(in[0]));
+  auto func = py::reinterpret_borrow<py::object>(hnd);
+
+  size_t idx = 1;
+  size_t nargs = *reinterpret_cast<uint64_t *>(in[idx++]);
+  py::list py_in;
+  for (size_t i=0; i<nargs; i++) {
+    // Getting type, rank, and shape of the input
+    auto dtp_a = tcdict.at(uint8_t(*reinterpret_cast<int64_t *>(in[idx++])));
+    size_t ndim_a = *reinterpret_cast<uint64_t *>(in[idx++]);
+    vector<size_t> shape_a;
+    for (size_t j=0; j<ndim_a; ++j) {
+      shape_a.push_back(*reinterpret_cast<uint64_t *>(in[idx++]));
+    }
+    // Building "pseudo" numpy.ndarays on top of the provided memory regions.
+    // This should be completely fine, as long as the called function does not
+    // keep any references to them.
+    py::array py_a (dtp_a, shape_a, in[idx++], dummy);
+    py_a.attr("flags").attr("writeable") = false;
+    py_in.append(py_a);
+  }
+
+  // if we have only one output, out_raw is a void * pointing to the data of this output
+  // otherwise, out_raw is a void ** pointing to an array of void * pointing to the individual data
+  void **out = reinterpret_cast<void **>(out_raw);
+  void *out_single = reinterpret_cast<void *>(out_raw);
+  size_t nout = *reinterpret_cast<uint64_t *>(in[idx++]);
+  py::list py_out;
+  for (size_t i=0; i<nout; i++) {
+    // Getting type, rank, and shape of the output
+    auto dtp_out = tcdict.at(uint8_t(*reinterpret_cast<int64_t *>(in[idx++])));
+    size_t ndim_out = *reinterpret_cast<uint64_t *>(in[idx++]);
+    vector<size_t> shape_out;
+    for (size_t j=0; j<ndim_out; ++j) {
+      shape_out.push_back(*reinterpret_cast<uint64_t *>(in[idx++]));
+    }
+    py::array py_o (dtp_out, shape_out, (nout==1) ? out_single : out[i], dummy);
+    py_out.append(py_o);
+  }
+
+  auto dtp_kwargs = tcdict.at(uint8_t(*reinterpret_cast<int64_t *>(in[idx++])));
+  size_t size_kwargs = *reinterpret_cast<uint64_t *>(in[idx++]);
+  py::array py_kwargs (dtp_kwargs, size_kwargs, in[idx++], dummy);
+
+  // Execute the Python function implementing the linear operation
+  func(py_out, py_in, py_kwargs);
+  }
+
+pybind11::dict Registrations()
+  {
+  pybind11::dict dict;
+  dict["cpu_pycall"] = EncapsulateFunction(pycall);
+  return dict;
+  }
+
+}
+
+PYBIND11_MODULE(_jaxbind, m) {
+  m.def("registrations", detail_pymodule_jax::Registrations);
+}
+


Test patch:
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
new file mode 100644
index 0000000..3fbd2e5
--- /dev/null
+++ b/.github/workflows/test.yml
@@ -0,0 +1,38 @@
+name: Test python application
+on:
+  push:
+    branches: [ "multi_arg" ]
+  pull_request:
+    branches: [ "multi_arg" ]
+permissions:
+  contents: read
+jobs:
+  pytest:
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v4
+    - uses: ./.github/actions/setup_package
+    - name: Test with pytest
+      run: |
+        pytest test
+  demos_01_linear_function:
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v4
+    - uses: ./.github/actions/setup_package
+    - name: Run linear function demo
+      run: python demos/01_linear_function.py
+  demos_02_multilinear_function:
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v4
+    - uses: ./.github/actions/setup_package
+    - name: Run linear function demo
+      run: python demos/02_multilinear_function.py
+  demos_03_nonlinear_function:
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v4
+    - uses: ./.github/actions/setup_package
+    - name: Run linear function demo
+      run: python demos/03_nonlinear_function.py
diff --git a/test/test_batching.py b/test/test_batching.py
new file mode 100644
index 0000000..3a047fa
--- /dev/null
+++ b/test/test_batching.py
@@ -0,0 +1,131 @@
+# SPDX-License-Identifier: BSD-2-Clause
+# Authors: Martin Reinecke, Jakob Roth, Gordian Edenhofer
+
+# Copyright(C) 2024 Max-Planck-Society
+
+import jax
+import numpy as np
+import pytest
+import scipy
+from jax.test_util import check_grads
+
+import jaxbind
+
+pmp = pytest.mark.parametrize
+
+jax.config.update("jax_enable_x64", True)
+
+
+def f(out, args, kwargs_dump):
+    x, y = args
+    kwargs = jaxbind.load_kwargs(kwargs_dump)
+    workers = kwargs.pop("workers", None)
+    batch_axes = kwargs.pop("batch_axes", None)
+    axes0 = list(range(len(x.shape)))
+    axes1 = list(range(len(x.shape)))
+    if batch_axes:
+        axes0 = [i for i in range(len(x.shape)) if not i in batch_axes[0]]
+        axes1 = [i for i in range(len(y.shape)) if not i in batch_axes[1]]
+    out[0][()] = scipy.fft.fftn(x, axes=axes0, norm="forward", workers=workers)
+    out[1][()] = scipy.fft.fftn(y, axes=axes1, norm="forward", workers=workers)
+
+
+def f_T(out, args, kwargs_dump):
+    x, y = args
+    kwargs = jaxbind.load_kwargs(kwargs_dump)
+    workers = kwargs.pop("workers", None)
+    batch_axes = kwargs.pop("batch_axes", None)
+    axes0 = list(range(len(x.shape)))
+    axes1 = list(range(len(x.shape)))
+    if batch_axes:
+        axes0 = [i for i in range(len(x.shape)) if not i in batch_axes[0]]
+        axes1 = [i for i in range(len(y.shape)) if not i in batch_axes[1]]
+    out[0][()] = scipy.fft.ifftn(
+        x.conj(), axes=axes0, norm="backward", workers=workers
+    ).conj()
+    out[1][()] = scipy.fft.ifftn(
+        y.conj(), axes=axes1, norm="backward", workers=workers
+    ).conj()
+
+
+def f_a(*args, **kwargs):
+    # Returns `shape` and `dtype` of output as well as the added batch_axes of the `output``
+    batch_axes = kwargs.pop("batch_axes", None)
+    x, y = args
+    out_ax = [(), ()]
+    if batch_axes:
+        if len(batch_axes[0]) > 0 and len(batch_axes[1]) > 0:
+            out_ax[0] = batch_axes[0][-1]
+            out_ax[1] = batch_axes[1][-1]
+        else:
+            raise RuntimeError("Batching along only 1 input axis not implemented.")
+    return ((x.shape, x.dtype, out_ax[0]), (y.shape, y.dtype, out_ax[1]))
+
+
+def f_a_T(*args, **kwargs):
+    # Returns `shape` and `dtype` of output as well as the added batch_axes of the `output``
+    batch_axes = kwargs.pop("batch_axes", None)
+    a, b = args
+    out_ax = [(), ()]
+    if batch_axes:
+        if len(batch_axes[0]) > 0 and len(batch_axes[1]) > 0:
+            out_ax[0] = batch_axes[0][-1]
+            out_ax[1] = batch_axes[1][-1]
+        else:
+            raise RuntimeError("Batching along only 1 input axis not implemented.")
+    return ((a.shape, a.dtype, out_ax[0]), (b.shape, b.dtype, out_ax[1]))
+
+
+f_jax = jaxbind.get_linear_call(f, f_T, f_a, f_a_T, func_can_batch=False)
+f_jax_can_batch = jaxbind.get_linear_call(f, f_T, f_a, f_a_T, func_can_batch=True)
+
+
+rng = np.random.default_rng(42)
+a1 = rng.random((5, 5, 5)) - 0.5 + 1j * (rng.random((5, 5, 5)) - 0.5)
+a2 = rng.random((5, 5, 5)) - 0.5 + 1j * (rng.random((5, 5, 5)) - 0.5)
+a = (a1, a2)
+av1 = rng.random((5, 5, 5, 5, 5)) - 0.5 + 1j * (rng.random((5, 5, 5, 5, 5)) - 0.5)
+av2 = rng.random((5, 5, 5, 5, 5)) - 0.5 + 1j * (rng.random((5, 5, 5, 5, 5)) - 0.5)
+av = (av1, av2)
+
+av1_diff_ax_len = rng.random((5, 5, 5, 3, 4)) - 0.5 + 1j * (rng.random((5, 5, 5, 3, 4)) - 0.5)
+av2_diff_ax_len = rng.random((5, 5, 5, 4, 3)) - 0.5 + 1j * (rng.random((5, 5, 5, 4, 3)) - 0.5)
+av_diff_ax_len = (av1_diff_ax_len, av2_diff_ax_len)
+
+@pmp("bt_a1", (0, 1, 2))
+@pmp("bt_a2", (0, 1, 2))
+@pmp("bt2_a1,bt2_a2", ((3, 4), (4, 3), (0, 0)))
+@pmp("o_a1", (0, 1))
+@pmp("o_a2", (0, 1))
+def test_vmap(bt_a1, bt_a2, bt2_a1, bt2_a2, o_a1, o_a2):
+    vj = jax.vmap(f_jax, in_axes=(bt_a1, bt_a2), out_axes=[o_a1, o_a2])
+    vb = jax.vmap(f_jax_can_batch, in_axes=(bt_a1, bt_a2), out_axes=[o_a1, o_a2])
+    rj = vj(*a)
+    rb = vb(*a)
+    np.testing.assert_allclose(rj, rb)
+    assert rj[0].shape == rb[0].shape and rj[1].shape == rb[1].shape
+    # check_grads(vj, a, order=2, modes=["fwd", "rev"], eps=1.0)
+    check_grads(vb, a, order=2, modes=["fwd", "rev"], eps=1.0)
+
+    vvj = jax.vmap(vj, in_axes=(bt2_a1, bt2_a2), out_axes=[o_a1, o_a2])
+    vvb = jax.vmap(vb, in_axes=(bt2_a1, bt2_a2), out_axes=[o_a1, o_a2])
+    rb = vvb(*av)
+    rj = vvj(*av)
+    np.testing.assert_allclose(rj[0], rb[0])
+    np.testing.assert_allclose(rj[1], rb[1])
+    assert rj[0].shape == rb[0].shape and rj[1].shape == rb[1].shape
+    # check_grads(vvj, av, order=2, modes=["fwd", "rev"], eps=1.0)
+    # check_grads(vvb, av, order=2, modes=["fwd", "rev"], eps=1.0)
+
+    if not (bt2_a1,bt2_a2) ==(0, 0):
+        vvj = jax.vmap(vj, in_axes=(bt2_a1, bt2_a2), out_axes=[o_a1, o_a2])
+        vvb = jax.vmap(vb, in_axes=(bt2_a1, bt2_a2), out_axes=[o_a1, o_a2])
+        rb = vvb(*av_diff_ax_len)
+        rj = vvj(*av_diff_ax_len)
+        np.testing.assert_allclose(rj[0], rb[0])
+        np.testing.assert_allclose(rj[1], rb[1])
+        assert rj[0].shape == rb[0].shape and rj[1].shape == rb[1].shape
+
+
+if __name__ == "__main__":
+    test_vmap(2, 2, 3, 4, 0, 1)
diff --git a/test/test_jax.py b/test/test_jax.py
deleted file mode 100644
index 7c53df3..0000000
--- a/test/test_jax.py
+++ /dev/null
@@ -1,366 +0,0 @@
-import ducc0
-import jax
-import numpy as np
-import pytest
-from jax.test_util import check_grads
-from numpy.testing import assert_allclose
-
-import jax_linop
-
-jax.config.update("jax_enable_x64", True)
-
-pmp = pytest.mark.parametrize
-
-r2cdict = {
-    np.dtype(np.float32): np.dtype(np.complex64),
-    np.dtype(np.float64): np.dtype(np.complex128)
-}
-
-c2rdict = {
-    np.dtype(np.complex64): np.dtype(np.float32),
-    np.dtype(np.complex128): np.dtype(np.float64)
-}
-
-
-def realtype(dtype):
-    return c2rdict[np.dtype(dtype)]
-
-
-def complextype(dtype):
-    return r2cdict[np.dtype(dtype)]
-
-
-def fht_operator(axes, nthreads):
-    def fhtfunc(inp, out, state):
-        # This function must _not_ keep any reference to 'inp' or 'out'!
-        # Also, it must not change 'inp' or 'state'.
-        ducc0.fft.genuine_fht(
-            inp, out=out, axes=state["axes"], nthreads=state["nthreads"]
-        )
-
-    def fhtfunc_abstract(shape, dtype, state):
-        return shape, dtype
-
-    return jax_linop.get_linear_call(
-        fhtfunc,
-        fhtfunc,
-        fhtfunc_abstract,
-        fhtfunc_abstract,
-        axes=tuple(axes),
-        nthreads=int(nthreads)
-    )
-
-
-def c2c_operator(axes, forward, nthreads):
-    def c2cfunc(inp, out, state):
-        ducc0.fft.c2c(
-            inp,
-            out=out,
-            axes=state["axes"],
-            nthreads=state["nthreads"],
-            forward=state["forward"]
-        )
-
-    def c2cfunc_abstract(shape, dtype, state):
-        return shape, dtype
-
-    return jax_linop.get_linear_call(
-        c2cfunc,
-        c2cfunc,  # The C2C FFT matrix is symmetric!
-        c2cfunc_abstract,
-        c2cfunc_abstract,
-        axes=tuple(axes),
-        forward=bool(forward),
-        nthreads=int(nthreads)
-    )
-
-
-def alm2realalm(alm, lmax, dtype, out=None):
-    if out is None:
-        out = np.empty((alm.shape[0], alm.shape[1] * 2 - lmax - 1), dtype=dtype)
-    out[:, 0:lmax + 1] = alm[:, 0:lmax + 1].real
-    out[:, lmax + 1:] = alm[:, lmax + 1:].view(dtype)
-    out[:, lmax + 1:] *= np.sqrt(2.)
-    return out
-
-
-def realalm2alm(alm, lmax, dtype, out=None):
-    if out is None:
-        out = np.empty((alm.shape[0], (alm.shape[1] + lmax + 1) // 2), dtype=dtype)
-    out[:, 0:lmax + 1] = alm[:, 0:lmax + 1]
-    out[:, lmax + 1:] = alm[:, lmax + 1:].view(dtype)
-    out[:, lmax + 1:] *= np.sqrt(2.) / 2
-    return out
-
-
-def sht2d_operator(lmax, mmax, ntheta, nphi, geometry, spin, nthreads):
-    def sht2dfunc(inp, out, state):
-        tmp = realalm2alm(inp, state["lmax"], complextype(inp.dtype))
-        ducc0.sht.synthesis_2d(
-            lmax=state["lmax"],
-            mmax=state["mmax"],
-            spin=state["spin"],
-            map=out,
-            alm=tmp,
-            nthreads=state["nthreads"],
-            geometry=state["geometry"]
-        )
-
-    def sht2dfunc_T(inp, out, state):
-        tmp = ducc0.sht.adjoint_synthesis_2d(
-            lmax=state["lmax"],
-            mmax=state["mmax"],
-            spin=state["spin"],
-            map=inp,
-            nthreads=state["nthreads"],
-            geometry=state["geometry"]
-        )
-        alm2realalm(tmp, state["lmax"], inp.dtype, out)
-
-    def sht2dfunc_abstract(shape_in, dtype_in, state):
-        spin = state["spin"]
-        ncomp = 1 if spin == 0 else 2
-        shape_out = (ncomp, state["ntheta"], state["nphi"])
-        return shape_out, dtype_in
-
-    def sht2dfunc_abstract_T(shape_in, dtype_in, state):
-        spin = state["spin"]
-        ncomp = 1 if spin == 0 else 2
-        lmax, mmax = state["lmax"], state["mmax"]
-        nalm = ((mmax + 1) * (mmax + 2)) // 2 + (mmax + 1) * (lmax - mmax)
-        nalm = nalm * 2 - lmax - 1
-        shape_out = (ncomp, nalm)
-        return shape_out, dtype_in
-
-    return jax_linop.get_linear_call(
-        sht2dfunc,
-        sht2dfunc_T,
-        sht2dfunc_abstract,
-        sht2dfunc_abstract_T,
-        lmax=int(lmax),
-        mmax=int(mmax),
-        spin=int(spin),
-        ntheta=int(ntheta),
-        nphi=int(nphi),
-        geometry=str(geometry),
-        nthreads=int(nthreads)
-    )
-
-def healpix_operator(lmax, mmax, nside, spin, nthreads):
-    def healpixfunc(inp, out, state):
-        tmp = realalm2alm(inp, state["lmax"], complextype(inp.dtype))
-        ducc0.sht.synthesis(
-            lmax=state["lmax"],
-            mmax=state["mmax"],
-            spin=state["spin"],
-            map=out,
-            alm=tmp,
-            nthreads=state["nthreads"],
-            **state["hpxparam"]
-        )
-
-    def healpixfunc_T(inp, out, state):
-        tmp = ducc0.sht.adjoint_synthesis(
-            lmax=state["lmax"],
-            mmax=state["mmax"],
-            spin=state["spin"],
-            map=inp,
-            nthreads=state["nthreads"],
-            **state["hpxparam"]
-        )
-        alm2realalm(tmp, state["lmax"], inp.dtype, out)
-
-    def healpixfunc_abstract(shape_in, dtype_in, state):
-        spin = state["spin"]
-        ncomp = 1 if spin == 0 else 2
-        shape_out = (ncomp, 12*state["nside"]**2)
-        return shape_out, dtype_in
-
-    def healpixfunc_abstract_T(shape_in, dtype_in, state):
-        spin = state["spin"]
-        ncomp = 1 if spin == 0 else 2
-        lmax, mmax = state["lmax"], state["mmax"]
-        nalm = ((mmax + 1) * (mmax + 2)) // 2 + (mmax + 1) * (lmax - mmax)
-        nalm = nalm * 2 - lmax - 1
-        shape_out = (ncomp, nalm)
-        return shape_out, dtype_in
-
-    nside = int(nside)
-    base = ducc0.healpix.Healpix_Base(nside, "RING")
-    return jax_linop.get_linear_call(
-        healpixfunc,
-        healpixfunc_T,
-        healpixfunc_abstract,
-        healpixfunc_abstract_T,
-        lmax=int(lmax),
-        mmax=int(mmax),
-        spin=int(spin),
-        nthreads=int(nthreads),
-        nside=nside,
-        hpxparam=base.sht_info()
-    )
-
-
-def _assert_close(a, b, epsilon):
-    assert_allclose(ducc0.misc.l2error(a, b), 0, atol=epsilon)
-
-
-@pmp("shape,axes", (((100, ), (0, )), ((10, 17), (0, 1)), ((10, 17, 3), (1, ))))
-@pmp("dtype", (np.float32, np.float64))
-@pmp("nthreads", (1, 2))
-def test_fht(shape, axes, dtype, nthreads):
-    rng = np.random.default_rng(42)
-
-    op = fht_operator(axes=axes, nthreads=nthreads)
-    a = (rng.random(shape) - 0.5).astype(dtype)
-    b1 = np.array(op(a)[0])
-    b2 = ducc0.fft.genuine_fht(a, axes=axes, nthreads=nthreads)
-    _assert_close(b1, b2, epsilon=1e-6 if dtype == np.float32 else 1e-14)
-
-    max_order = 2
-    check_grads(op, (a, ), order=max_order, modes=("fwd", "rev"), eps=1.)
-
-
-@pmp("shape,axes", (((100, ), (0, )), ((10, 17), (0, 1)), ((10, 17, 3), (1, ))))
-@pmp("forward", (False, True))
-@pmp("dtype", (np.complex64, np.complex128))
-@pmp("nthreads", (1, 2))
-def test_c2c(shape, axes, forward, dtype, nthreads):
-    rng = np.random.default_rng(42)
-
-    op = c2c_operator(axes=axes, forward=forward, nthreads=nthreads)
-    a = (rng.random(shape) -
-         0.5).astype(dtype) + (1j * (rng.random(shape) - 0.5)).astype(dtype)
-    b1 = np.array(op(a)[0])
-    b2 = ducc0.fft.c2c(a, axes=axes, forward=forward, nthreads=nthreads)
-    _assert_close(b1, b2, epsilon=1e-6 if dtype == np.complex64 else 1e-14)
-
-    max_order = 2
-    check_grads(op, (a, ), order=max_order, modes=("fwd", "rev"), eps=1.)
-
-
-def nalm(lmax, mmax):
-    return ((mmax + 1) * (mmax + 2)) // 2 + (mmax + 1) * (lmax - mmax)
-
-
-def random_alm(lmax, mmax, spin, ncomp, rng):
-    res = rng.uniform(-1., 1., (ncomp, nalm(lmax, mmax))) \
-     + 1j*rng.uniform(-1., 1., (ncomp, nalm(lmax, mmax)))
-    # make a_lm with m==0 real-valued
-    res[:, 0:lmax + 1].imag = 0.
-    ofs = 0
-    for s in range(spin):
-        res[:, ofs:ofs + spin - s] = 0.
-        ofs += lmax + 1 - s
-    return res
-
-
-@pmp("lmmax", ((10, 10), (20, 5)))
-@pmp("geometry", ("GL", "F1", "F2", "CC", "DH", "MW", "MWflip"))
-@pmp("ntheta", (20, ))
-@pmp("nphi", (30, ))
-@pmp("spin", (0, 2))
-@pmp("dtype", (np.float32, np.float64))
-@pmp("nthreads", (1, 2))
-def test_sht2d(lmmax, geometry, ntheta, nphi, spin, dtype, nthreads):
-    rng = np.random.default_rng(42)
-
-    lmax, mmax = lmmax
-    ncomp = 1 if spin == 0 else 2
-    alm0 = random_alm(lmax, mmax, spin, ncomp, rng).astype(complextype(dtype))
-    alm0r = alm2realalm(alm0, lmax, dtype)
-
-    op = sht2d_operator(
-        lmax=lmax,
-        mmax=mmax,
-        ntheta=ntheta,
-        nphi=nphi,
-        geometry=geometry,
-        spin=spin,
-        nthreads=nthreads
-    )
-    op_adj = lambda x: jax.linear_transpose(lambda y: op(y)[0], alm0r
-                                           )(x.conj())[0].conj()
-
-    map1 = np.array(op(alm0r)[0])
-    map2 = ducc0.sht.synthesis_2d(
-        alm=alm0,
-        lmax=lmax,
-        mmax=mmax,
-        spin=spin,
-        geometry=geometry,
-        ntheta=ntheta,
-        nphi=nphi,
-        nthreads=nthreads
-    )
-    _assert_close(map1, map2, epsilon=1e-6 if dtype == np.float32 else 1e-14)
-
-    map0 = (rng.random((ncomp, ntheta, nphi)) - 0.5).astype(dtype)
-    alm1r = np.array(op_adj(map0))
-    alm1 = realalm2alm(alm1r, lmax, complextype(dtype))
-    alm2 = ducc0.sht.adjoint_synthesis_2d(
-        map=map0,
-        lmax=lmax,
-        mmax=mmax,
-        spin=spin,
-        geometry=geometry,
-        nthreads=nthreads
-    )
-    _assert_close(alm1, alm2, epsilon=1e-6 if dtype == np.float32 else 1e-14)
-
-    max_order = 2
-    check_grads(op, (alm0r, ), order=max_order, modes=("fwd", "rev"), eps=1.)
-    check_grads(op_adj, (map0, ), order=max_order, modes=("fwd", "rev"), eps=1.)
-
-@pmp("lmmax", ((10, 10), (20, 5)))
-@pmp("nside", (16, 2))
-@pmp("spin", (0, 2))
-@pmp("dtype", (np.float32, np.float64))
-@pmp("nthreads", (1, 2))
-def test_healpix(lmmax, nside, spin, dtype, nthreads):
-    rng = np.random.default_rng(42)
-
-    lmax, mmax = lmmax
-    ncomp = 1 if spin == 0 else 2
-    alm0 = random_alm(lmax, mmax, spin, ncomp, rng).astype(complextype(dtype))
-    alm0r = alm2realalm(alm0, lmax, dtype)
-    base = ducc0.healpix.Healpix_Base(nside, "RING")
-    hpxparam = base.sht_info()
-
-    op = healpix_operator(
-        lmax=lmax,
-        mmax=mmax,
-        nside=nside,
-        spin=spin,
-        nthreads=nthreads
-    )
-    op_adj = lambda x: jax.linear_transpose(lambda y: op(y)[0], alm0r
-                                           )(x.conj())[0].conj()
-
-    map1 = np.array(op(alm0r)[0])
-    map2 = ducc0.sht.synthesis(
-        alm=alm0,
-        lmax=lmax,
-        mmax=mmax,
-        spin=spin,
-        nthreads=nthreads,
-        **hpxparam
-    )
-    _assert_close(map1, map2, epsilon=1e-6 if dtype == np.float32 else 1e-14)
-
-    map0 = (rng.random((ncomp, 12*nside**2)) - 0.5).astype(dtype)
-    alm1r = np.array(op_adj(map0))
-    alm1 = realalm2alm(alm1r, lmax, complextype(dtype))
-    alm2 = ducc0.sht.adjoint_synthesis(
-        map=map0,
-        lmax=lmax,
-        mmax=mmax,
-        spin=spin,
-        nthreads=nthreads,
-        **hpxparam
-    )
-    _assert_close(alm1, alm2, epsilon=1e-6 if dtype == np.float32 else 1e-14)
-
-    max_order = 2
-    check_grads(op, (alm0r, ), order=max_order, modes=("fwd", "rev"), eps=1.)
-    check_grads(op_adj, (map0, ), order=max_order, modes=("fwd", "rev"), eps=1.)
diff --git a/test/test_lin.py b/test/test_lin.py
new file mode 100644
index 0000000..2a0b402
--- /dev/null
+++ b/test/test_lin.py
@@ -0,0 +1,421 @@
+# SPDX-License-Identifier: BSD-2-Clause
+# Authors: Martin Reinecke, Jakob Roth, Gordian Edenhofer
+
+# Copyright(C) 2024 Max-Planck-Society
+
+from functools import partial
+
+have_ducc0 = True
+try:
+    import ducc0
+except ImportError:
+    have_ducc0 = False
+
+import jax
+import scipy
+import numpy as np
+import pytest
+from jax.test_util import check_grads
+from numpy.testing import assert_allclose
+
+import jaxbind
+
+pmp = pytest.mark.parametrize
+
+jax.config.update("jax_enable_x64", True)
+
+r2cdict = {
+    np.dtype(np.float32): np.dtype(np.complex64),
+    np.dtype(np.float64): np.dtype(np.complex128),
+}
+
+c2rdict = {
+    np.dtype(np.complex64): np.dtype(np.float32),
+    np.dtype(np.complex128): np.dtype(np.float64),
+}
+
+
+def realtype(dtype):
+    return c2rdict[np.dtype(dtype)]
+
+
+def complextype(dtype):
+    return r2cdict[np.dtype(dtype)]
+
+
+def fhtfunc(out, args, kwargs_dump):
+    kwargs = jaxbind.load_kwargs(kwargs_dump)
+    # This function must _not_ keep any reference to 'inp' or 'args'!
+    # Also, it must not change 'args'.
+    if have_ducc0:
+        ducc0.fft.genuine_fht(args[0], out=out[0], **kwargs)
+    else:
+        tmp = scipy.fft.fftn(args[0], axes=kwargs["axes"])
+        out[0][()] = tmp.real - tmp.imag
+
+
+def fhtfunc_abstract(*args, **kwargs):
+    (x,) = args
+    return ((x.shape, x.dtype),)
+
+
+def c2cfunc(out, args, kwargs_dump):
+    kwargs = jaxbind.load_kwargs(kwargs_dump)
+    (x,) = args
+    if have_ducc0:
+        ducc0.fft.c2c(x, out=out[0], **kwargs)
+    else:
+        func = scipy.fft.fftn if kwargs["forward"] else scipy.fft.ifftn
+        norm = "backward" if kwargs["forward"] else "forward"
+        out[0][()] = func(x, norm=norm, axes=kwargs["axes"])
+
+
+def c2cfunc_abstract(*args, **kwargs):
+    (x,) = args
+    return ((x.shape, x.dtype),)
+
+
+def alm2realalm(alm, lmax, dtype, out=None):
+    if out is None:
+        out = np.empty((alm.shape[0], alm.shape[1] * 2 - lmax - 1), dtype=dtype)
+    out[:, 0 : lmax + 1] = alm[:, 0 : lmax + 1].real
+    out[:, lmax + 1 :] = alm[:, lmax + 1 :].view(dtype)
+    out[:, lmax + 1 :] *= np.sqrt(2.0)
+    return out
+
+
+def realalm2alm(alm, lmax, dtype, out=None):
+    if out is None:
+        out = np.empty((alm.shape[0], (alm.shape[1] + lmax + 1) // 2), dtype=dtype)
+    out[:, 0 : lmax + 1] = alm[:, 0 : lmax + 1]
+    out[:, lmax + 1 :] = alm[:, lmax + 1 :].view(dtype)
+    out[:, lmax + 1 :] *= np.sqrt(2.0) / 2
+    return out
+
+
+def sht2d_operator(lmax, mmax, ntheta, nphi, geometry, spin, nthreads):
+    def sht2dfunc(out, args, kwargs_dump):
+        inp = args[0]
+        state = jaxbind.load_kwargs(kwargs_dump)
+        tmp = realalm2alm(inp, state["lmax"], complextype(inp.dtype))
+        ducc0.sht.synthesis_2d(
+            lmax=state["lmax"],
+            mmax=state["mmax"],
+            spin=state["spin"],
+            map=out[0],
+            alm=tmp,
+            nthreads=state["nthreads"],
+            geometry=state["geometry"],
+            ntheta=state["ntheta"],
+            nphi=state["nphi"],
+        )
+
+    def sht2dfunc_T(out, args, kwargs_dump):
+        inp = args[0]
+        state = jaxbind.load_kwargs(kwargs_dump)
+        tmp = ducc0.sht.adjoint_synthesis_2d(
+            lmax=state["lmax"],
+            mmax=state["mmax"],
+            spin=state["spin"],
+            map=inp,
+            nthreads=state["nthreads"],
+            geometry=state["geometry"],
+        )
+        alm2realalm(tmp, state["lmax"], inp.dtype, out[0])
+
+    def sht2dfunc_abstract(*args, **kwargs):
+        spin = kwargs["spin"]
+        ncomp = 1 if spin == 0 else 2
+        shape_out = (ncomp, kwargs["ntheta"], kwargs["nphi"])
+        return ((shape_out, args[0].dtype),)
+
+    def sht2dfunc_abstract_T(*args, **kwargs):
+        spin = kwargs["spin"]
+        ncomp = 1 if spin == 0 else 2
+        lmax, mmax = kwargs["lmax"], kwargs["mmax"]
+        nalm = ((mmax + 1) * (mmax + 2)) // 2 + (mmax + 1) * (lmax - mmax)
+        nalm = nalm * 2 - lmax - 1
+        shape_out = (ncomp, nalm)
+        return ((shape_out, args[0].dtype),)
+
+    func = jaxbind.get_linear_call(
+        sht2dfunc,
+        sht2dfunc_T,
+        sht2dfunc_abstract,
+        sht2dfunc_abstract_T,
+    )
+
+    return partial(
+        func,
+        geometry=geometry,
+        ntheta=ntheta,
+        nphi=nphi,
+        lmax=lmax,
+        mmax=mmax,
+        spin=spin,
+        nthreads=nthreads,
+    )
+
+
+def healpixfunc(out, args, kwargs_dump):
+    kwargs = jaxbind.load_kwargs(kwargs_dump).copy()
+    theta, phi0, nphi, ringstart, x = args
+    tmp = realalm2alm(x, kwargs["lmax"], complextype(x.dtype))
+    ducc0.sht.synthesis(
+        map=out[0],
+        alm=tmp,
+        theta=theta,
+        phi0=phi0,
+        nphi=nphi,
+        ringstart=ringstart,
+        spin=kwargs["spin"],
+        lmax=kwargs["lmax"],
+        mmax=kwargs["mmax"],
+        nthreads=kwargs["nthreads"],
+    )
+
+
+def healpixfunc_T(out, args, kwargs_dump):
+    kwargs = jaxbind.load_kwargs(kwargs_dump).copy()
+    theta, phi0, nphi, ringstart, x = args
+    tmp = ducc0.sht.adjoint_synthesis(
+        map=x,
+        theta=theta,
+        phi0=phi0,
+        nphi=nphi,
+        ringstart=ringstart,
+        spin=kwargs["spin"],
+        lmax=kwargs["lmax"],
+        mmax=kwargs["mmax"],
+        nthreads=kwargs["nthreads"],
+    )
+    alm2realalm(tmp, kwargs["lmax"], x.dtype, out[0])
+
+
+def healpixfunc_abstract(*args, **kwargs):
+    _, _, _, _, x = args
+    spin = kwargs["spin"]
+    ncomp = 1 if spin == 0 else 2
+    shape_out = (ncomp, 12 * kwargs["nside"] ** 2)
+    return ((shape_out, x.dtype),)
+
+
+def healpixfunc_abstract_T(*args, **kwargs):
+    _, _, _, _, x = args
+    spin = kwargs["spin"]
+    ncomp = 1 if spin == 0 else 2
+    lmax, mmax = kwargs["lmax"], kwargs["mmax"]
+    nalm = ((mmax + 1) * (mmax + 2)) // 2 + (mmax + 1) * (lmax - mmax)
+    nalm = nalm * 2 - lmax - 1
+    shape_out = (ncomp, nalm)
+    return ((shape_out, x.dtype),)
+
+
+def _assert_close(a, b, epsilon):
+    if have_ducc0:
+        assert_allclose(ducc0.misc.l2error(a, b), 0, atol=epsilon)
+    else:
+        assert_allclose(
+            scipy.linalg.norm(a - b) / scipy.linalg.norm(a), 0, atol=epsilon
+        )
+
+
+@pmp("shape,axes", (((100,), (0,)), ((10, 17), (0, 1)), ((10, 17, 3), (1,))))
+@pmp("dtype", (np.float32, np.float64))
+@pmp("nthreads", (1, 2))
+def test_fht(shape, axes, dtype, nthreads):
+    rng = np.random.default_rng(42)
+
+    fht = jaxbind.get_linear_call(
+        fhtfunc, fhtfunc, fhtfunc_abstract, fhtfunc_abstract
+    )
+    kw = dict(axes=axes, nthreads=nthreads)
+
+    a = (rng.random(shape) - 0.5).astype(dtype)
+    b1 = np.array(fht(a, **kw)[0])
+    if have_ducc0:
+        b2 = ducc0.fft.genuine_fht(a, **kw)
+    else:
+        b2 = scipy.fft.fftn(a, axes=kw["axes"])
+        b2 = b2.real - b2.imag
+
+    _assert_close(b1, b2, epsilon=1e-6 if dtype == np.float32 else 1e-14)
+
+    max_order = 2
+    check_grads(
+        partial(fht, **kw), (a,), order=max_order, modes=("fwd", "rev"), eps=1.0
+    )
+
+
+@pmp("shape,axes", (((100,), (0,)), ((10, 17), (0, 1)), ((10, 17, 3), (1,))))
+@pmp("forward", (False, True))
+@pmp("dtype", (np.complex64, np.complex128))
+@pmp("nthreads", (1, 2))
+def test_c2c(shape, axes, forward, dtype, nthreads):
+    rng = np.random.default_rng(42)
+
+    # The C2C FFT matrix is symmetric!
+    c2c = jaxbind.get_linear_call(
+        c2cfunc, c2cfunc, c2cfunc_abstract, c2cfunc_abstract
+    )
+    kw = dict(axes=axes, forward=forward, nthreads=nthreads)
+
+    a = (rng.random(shape) - 0.5).astype(dtype) + (
+        1j * (rng.random(shape) - 0.5)
+    ).astype(dtype)
+    b1 = np.array(c2c(a, **kw)[0])
+    if have_ducc0:
+        b2 = ducc0.fft.c2c(a, **kw)
+    else:
+        if forward:
+            b2 = scipy.fft.fftn(a, axes=kw["axes"])
+        else:
+            b2 = scipy.fft.ifftn(a, norm="forward", axes=kw["axes"])
+    _assert_close(b1, b2, epsilon=1e-6 if dtype == np.complex64 else 1e-14)
+
+    max_order = 2
+    check_grads(
+        partial(c2c, **kw), (a,), order=max_order, modes=("fwd", "rev"), eps=1.0
+    )
+
+
+def nalm(lmax, mmax):
+    return ((mmax + 1) * (mmax + 2)) // 2 + (mmax + 1) * (lmax - mmax)
+
+
+def random_alm(lmax, mmax, spin, ncomp, rng):
+    res = rng.uniform(-1.0, 1.0, (ncomp, nalm(lmax, mmax))) + 1j * rng.uniform(
+        -1.0, 1.0, (ncomp, nalm(lmax, mmax))
+    )
+    # make a_lm with m==0 real-valued
+    res[:, 0 : lmax + 1].imag = 0.0
+    ofs = 0
+    for s in range(spin):
+        res[:, ofs : ofs + spin - s] = 0.0
+        ofs += lmax + 1 - s
+    return res
+
+
+@pmp("lmmax", ((10, 10), (20, 5)))
+@pmp("geometry", ("GL", "F1", "F2", "CC", "DH", "MW", "MWflip"))
+@pmp("ntheta", (20,))
+@pmp("nphi", (30,))
+@pmp("spin", (0, 2))
+@pmp("dtype", (np.float32, np.float64))
+@pmp("nthreads", (1, 2))
+def test_sht2d(lmmax, geometry, ntheta, nphi, spin, dtype, nthreads):
+    if not have_ducc0:
+        pytest.skip()
+    rng = np.random.default_rng(42)
+
+    lmax, mmax = lmmax
+    ncomp = 1 if spin == 0 else 2
+    alm0 = random_alm(lmax, mmax, spin, ncomp, rng).astype(complextype(dtype))
+    alm0r = alm2realalm(alm0, lmax, dtype)
+
+    op = sht2d_operator(
+        lmax=lmax,
+        mmax=mmax,
+        ntheta=ntheta,
+        nphi=nphi,
+        geometry=geometry,
+        spin=spin,
+        nthreads=nthreads,
+    )
+    # The conjugations are only necessary if the input or output data types
+    # are complex. Leaving them out makes things (surprisingly) faster.
+    #    op_adj = lambda x: jax.linear_transpose(lambda y: op(y)[0], alm0r)(x.conj())[
+    #        0
+    #    ].conj()
+    op_adj = lambda x: jax.linear_transpose(lambda y: op(y)[0], alm0r)(x)[0]
+
+    map1 = np.array(op(alm0r)[0])
+    map2 = ducc0.sht.synthesis_2d(
+        alm=alm0,
+        lmax=lmax,
+        mmax=mmax,
+        spin=spin,
+        geometry=geometry,
+        ntheta=ntheta,
+        nphi=nphi,
+        nthreads=nthreads,
+    )
+    _assert_close(map1, map2, epsilon=1e-6 if dtype == np.float32 else 1e-14)
+
+    map0 = (rng.random((ncomp, ntheta, nphi)) - 0.5).astype(dtype)
+    alm1r = np.array(op_adj(map0))
+    alm1 = realalm2alm(alm1r, lmax, complextype(dtype))
+    alm2 = ducc0.sht.adjoint_synthesis_2d(
+        map=map0, lmax=lmax, mmax=mmax, spin=spin, geometry=geometry, nthreads=nthreads
+    )
+    _assert_close(alm1, alm2, epsilon=1e-6 if dtype == np.float32 else 1e-14)
+
+    max_order = 2
+    check_grads(op, (alm0r,), order=max_order, modes=("fwd", "rev"), eps=1.0)
+    check_grads(op_adj, (map0,), order=max_order, modes=("fwd", "rev"), eps=1.0)
+
+
+@pmp("lmmax", ((10, 10), (20, 5)))
+@pmp("nside", (16, 2))
+@pmp("spin", (0, 2))
+@pmp("dtype", (np.float32, np.float64))
+@pmp("nthreads", (1, 2))
+def test_healpix(lmmax, nside, spin, dtype, nthreads):
+    if not have_ducc0:
+        pytest.skip()
+    rng = np.random.default_rng(42)
+
+    lmax, mmax = lmmax
+    ncomp = 1 if spin == 0 else 2
+    alm0 = random_alm(lmax, mmax, spin, ncomp, rng).astype(complextype(dtype))
+    alm0r = alm2realalm(alm0, lmax, dtype)
+    base = ducc0.healpix.Healpix_Base(nside, "RING")
+    hpxparam = base.sht_info()
+
+    hp = jaxbind.get_linear_call(
+        healpixfunc,
+        healpixfunc_T,
+        healpixfunc_abstract,
+        healpixfunc_abstract_T,
+        first_n_args_fixed=4,
+    )
+
+    def hpp(x):
+        # Partial insert where the first parameter is not inserted
+        return hp(
+            hpxparam["theta"],
+            hpxparam["phi0"],
+            hpxparam["nphi"],
+            hpxparam["ringstart"],
+            x,
+            lmax=lmax,
+            mmax=mmax,
+            spin=spin,
+            nthreads=nthreads,
+            nside=nside,
+        )
+
+    # The conjugations are only necessary if the input or output data types
+    # are complex. Leaving them out makes things (surprisingly) faster.
+    #    hp_adj = lambda x: jax.linear_transpose(lambda y: hp(y)[0], alm0r)(x.conj())[
+    #        0
+    #    ].conj()
+    hpp_adj = lambda x: jax.linear_transpose(lambda y: hpp(y)[0], alm0r)(x)[0]
+
+    map1 = np.array(hpp(alm0r)[0])
+    map2 = ducc0.sht.synthesis(
+        alm=alm0, lmax=lmax, mmax=mmax, spin=spin, nthreads=nthreads, **hpxparam
+    )
+    _assert_close(map1, map2, epsilon=1e-6 if dtype == np.float32 else 1e-14)
+
+    map0 = (rng.random((ncomp, 12 * nside**2)) - 0.5).astype(dtype)
+    alm1r = np.array(hpp_adj(map0))
+    alm1 = realalm2alm(alm1r, lmax, complextype(dtype))
+    alm2 = ducc0.sht.adjoint_synthesis(
+        map=map0, lmax=lmax, mmax=mmax, spin=spin, nthreads=nthreads, **hpxparam
+    )
+    _assert_close(alm1, alm2, epsilon=1e-6 if dtype == np.float32 else 1e-14)
+
+    max_order = 2
+    check_grads(hpp, (alm0r,), order=max_order, modes=("fwd", "rev"), eps=1.0)
+    check_grads(hpp_adj, (map0,), order=max_order, modes=("fwd", "rev"), eps=1.0)
diff --git a/test/test_mlin.py b/test/test_mlin.py
new file mode 100644
index 0000000..76e23ac
--- /dev/null
+++ b/test/test_mlin.py
@@ -0,0 +1,133 @@
+# SPDX-License-Identifier: BSD-2-Clause
+# Authors: Martin Reinecke, Jakob Roth, Gordian Edenhofer
+
+# Copyright(C) 2024 Max-Planck-Society
+
+import jax
+import numpy as np
+from jax import numpy as jnp
+from jax.test_util import check_grads
+
+import jaxbind
+
+from functools import partial
+
+jax.config.update("jax_enable_x64", True)
+
+
+def test_mlin():
+    def mlin(out, args, kwargs_dump):
+        x, y = args[0], args[1]
+        out[0][()] = x * y
+        out[1][()] = x * y
+
+    def mlin_T1(out, args, kwargs_dump):
+        y, da, db = args[0], args[1], args[2]
+        out[0][()] = y * da + y * db
+
+    def mlin_T2(out, args, kwargs_dump):
+        x, da, db = args[0], args[1], args[2]
+        out[0][()] = x * da + x * db
+
+    def mlin_abstract(*args, **kwargs):
+        assert args[0].shape == args[1].shape
+        return (
+            (args[0].shape, args[0].dtype, None),
+            (args[0].shape, args[0].dtype, None),
+        )
+
+    def mlin_abstract_T1(*args, **kwargs):
+        assert args[0].shape == args[1].shape
+        return ((args[0].shape, args[0].dtype, None),)
+
+    def mlin_abstract_T2(*args, **kwargs):
+        out_axes = kwargs.pop("batch_axes", ())
+        assert args[0].shape == args[1].shape
+        return ((args[0].shape, args[0].dtype, out_axes),)
+
+    func_T = (mlin_T1, mlin_T2)
+    func_abstract_T = (mlin_abstract_T1, mlin_abstract_T2)
+    mlin_jax = jaxbind.get_linear_call(
+        mlin,
+        func_T,
+        mlin_abstract,
+        func_abstract_T,
+    )
+
+    inp = (4 + jnp.zeros((2, 2)), 1 + jnp.zeros((2, 2)))
+    check_grads(partial(mlin_jax, axes=(3, 4)), inp, order=2, modes=["fwd"], eps=1.0)
+    check_grads(partial(mlin_jax, axes=(3, 4)), inp, order=1, modes=["rev"], eps=1.0)
+
+    inp2 = (7 + jnp.zeros((2, 2)), -3 + jnp.zeros((2, 2)))
+    inp3 = (10 + jnp.zeros((2, 2)), 5 + jnp.zeros((2, 2)))
+
+    def mlin_purejax(x, y):
+        return [x * y, x * y]
+
+    primals_jax, njf_vjp = jax.vjp(mlin_purejax, *inp)
+    res_vjp_jax = njf_vjp(mlin_purejax(*inp))
+    res_jvp_jax = jax.jvp(mlin_purejax, inp2, inp3)
+
+    primals, f_vjp = jax.vjp(mlin_jax, *inp)
+    res_vjp = f_vjp(mlin_jax(*inp))
+    res_jvp = jax.jvp(mlin_jax, inp2, inp3)
+
+    np.testing.assert_allclose(primals, primals_jax)
+    np.testing.assert_allclose(res_vjp, res_vjp_jax)
+    np.testing.assert_allclose(res_jvp, res_jvp_jax)
+
+
+def test_mlin_fixed_arg():
+    def mlin(out, args, kwargs_dump):
+        x, y = args[0], args[1]
+        out[0][()] = x * x * y
+        out[1][()] = x * y
+
+    def mlin_T2(out, args, kwargs_dump):
+        x, da, db = args[0], args[1], args[2]
+        out[0][()] = x * x * da + x * db
+
+    def mlin_abstract(*args, **kwargs):
+        assert args[0].shape == args[1].shape
+        return (
+            (args[0].shape, args[0].dtype, None),
+            (args[0].shape, args[0].dtype, None),
+        )
+
+    def mlin_abstract_T2(*args, **kwargs):
+        assert args[0].shape == args[1].shape
+        return ((args[0].shape, args[0].dtype, None),)
+
+    func_T = (None, mlin_T2)
+    func_abstract_T = (None, mlin_abstract_T2)
+    mlin_jax = jaxbind.get_linear_call(
+        mlin,
+        func_T,
+        mlin_abstract,
+        func_abstract_T,
+        first_n_args_fixed=1,
+    )
+
+    inp1 = 4 + jnp.zeros((2, 2))
+    inp2 = 1 + jnp.zeros((2, 2))
+
+    def mlin_purejax(x, y):
+        return [x * y * x, x * y]
+
+    mlin_jax_pt = partial(mlin_jax, inp1, axes=(3, 4))
+    mlin_purejax_pt = partial(mlin_purejax, inp1)
+
+    check_grads(mlin_jax_pt, (inp2,), order=2, modes=["fwd"], eps=1.0)
+    check_grads(mlin_jax_pt, (inp2,), order=1, modes=["rev"], eps=1.0)
+
+    primals, f_vjp = jax.vjp(mlin_jax_pt, inp2)
+    res_vjp = f_vjp(mlin_jax_pt(inp2))
+    res_jvp = jax.jvp(mlin_jax_pt, (inp1,), (inp2,))
+
+    primals_jax, njf_vjp = jax.vjp(mlin_purejax_pt, inp2)
+    res_vjp_jax = njf_vjp(mlin_purejax_pt(inp2))
+    res_jvp_jax = jax.jvp(mlin_purejax_pt, (inp1,), (inp2,))
+
+    np.testing.assert_allclose(primals, primals_jax)
+    np.testing.assert_allclose(res_vjp, res_vjp_jax)
+    np.testing.assert_allclose(res_jvp, res_jvp_jax)
diff --git a/test/test_nonlin.py b/test/test_nonlin.py
new file mode 100644
index 0000000..c5eff74
--- /dev/null
+++ b/test/test_nonlin.py
@@ -0,0 +1,240 @@
+# SPDX-License-Identifier: BSD-2-Clause
+# Authors: Martin Reinecke, Jakob Roth, Gordian Edenhofer
+
+# Copyright(C) 2024 Max-Planck-Society
+
+from functools import partial
+
+import jax
+import numpy as np
+import pytest
+from jax import random
+from jax.test_util import check_grads
+
+import jaxbind
+
+pmp = pytest.mark.parametrize
+
+jax.config.update("jax_enable_x64", True)
+
+precision_dict = {
+    np.dtype(np.complex64): np.dtype(np.float32),
+    np.dtype(np.complex128): np.dtype(np.float64),
+    np.dtype(np.float32): np.dtype(np.float32),
+    np.dtype(np.float64): np.dtype(np.float64),
+}
+
+
+def real_type_same_precision(dtype):
+    return precision_dict[np.dtype(dtype)]
+
+
+def is_complex_type(dtype):
+    if dtype is np.float64 or dtype is np.float32:
+        return False
+    else:
+        return True
+
+
+def f_call(x, y):
+    return x * y, y * y
+
+
+def f(out, args, kwargs_dump):
+    out[0][()], out[1][()] = f_call(*args)
+
+
+def f_jvp_call(x, y, dx, dy):
+    return y * dx + x * dy, 2 * y * dy
+
+
+def f_jvp(out, args, kwargs_dump):
+    out[0][()], out[1][()] = f_jvp_call(*args)
+
+
+def f_vjp_call(x, y, da, db):
+    return y * da, x * da + 2 * y * db
+
+
+def f_vjp(out, args, kwargs_dump):
+    out[0][()], out[1][()] = f_vjp_call(*args)
+
+
+def f_jvp_fix_x_call(x, y, dy):
+    return x * dy, 2 * y * dy
+
+
+def f_jvp_fix_x(out, args, kwargs_dump):
+    out[0][()], out[1][()] = f_jvp_fix_x_call(*args)
+
+
+def f_vjp_fix_x_call(x, y, da, db):
+    return x * da + 2 * y * db
+
+
+def f_vjp_fix_x(out, args, kwargs_dump):
+    out[0][()] = f_vjp_fix_x_call(*args)
+
+
+def f_abstract(*args, **kwargs):
+    out = jax.eval_shape(f_call, *args[0:2])
+    return tuple((o.shape, o.dtype, 0) for o in out)
+
+
+def f_abstract_vjp_fix_x(*args, **kwargs):
+    out = jax.eval_shape(f_vjp_fix_x_call, *args[0:4])
+    return ((out.shape, out.dtype, ()),)
+
+
+@pmp("first_n_args_fixed", (0, 1))
+@pmp("dtype", (np.float32, np.float64, np.complex64, np.complex128))
+@pmp("shape", ((2,), (3, 4)))
+@pmp("seed", (42,))
+@pmp("jit", (False, True))
+def test_derivatives(first_n_args_fixed, dtype, shape, seed, jit):
+    if first_n_args_fixed == 0:
+        funcs_deriv = (
+            f_jvp,
+            f_vjp,
+        )
+        absr_T = f_abstract
+    else:
+        funcs_deriv = (
+            f_jvp_fix_x,
+            f_vjp_fix_x,
+        )
+        absr_T = f_abstract_vjp_fix_x
+
+    f_jax = jaxbind.get_nonlinear_call(
+        f,
+        funcs_deriv,
+        f_abstract,
+        absr_T,
+        first_n_args_fixed=first_n_args_fixed,
+    )
+    if jit:
+        f_jax = jax.jit(f_jax)
+    key = random.PRNGKey(seed)
+    key, subkey = random.split(key)
+    inp1 = jax.random.uniform(
+        subkey, shape=shape, dtype=real_type_same_precision(dtype)
+    )
+    if is_complex_type(dtype):
+        key, subkey = random.split(key)
+        inp1 = inp1 + jax.random.uniform(
+            subkey, shape=shape, dtype=real_type_same_precision(dtype)
+        )
+
+    key, subkey = random.split(key)
+    inp2 = jax.random.uniform(
+        subkey, shape=shape, dtype=real_type_same_precision(dtype)
+    )
+    if is_complex_type(dtype):
+        key, subkey = random.split(key)
+        inp2 = inp2 + jax.random.uniform(
+            subkey, shape=shape, dtype=real_type_same_precision(dtype)
+        )
+
+    if first_n_args_fixed == 0:
+        inp = (inp1, inp2)
+    else:
+        f_jax = partial(f_jax, inp1)
+        inp = (inp2,)
+
+    check_grads(f_jax, inp, order=1, modes=["fwd", "rev"], eps=1e-3)
+
+
+@pmp("in_axes", ((0, 0), (0, None), (None, 0)))
+@pmp("dtype", (np.float32, np.float64, np.complex64, np.complex128))
+@pmp("shape", ((2,), (3, 4)))
+@pmp("seed", (42,))
+@pmp("jit", (False, True))
+def test_vmap(in_axes, dtype, shape, seed, jit):
+    funcs_deriv = (
+        f_jvp,
+        f_vjp,
+    )
+    absr_T = f_abstract
+
+    f_jax = jaxbind.get_nonlinear_call(
+        f,
+        funcs_deriv,
+        f_abstract,
+        absr_T,
+    )
+    if jit:
+        f_jax = jax.jit(f_jax)
+    key = random.PRNGKey(seed)
+
+    key, subkey = random.split(key)
+    inp1 = jax.random.uniform(
+        subkey, shape=shape, dtype=real_type_same_precision(dtype)
+    )
+    if is_complex_type(dtype):
+        key, subkey = random.split(key)
+        inp1 = inp1 + jax.random.uniform(
+            subkey, shape=shape, dtype=real_type_same_precision(dtype)
+        )
+
+    key, subkey = random.split(key)
+    inp2 = jax.random.uniform(
+        subkey, shape=shape, dtype=real_type_same_precision(dtype)
+    )
+    if is_complex_type(dtype):
+        key, subkey = random.split(key)
+        inp2 = inp2 + jax.random.uniform(
+            subkey, shape=shape, dtype=real_type_same_precision(dtype)
+        )
+
+    inp = (inp1, inp2)
+
+    # test consistency against JAX native function
+    res1 = f_jax(*inp)
+    res2 = f_call(*inp)
+    np.testing.assert_allclose(res1, res2)
+
+    if in_axes[0] == 0:
+        i1_shape = (10,) + shape
+    else:
+        i1_shape = shape
+    if in_axes[1] == 0:
+        i2_shape = (10,) + shape
+    else:
+        i2_shape = shape
+
+    key, subkey = random.split(key)
+    inp1 = jax.random.uniform(
+        subkey, shape=i1_shape, dtype=real_type_same_precision(dtype)
+    )
+    if is_complex_type(dtype):
+        key, subkey = random.split(key)
+        inp1 = inp1 + jax.random.uniform(
+            subkey, shape=i1_shape, dtype=real_type_same_precision(dtype)
+        )
+
+    key, subkey = random.split(key)
+    inp2 = jax.random.uniform(
+        subkey, shape=i2_shape, dtype=real_type_same_precision(dtype)
+    )
+    if is_complex_type(dtype):
+        key, subkey = random.split(key)
+        inp2 = inp2 + jax.random.uniform(
+            subkey, shape=i2_shape, dtype=real_type_same_precision(dtype)
+        )
+
+    inp = (inp1, inp2)
+
+    f_jax_vmap = jax.vmap(f_jax, in_axes=in_axes)
+    f_call_vmap = jax.vmap(f_call, in_axes=in_axes)
+
+    if jit:
+        f_jax_vmap = jax.jit(f_jax_vmap)
+        f_call_vmap = jax.jit(f_call_vmap)
+
+    # test consistency against JAX native function with vmap
+    res1 = f_jax_vmap(*inp)
+    res2 = f_call_vmap(*inp)
+    np.testing.assert_allclose(res1, res2)
+
+    # test derivatives of vmap function
+    check_grads(f_jax_vmap, inp, order=1, modes=["fwd", "rev"], eps=1e-3)
diff --git a/test/test_reference.py b/test/test_reference.py
new file mode 100644
index 0000000..1c86442
--- /dev/null
+++ b/test/test_reference.py
@@ -0,0 +1,47 @@
+# SPDX-License-Identifier: BSD-2-Clause
+# Authors: Martin Reinecke, Jakob Roth, Gordian Edenhofer
+
+# Copyright(C) 2024 Max-Planck-Society
+
+# %%
+import jax
+from jax import numpy as jnp
+
+import jaxbind
+
+jax.config.update("jax_enable_x64", True)
+
+def test_reference():
+    def lin(out, args, kwargs_dump):
+        x, y = args
+        out[0][()] = x + y
+        out[1][()] = x + y
+
+
+    def lin_T(out, args, kwargs_dump):
+        a, b = args
+        out[0][()] = a + b
+        out[1][()] = a + b
+
+
+    def lin_abstract(*args, **kwargs):
+        x, y = args
+        assert x.shape == y.shape
+        return ((x.shape, x.dtype), (x.shape, x.dtype))
+
+
+    lin_jax = jaxbind.get_linear_call(lin, lin_T, lin_abstract, lin_abstract)
+    inp = (4 + jnp.zeros((2, 2)), 1 + jnp.zeros((2, 2)))
+    lin_jax = jax.jit(lin_jax, static_argnames=("axes",))
+
+    del lin, lin_abstract, lin_T
+
+
+    def lin(*a):
+        raise ValueError()
+
+
+    _ = {"a": jnp.arange(10)}
+
+    lin_jax(*inp, axes=(3, 4))
+    jax.linear_transpose(lin_jax, *inp)(list(inp))