Skip to main content

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, 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.

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

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:

# 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.

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.

# 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

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.

Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

entangle_jax-0.1.0.tar.gz (58.9 kB view details)

Uploaded Source

Built Distributions

If you're not sure about the file name format, learn more about wheel file names.

entangle_jax-0.1.0-cp313-cp313-win_amd64.whl (228.0 kB view details)

Uploaded CPython 3.13Windows x86-64

entangle_jax-0.1.0-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl (69.6 kB view details)

Uploaded CPython 3.13manylinux: glibc 2.24+ x86-64manylinux: glibc 2.28+ x86-64

entangle_jax-0.1.0-cp313-cp313-macosx_11_0_arm64.whl (43.4 kB view details)

Uploaded CPython 3.13macOS 11.0+ ARM64

entangle_jax-0.1.0-cp312-cp312-win_amd64.whl (228.8 kB view details)

Uploaded CPython 3.12Windows x86-64

entangle_jax-0.1.0-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl (70.3 kB view details)

Uploaded CPython 3.12manylinux: glibc 2.24+ x86-64manylinux: glibc 2.28+ x86-64

entangle_jax-0.1.0-cp312-cp312-macosx_11_0_arm64.whl (44.1 kB view details)

Uploaded CPython 3.12macOS 11.0+ ARM64

entangle_jax-0.1.0-cp311-cp311-win_amd64.whl (228.1 kB view details)

Uploaded CPython 3.11Windows x86-64

entangle_jax-0.1.0-cp311-cp311-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl (64.8 kB view details)

Uploaded CPython 3.11manylinux: glibc 2.24+ x86-64manylinux: glibc 2.28+ x86-64

entangle_jax-0.1.0-cp311-cp311-macosx_11_0_arm64.whl (43.2 kB view details)

Uploaded CPython 3.11macOS 11.0+ ARM64

File details

Details for the file entangle_jax-0.1.0.tar.gz.

File metadata

  • Download URL: entangle_jax-0.1.0.tar.gz
  • Upload date:
  • Size: 58.9 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: twine/7.0.0 CPython/3.13.14

File hashes

Hashes for entangle_jax-0.1.0.tar.gz
Algorithm Hash digest
SHA256 cca69518cd32774dc67c9eb98341e9e1dec10b1a2e83feaec4fa266118094d24
MD5 892a3e3d4c8da727f4c6b13ba79feafa
BLAKE2b-256 90c1a6ad9f796d5f4ac5eaea7d1cb1b6ffeeac9fc28a8da6aeac19af8a8cd097

See more details on using hashes here.

Provenance

The following attestation bundles were made for entangle_jax-0.1.0.tar.gz:

Publisher: publish.yml on nardi/entangle-jax

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file entangle_jax-0.1.0-cp313-cp313-win_amd64.whl.

File metadata

File hashes

Hashes for entangle_jax-0.1.0-cp313-cp313-win_amd64.whl
Algorithm Hash digest
SHA256 ee8725f91f1568de4dc237108bfe6ea8f2ed01cd0257a7e5c430cf76fb575501
MD5 91dbbf160ed898adf47dbce6cd6aa663
BLAKE2b-256 c1dc03f72921af53f85aebcf28ed9c473644fd28bdc3b4f50dd109af2a6b10da

See more details on using hashes here.

Provenance

The following attestation bundles were made for entangle_jax-0.1.0-cp313-cp313-win_amd64.whl:

Publisher: publish.yml on nardi/entangle-jax

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file entangle_jax-0.1.0-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl.

File metadata

File hashes

Hashes for entangle_jax-0.1.0-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl
Algorithm Hash digest
SHA256 b2d4607840db069fb8296a839d55a48dd664601d758fb59df6588ec45110e597
MD5 a30ed376141be7aed869f9f22ecab5c3
BLAKE2b-256 903e116d5382be56ee9d9516ae72f9a39f733f9876d933b8ce63e1fa07a93722

See more details on using hashes here.

Provenance

The following attestation bundles were made for entangle_jax-0.1.0-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl:

Publisher: publish.yml on nardi/entangle-jax

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file entangle_jax-0.1.0-cp313-cp313-macosx_11_0_arm64.whl.

File metadata

File hashes

Hashes for entangle_jax-0.1.0-cp313-cp313-macosx_11_0_arm64.whl
Algorithm Hash digest
SHA256 ceaa7be28d3e69ae4cc355ff51c6605ea749c38cf1bb825c9bfe418dc3079f95
MD5 4b4a87655843ec86104ccd8ec0b45fb2
BLAKE2b-256 2d5c79cb0d16b8499ad799a631d57f9b915f070ea641ba2c1f8ca40e64864a73

See more details on using hashes here.

Provenance

The following attestation bundles were made for entangle_jax-0.1.0-cp313-cp313-macosx_11_0_arm64.whl:

Publisher: publish.yml on nardi/entangle-jax

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file entangle_jax-0.1.0-cp312-cp312-win_amd64.whl.

