entangle-jax 0.1.0__tar.gz

entangle_jax-0.1.0/.github/workflows/publish.yml

@@ -0,0 +1,44 @@
+name: Publish
+on:
+  release:
+    types: [published]
+jobs:
+  build_wheels:
+    strategy:
+      fail-fast: false
+      matrix:
+        os: [ubuntu-latest, windows-latest, macos-latest]
+    runs-on: ${{ matrix.os }}
+    steps:
+      - uses: actions/checkout@v7
+      - uses: pypa/cibuildwheel@v4.1.0
+      - uses: actions/upload-artifact@v4
+        with:
+          name: wheels-${{ matrix.os }}
+          path: wheelhouse/*.whl
+
+  build_sdist:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v7
+      - uses: astral-sh/setup-uv@v7
+      - run: uv sync --group dev
+      - run: uv build --sdist
+      - uses: actions/upload-artifact@v4
+        with:
+          name: sdist
+          path: dist/*.tar.gz
+
+  publish:
+    needs: [build_wheels, build_sdist]
+    runs-on: ubuntu-latest
+    environment: pypi
+    permissions:
+      id-token: write
+    steps:
+      - uses: actions/download-artifact@v4
+        with:
+          pattern: "*"
+          path: dist/
+          merge-multiple: true
+      - uses: pypa/gh-action-pypi-publish@release/v1

entangle_jax-0.1.0/.github/workflows/tests.yml

@@ -0,0 +1,27 @@
+name: Tests
+on:
+  pull_request:
+  push:
+    branches:
+      - main
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v7
+      - uses: astral-sh/setup-uv@v7
+      - run: uv sync --group dev
+      - run: uv run pytest
+      - run: uv run ruff check .
+      - run: uv run ruff format --check .
+      - run: uv run ty check
+
+  build_wheels:
+    strategy:
+      fail-fast: false
+      matrix:
+        os: [ubuntu-latest, windows-latest, macos-latest]
+    runs-on: ${{ matrix.os }}
+    steps:
+      - uses: actions/checkout@v7
+      - uses: pypa/cibuildwheel@v4.1.0

entangle_jax-0.1.0/.gitignore

@@ -0,0 +1,10 @@
+.venv/
+.probe/
+__pycache__/
+*.pyc
+build/
+dist/
+*.egg-info/
+.pytest_cache/
+.ruff_cache/
+site/

entangle_jax-0.1.0/LICENSE

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

entangle_jax-0.1.0/PKG-INFO

@@ -0,0 +1,108 @@
+Metadata-Version: 2.4
+Name: entangle-jax
+Version: 0.1.0
+Summary: A generic JAX primitive that is opaque to XLA reordering optimizations
+Author-Email: Nardi Lam <mail@nardilam.nl>
+License-Expression: MIT
+Project-URL: Homepage, https://github.com/nardi/entangle-jax
+Requires-Python: >=3.11
+Requires-Dist: jax>=0.10
+Requires-Dist: jaxtyping>=0.2
+Description-Content-Type: text/markdown
+
+# entangle-jax
+
+A generic JAX primitive that is opaque to XLA reordering optimizations.
+
+`entangle(payload, *witnesses)` returns `payload` unchanged, but ordered under
+`jit` after every witness's producer. Useful for sequencing native
+side-effecting calls that XLA would otherwise be free to reorder, since they
+share no ordinary data dependency.
+
+Install with `uv add entangle-jax` or `pip install entangle-jax`.
+
+## Why
+
+When interacting with external native libraries, it is sometimes necessary to
+interact with objects they manage through opaque pointers. This conflicts with
+the XLA memory model where everything is a static-sized buffer that can be
+arbitrarily copied or reused when needed. One way to achieve this interaction is
+by creating a token that XLA can pass around that references this object, but
+XLA is not aware that e.g. multiple copies of this token in fact reference the
+same memory, and so will perform optimizations that are at worst unsafe and at
+best inefficient.
+
+Unfortunately, there is no way to tell XLA that a buffer will be modified
+in-place and that it should maintain consistent order of operations involving
+this buffer. Marking a primitive as side-effecting will stop XLA from removing
+it as dead code, but it will still assume that the input and output buffers are
+distinct, independent objects. At most, it is possible to create an in-place
+modifying primitive by using [input-output aliasing](https://openxla.org/xla/aliasing), in which case the input and
+output buffers are the same, but then XLA will simply copy the input buffer if
+it is used by another function to ensure both calls are safe.
+
+To work around this, we can create an artificial data dependency with
+`entangle`. The contents of a custom call are not visible to the XLA optimizer, so
+it has to treat the output `payload` as a new distinct variable from the input
+`payload` that depends on all `witness` values. A standard-library alternative
+would be `jax.lax.optimization_barrier`, but this is currently 
+[unreliable on CPU without a non-default XLA flag](https://github.com/openxla/xla/issues/20440).
+
+The actual runtime behavior of the primitive is a no-op: it uses aliasing to tell XLA that the input and output `payload` should be the same buffer, so it does literally nothing. This also means it is supported on every platform. The usual pattern is to overwrite the `payload` variable so that its buffer can always be donated, but if it cannot be (for example because it is owned by a caller outside of the JIT context) XLA will insert a copy of the `payload` buffer.
+
+## Example
+
+```python
+import jax.numpy as jnp
+
+token1 = jnp.asarray(0, jnp.int32)
+x = some_native_read(token1)
+token2 = modify_token(token1)
+y = some_native_read(token2)
+```
+
+Here `token*` is a value that has some library object associated with it, and `modify_token` changes something about this backing object. Even if `modify_token` is marked as side-effecting, `x = some_native_read(token1)` and `token2 = modify_token(token1)` read the same token variable. This means XLA could choose to reorder them, which would lead to the following order:
+
+```python
+# Equivalent reordering:
+token1 = jnp.asarray(0, jnp.int32)
+token2 = modify_token(token1)
+x = some_native_read(token1)
+y = some_native_read(token2)
+```
+
+In this case, both `x` and `y` read the same value. To enforce ordering, we can entangle `x` and the `token` used in its computation, before modifying it.
+
+```python
+from entangle_jax import entangle
+
+token1 = jnp.asarray(0, jnp.int32)
+x = some_native_read(token1)
+token2 = entangle(token1, x)
+token3 = modify_token(token2)
+y = some_native_read(token3)
+```
+
+This means that `modify_token` cannot be reordered to be before the line that determines `x`, since it uses a value that depends on `x`.
+
+```python
+# Invalid reordering:
+token1 = jnp.asarray(0, jnp.int32)
+token2 = entangle(token1, x)  # x doesn't exist yet!
+token3 = modify_token(token2)
+x = some_native_read(token1)
+y = some_native_read(token3)
+```
+
+## Development
+
+```bash
+uv sync
+uv run pytest
+uv run ruff check .
+uv run ruff format --check .
+uv run ty check
+```
+
+The compiled extension rebuilds automatically on import while developing, so the
+default editable install (`uv sync`) is what you want day to day.

entangle_jax-0.1.0/README.md

@@ -0,0 +1,96 @@
+# entangle-jax
+
+A generic JAX primitive that is opaque to XLA reordering optimizations.
+
+`entangle(payload, *witnesses)` returns `payload` unchanged, but ordered under
+`jit` after every witness's producer. Useful for sequencing native
+side-effecting calls that XLA would otherwise be free to reorder, since they
+share no ordinary data dependency.
+
+Install with `uv add entangle-jax` or `pip install entangle-jax`.
+
+## Why
+
+When interacting with external native libraries, it is sometimes necessary to
+interact with objects they manage through opaque pointers. This conflicts with
+the XLA memory model where everything is a static-sized buffer that can be
+arbitrarily copied or reused when needed. One way to achieve this interaction is
+by creating a token that XLA can pass around that references this object, but
+XLA is not aware that e.g. multiple copies of this token in fact reference the
+same memory, and so will perform optimizations that are at worst unsafe and at
+best inefficient.
+
+Unfortunately, there is no way to tell XLA that a buffer will be modified
+in-place and that it should maintain consistent order of operations involving
+this buffer. Marking a primitive as side-effecting will stop XLA from removing
+it as dead code, but it will still assume that the input and output buffers are
+distinct, independent objects. At most, it is possible to create an in-place
+modifying primitive by using [input-output aliasing](https://openxla.org/xla/aliasing), in which case the input and
+output buffers are the same, but then XLA will simply copy the input buffer if
+it is used by another function to ensure both calls are safe.
+
+To work around this, we can create an artificial data dependency with
+`entangle`. The contents of a custom call are not visible to the XLA optimizer, so
+it has to treat the output `payload` as a new distinct variable from the input
+`payload` that depends on all `witness` values. A standard-library alternative
+would be `jax.lax.optimization_barrier`, but this is currently 
+[unreliable on CPU without a non-default XLA flag](https://github.com/openxla/xla/issues/20440).
+
+The actual runtime behavior of the primitive is a no-op: it uses aliasing to tell XLA that the input and output `payload` should be the same buffer, so it does literally nothing. This also means it is supported on every platform. The usual pattern is to overwrite the `payload` variable so that its buffer can always be donated, but if it cannot be (for example because it is owned by a caller outside of the JIT context) XLA will insert a copy of the `payload` buffer.
+
+## Example
+
+```python
+import jax.numpy as jnp
+
+token1 = jnp.asarray(0, jnp.int32)
+x = some_native_read(token1)
+token2 = modify_token(token1)
+y = some_native_read(token2)
+```
+
+Here `token*` is a value that has some library object associated with it, and `modify_token` changes something about this backing object. Even if `modify_token` is marked as side-effecting, `x = some_native_read(token1)` and `token2 = modify_token(token1)` read the same token variable. This means XLA could choose to reorder them, which would lead to the following order:
+
+```python
+# Equivalent reordering:
+token1 = jnp.asarray(0, jnp.int32)
+token2 = modify_token(token1)
+x = some_native_read(token1)
+y = some_native_read(token2)
+```
+
+In this case, both `x` and `y` read the same value. To enforce ordering, we can entangle `x` and the `token` used in its computation, before modifying it.
+
+```python
+from entangle_jax import entangle
+
+token1 = jnp.asarray(0, jnp.int32)
+x = some_native_read(token1)
+token2 = entangle(token1, x)
+token3 = modify_token(token2)
+y = some_native_read(token3)
+```
+
+This means that `modify_token` cannot be reordered to be before the line that determines `x`, since it uses a value that depends on `x`.
+
+```python
+# Invalid reordering:
+token1 = jnp.asarray(0, jnp.int32)
+token2 = entangle(token1, x)  # x doesn't exist yet!
+token3 = modify_token(token2)
+x = some_native_read(token1)
+y = some_native_read(token3)
+```
+
+## Development
+
+```bash
+uv sync
+uv run pytest
+uv run ruff check .
+uv run ruff format --check .
+uv run ty check
+```
+
+The compiled extension rebuilds automatically on import while developing, so the
+default editable install (`uv sync`) is what you want day to day.

entangle_jax-0.1.0/meson.build

@@ -0,0 +1,10 @@
+project(
+  'entangle-jax',
+  'cpp', 'cython',
+  version: '0.1.0',
+  default_options: ['cpp_std=c++17', 'buildtype=release'],
+)
+
+python = import('python').find_installation(pure: false)
+
+subdir('src/entangle_jax')

entangle_jax-0.1.0/pyproject.toml

@@ -0,0 +1,48 @@
+[build-system]
+build-backend = "mesonpy"
+requires = ["meson-python>=0.16", "Cython>=3.0", "jax>=0.10"]
+
+[project]
+name = "entangle-jax"
+version = "0.1.0"
+description = "A generic JAX primitive that is opaque to XLA reordering optimizations"
+readme = "README.md"
+license = "MIT"
+requires-python = ">=3.11"
+authors = [{ name = "Nardi Lam", email = "mail@nardilam.nl" }]
+dependencies = ["jax>=0.10", "jaxtyping>=0.2"]
+
+[project.urls]
+Homepage = "https://github.com/nardi/entangle-jax"
+
+[dependency-groups]
+dev = [
+    "pytest>=8.0",
+    "ruff>=0.6",
+    "meson>=1.3",
+    "ninja>=1.11",
+    "meson-python>=0.16",
+    "Cython>=3.0",
+    "ty>=0.0.1a1",
+]
+
+[tool.uv]
+no-build-isolation-package = ["entangle-jax"]
+
+[tool.meson-python.args]
+setup = ["--vsenv"]
+
+[tool.cibuildwheel]
+build = "{cp311-manylinux_x86_64,cp312-manylinux_x86_64,cp313-manylinux_x86_64,cp311-win_amd64,cp312-win_amd64,cp313-win_amd64,cp311-macosx_x86_64,cp312-macosx_x86_64,cp313-macosx_x86_64,cp311-macosx_arm64,cp312-macosx_arm64,cp313-macosx_arm64}"
+test-requires = "pytest"
+test-command = "pytest {project}/tests"
+
+[tool.ruff]
+line-length = 100
+target-version = "py311"
+
+[tool.ruff.lint]
+select = ["E", "F", "I", "UP", "B"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]

entangle_jax-0.1.0/src/entangle_jax/__init__.py

@@ -0,0 +1,15 @@
+"""A generic ordering primitive for JAX, opaque to XLA's algebraic simplifier.
+
+`entangle(payload, *witnesses)` returns `payload` unchanged, but ordered under `jit`
+after every witness's producer. Any XLA custom call is opaque to the algebraic
+simplifier by construction, so this holds without needing `has_side_effect=True`.
+It is more reliable than `jax.lax.optimization_barrier` which needs a non-default flag on
+CPU.
+
+Intended for sequencing native side-effecting custom calls that touch the same
+resource but share no ordinary data dependency.
+"""
+
+from entangle_jax._primitive import entangle
+
+__all__ = ["entangle"]

entangle_jax-0.1.0/src/entangle_jax/_entangle_ffi.cc

@@ -0,0 +1,30 @@
+#include "xla/ffi/api/c_api.h"
+#include "xla/ffi/api/ffi.h"
+
+#include "_entangle_ffi.h"
+
+namespace ffi = xla::ffi;
+
+namespace {
+
+ffi::Error entangle_impl(const ffi::AnyBuffer payload, const ffi::AnyBuffer witness,
+                          ffi::Result<ffi::AnyBuffer> out) {
+  (void)witness;
+  (void)payload;
+  (void)out;
+  return ffi::Error::Success();
+}
+
+XLA_FFI_DEFINE_HANDLER_SYMBOL(
+    entangle_handler, entangle_impl,
+    ffi::Ffi::Bind()
+        .Arg<ffi::AnyBuffer>()
+        .Arg<ffi::AnyBuffer>()
+        .Ret<ffi::AnyBuffer>()
+);
+
+}  // namespace
+
+extern "C" void* entangle_handler_address() {
+  return reinterpret_cast<void*>(entangle_handler);
+}

entangle_jax-0.1.0/src/entangle_jax/_entangle_ffi.h

@@ -0,0 +1,3 @@
+#pragma once
+
+extern "C" void* entangle_handler_address();

entangle_jax-0.1.0/src/entangle_jax/_ffi.pyi

@@ -0,0 +1,6 @@
+# Type stub for the compiled _ffi extension (built from _ffi.pyx).
+#
+# Importing this module registers the entangle_jax* XLA FFI targets as a side effect.
+# The only name meant for use from Python is the test-only slot reset below.
+
+def testing_reset_slots() -> None: ...

entangle_jax-0.1.0/src/entangle_jax/_ffi.pyx

@@ -0,0 +1,47 @@
+# distutils: language = c++
+# cython: language_level=3
+"""Registers the compiled entangle FFI handlers as JAX custom call targets.
+
+The handler addresses come from _entangle_ffi.cc and _testing_ffi.cc, both compiled
+directly into this extension. Each address is wrapped in a PyCapsule, the format JAX's
+FFI registration expects, and handed to jax.ffi.register_ffi_target. The _testing_*
+handlers back entangle_jax._testing, a private module used only by this package's own
+hazard regression test and not part of the public API.
+"""
+
+from cpython.pycapsule cimport PyCapsule_New
+
+
+cdef extern from "_entangle_ffi.h":
+    void* entangle_handler_address()
+
+cdef extern from "_testing_ffi.h":
+    void* write_handler_address()
+    void* read_handler_address()
+    void reset_slots()
+
+
+cdef object _capsule(void* address):
+    return PyCapsule_New(address, NULL, NULL)
+
+
+def _register_targets():
+    import jax
+
+    for platform in ("cpu", "cuda", "rocm", "tpu"):
+        jax.ffi.register_ffi_target(
+            "entangle_jax", _capsule(entangle_handler_address()), platform=platform
+        )
+        jax.ffi.register_ffi_target(
+            "entangle_jax_testing_write", _capsule(write_handler_address()), platform=platform
+        )
+        jax.ffi.register_ffi_target(
+            "entangle_jax_testing_read", _capsule(read_handler_address()), platform=platform
+        )
+
+
+_register_targets()
+
+
+def testing_reset_slots():
+    reset_slots()

entangle_jax-0.1.0/src/entangle_jax/_primitive.py

@@ -0,0 +1,91 @@
+"""Primitive definition for the entangle ordering operation."""
+
+from typing import TypeVar
+
+import jax.custom_batching
+import jax.extend.core
+import jax.interpreters.ad as ad
+import jax.interpreters.batching
+import jax.interpreters.mlir as mlir
+import jax.tree_util as jtu
+from jaxtyping import Array, PyTree
+
+from entangle_jax import _ffi  # noqa: F401
+
+_Payload = TypeVar("_Payload", bound=PyTree[Array])
+"""The pytree of arrays passed as `entangle`'s payload, returned unchanged so a caller
+keeps the exact type it passed in."""
+
+entangle_p = jax.extend.core.Primitive("entangle_jax")
+
+
+@entangle_p.def_impl
+def _entangle_impl(payload: Array, witness: Array) -> Array:
+    return jax.ffi.ffi_call(
+        "entangle_jax",
+        jax.ShapeDtypeStruct(payload.shape, payload.dtype),
+        input_output_aliases={0: 0},
+    )(payload, witness)
+
+
+@entangle_p.def_abstract_eval
+def _entangle_abstract_eval(payload, witness):
+    del witness
+    return payload
+
+
+mlir.register_lowering(entangle_p, mlir.lower_fun(_entangle_impl, multiple_results=False))
+
+
+def _entangle_p_vmap(vector_arg_values, batch_axes):
+    # entangle is a cheap passthrough, so a sequential map over the batch keeps
+    # it correct without a bespoke batching rule.
+    out = jax.vmap(jax.custom_batching.sequential_vmap(entangle_p.bind), in_axes=batch_axes)(
+        *vector_arg_values
+    )
+    return out, 0
+
+
+jax.interpreters.batching.primitive_batchers[entangle_p] = _entangle_p_vmap
+
+
+def _entangle_jvp(primals, tangents):
+    payload, witness = primals
+    t_payload, t_witness = tangents
+    if not isinstance(t_witness, ad.Zero):
+        raise NotImplementedError(
+            "entangle: cannot differentiate w.r.t. a witness argument. Witnesses only "
+            "establish an execution-order dependency. You should stop-gradient anything passed as "
+            "a witness if it originates from a differentiated computation."
+        )
+    out = entangle_p.bind(payload, witness)
+    if isinstance(t_payload, ad.Zero):
+        return out, ad.Zero(out.aval.to_tangent_aval())
+    return out, entangle_p.bind(t_payload, witness)
+
+
+ad.primitive_jvps[entangle_p] = _entangle_jvp
+
+
+def _entangle_transpose(cts_out, payload, witness):
+    if ad.is_undefined_primal(witness):
+        raise NotImplementedError("entangle: cannot differentiate w.r.t. a witness")
+    if ad.is_undefined_primal(payload):
+        return cts_out, None
+    return None, None
+
+
+ad.primitive_transposes[entangle_p] = _entangle_transpose
+
+
+def entangle(payload: _Payload, *witnesses: PyTree[Array]) -> _Payload:
+    """Return `payload` unchanged, ordered under `jit` after every witness's producer."""
+    payload_leaves, treedef = jtu.tree_flatten(payload)
+    witness_leaves = [leaf for w in witnesses for leaf in jtu.tree_leaves(w)]
+    if not witness_leaves:
+        return payload
+    signal, *rest = witness_leaves
+    for w in rest:
+        signal = entangle_p.bind(signal, w)
+    entangled_leaves = [entangle_p.bind(leaf, signal) for leaf in payload_leaves]
+    return jtu.tree_unflatten(treedef, entangled_leaves)

entangle_jax-0.1.0/src/entangle_jax/_testing.py

@@ -0,0 +1,33 @@
+"""Private test-only stand-ins for a mutable native resource.
+
+Used only by this package's own hazard regression test, to reproduce the
+sibling-calls reordering hazard entangle exists to fix, without needing a real
+consumer package. Not part of the public API. Do not import from outside this
+package's own test suite.
+"""
+
+import jax
+import jax.numpy as jnp
+
+from entangle_jax import _ffi  # noqa: F401
+
+
+def reset_slots():
+    _ffi.testing_reset_slots()
+
+
+def write(slot_id, value):
+    """Mutate slot `slot_id` to `value`, in place. Has side effect, mirrors refactor."""
+    return jax.ffi.ffi_call(
+        "entangle_jax_testing_write",
+        jax.ShapeDtypeStruct((), jnp.int32),
+        has_side_effect=True,
+    )(jnp.asarray(slot_id, jnp.int32), jnp.asarray(value, jnp.float32))
+
+
+def read(slot_id):
+    """Read slot `slot_id`. No side effect declared, mirrors solve."""
+    return jax.ffi.ffi_call(
+        "entangle_jax_testing_read",
+        jax.ShapeDtypeStruct((), jnp.float32),
+    )(jnp.asarray(slot_id, jnp.int32))

entangle_jax-0.1.0/src/entangle_jax/_testing_ffi.cc

@@ -0,0 +1,54 @@
+#include <cstdint>
+#include <cstring>
+
+#include "xla/ffi/api/c_api.h"
+#include "xla/ffi/api/ffi.h"
+
+#include "_testing_ffi.h"
+
+namespace ffi = xla::ffi;
+
+static double g_slots[16] = {0};
+
+ffi::Error write_impl(const ffi::Buffer<ffi::DataType::S32> id,
+                       const ffi::Buffer<ffi::DataType::F32> value,
+                       ffi::Result<ffi::Buffer<ffi::DataType::S32>> out_id) {
+  int32_t slot = *id.typed_data();
+  g_slots[slot] = *value.typed_data();
+  *out_id->typed_data() = slot;
+  return ffi::Error::Success();
+}
+
+XLA_FFI_DEFINE_HANDLER_SYMBOL(
+    write_handler, write_impl,
+    ffi::Ffi::Bind()
+        .Arg<ffi::Buffer<ffi::DataType::S32>>()
+        .Arg<ffi::Buffer<ffi::DataType::F32>>()
+        .Ret<ffi::Buffer<ffi::DataType::S32>>());
+
+ffi::Error read_impl(const ffi::Buffer<ffi::DataType::S32> id,
+                      ffi::Result<ffi::Buffer<ffi::DataType::F32>> out_value) {
+  int32_t slot = *id.typed_data();
+  *out_value->typed_data() = g_slots[slot];
+  return ffi::Error::Success();
+}
+
+XLA_FFI_DEFINE_HANDLER_SYMBOL(
+    read_handler, read_impl,
+    ffi::Ffi::Bind()
+        .Arg<ffi::Buffer<ffi::DataType::S32>>()
+        .Ret<ffi::Buffer<ffi::DataType::F32>>());
+
+extern "C" {
+void* write_handler_address() {
+  return reinterpret_cast<void*>(write_handler);
+}
+
+void* read_handler_address() {
+  return reinterpret_cast<void*>(read_handler);
+}
+
+void reset_slots() {
+  std::memset(g_slots, 0, sizeof(g_slots));
+}
+}

entangle_jax-0.1.0/src/entangle_jax/_testing_ffi.h

@@ -0,0 +1,7 @@
+#pragma once
+
+extern "C" {
+void* write_handler_address();
+void* read_handler_address();
+void reset_slots();
+}