File metadata

File hashes

Hashes for entangle_jax-0.1.0-cp312-cp312-win_amd64.whl
Algorithm Hash digest
SHA256 e9b80e0c5a35d63b1734bebcb5148942834f64988534992c700d97e5791378a1
MD5 0c4fcb5dfd33ac14256f7a4a028155d2
BLAKE2b-256 bd6a8eb75f6d05df1e584bf1626eab191ff10579b9bd7be8621ba3c30b9b7498

See more details on using hashes here.

Provenance

The following attestation bundles were made for entangle_jax-0.1.0-cp312-cp312-win_amd64.whl:

Publisher: publish.yml on nardi/entangle-jax

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file entangle_jax-0.1.0-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl.

File metadata

File hashes

Hashes for entangle_jax-0.1.0-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl
Algorithm Hash digest
SHA256 f00544883d02f761b4d9bcb9715ade2e06c5cae3ac457fa0e8ca2b7dfa239531
MD5 f29df56fc77351bd4bf4fe3512fbe244
BLAKE2b-256 48ffc4eb88e28acaec415cf2c2d034a6111ba7c1762b0384a094d614148e0f6f

See more details on using hashes here.

Provenance

The following attestation bundles were made for entangle_jax-0.1.0-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl:

Publisher: publish.yml on nardi/entangle-jax

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file entangle_jax-0.1.0-cp312-cp312-macosx_11_0_arm64.whl.

File metadata

File hashes

Hashes for entangle_jax-0.1.0-cp312-cp312-macosx_11_0_arm64.whl
Algorithm Hash digest
SHA256 d62052d8589ec8c340b93bc9bea91fd2a38bbc8830bb61507d6440b9eb0bf1af
MD5 06e5e439b30b897080a619ac1bf3a851
BLAKE2b-256 01fbb0df97e3ce2401ebef02ddddd5732dd00962717be2f2635b27020aae54f4

See more details on using hashes here.

Provenance

The following attestation bundles were made for entangle_jax-0.1.0-cp312-cp312-macosx_11_0_arm64.whl:

Publisher: publish.yml on nardi/entangle-jax

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file entangle_jax-0.1.0-cp311-cp311-win_amd64.whl.

File metadata

File hashes

Hashes for entangle_jax-0.1.0-cp311-cp311-win_amd64.whl
Algorithm Hash digest
SHA256 6faa3383f178a8b4625ad4191886e46d90d4759519ce1dc0defb8ee4e7716a5a
MD5 891e714208a0b3e8e2328a356c79ecd6
BLAKE2b-256 222a032c3feb5b39f297072dbec0750fadd5d4348bbbd990d6d0ea0664e3a2d4

See more details on using hashes here.

Provenance

The following attestation bundles were made for entangle_jax-0.1.0-cp311-cp311-win_amd64.whl:

Publisher: publish.yml on nardi/entangle-jax

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file entangle_jax-0.1.0-cp311-cp311-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl.

File metadata

File hashes

Hashes for entangle_jax-0.1.0-cp311-cp311-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl
Algorithm Hash digest
SHA256 07c10c8413e0c53957832ed907b7ff0ea976f6179ab75b311b34dabf98fc412e
MD5 82db69f0f02529d1acf01a9817ad6221
BLAKE2b-256 a6c72c0a03709e6b34c78276df046fedc06f7faeb17deec5994ff15f252c44e2

See more details on using hashes here.

Provenance

The following attestation bundles were made for entangle_jax-0.1.0-cp311-cp311-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl:

Publisher: publish.yml on nardi/entangle-jax

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

File details

Details for the file entangle_jax-0.1.0-cp311-cp311-macosx_11_0_arm64.whl.

File metadata

File hashes

Hashes for entangle_jax-0.1.0-cp311-cp311-macosx_11_0_arm64.whl
Algorithm Hash digest
SHA256 fe7e781a0f649d0a5c5b6d1623cb22b690daa1855f74225f7d018b681fa761bc
MD5 1533958c61ad566d5274a9951739527e
BLAKE2b-256 a9259b55ccfe153982a9961e5011267f2e53ad3ceb33f52c5e78179499bf6582

See more details on using hashes here.

Provenance

The following attestation bundles were made for entangle_jax-0.1.0-cp311-cp311-macosx_11_0_arm64.whl:

Publisher: publish.yml on nardi/entangle-jax

Attestations: Values shown here reflect the state when the release was signed and may no longer be current.

Release history Release notifications | RSS feed

This release

0.1.0 This release

10 files

Anthropic, PBC Visionary sponsor Bloomberg Visionary sponsor Hudson River Trading Visionary sponsor Meta Visionary sponsor NVIDIA Visionary sponsor Microsoft Sustainability sponsor Depot Continuous Integration AWS Cloud computing and Security Sponsor Datadog Monitoring Fastly CDN Google Download Analytics Sentry Error logging StatusPage Status page