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
Built Distributions
Filter files by name, interpreter, ABI, and platform.
If you're not sure about the file name format, learn more about wheel file names.
Copy a direct link to the current filters
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
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
cca69518cd32774dc67c9eb98341e9e1dec10b1a2e83feaec4fa266118094d24
|
|
| MD5 |
892a3e3d4c8da727f4c6b13ba79feafa
|
|
| BLAKE2b-256 |
90c1a6ad9f796d5f4ac5eaea7d1cb1b6ffeeac9fc28a8da6aeac19af8a8cd097
|
Provenance
The following attestation bundles were made for entangle_jax-0.1.0.tar.gz:
Publisher:
publish.yml on nardi/entangle-jax
-
Statement:
-
Statement type:
https://in-toto.io/Statement/v1 -
Predicate type:
https://docs.pypi.org/attestations/publish/v1 -
Subject name:
entangle_jax-0.1.0.tar.gz -
Subject digest:
cca69518cd32774dc67c9eb98341e9e1dec10b1a2e83feaec4fa266118094d24 - Sigstore transparency entry: 2818322113
- Sigstore integration time:
-
Permalink:
nardi/entangle-jax@948eea4273f2332b131361b964c755cbf3f835b8 -
Branch / Tag:
refs/tags/v0.1.0 - Owner: https://github.com/nardi
-
Access:
public
-
Token Issuer:
https://token.actions.githubusercontent.com -
Runner Environment:
github-hosted -
Publication workflow:
publish.yml@948eea4273f2332b131361b964c755cbf3f835b8 -
Trigger Event:
release
-
Statement type:
File details
Details for the file entangle_jax-0.1.0-cp313-cp313-win_amd64.whl.
File metadata
- Download URL: entangle_jax-0.1.0-cp313-cp313-win_amd64.whl
- Upload date:
- Size: 228.0 kB
- Tags: CPython 3.13, Windows x86-64
- Uploaded using Trusted Publishing? Yes
- Uploaded via:
twine/7.0.0 CPython/3.13.14
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
ee8725f91f1568de4dc237108bfe6ea8f2ed01cd0257a7e5c430cf76fb575501
|
|
| MD5 |
91dbbf160ed898adf47dbce6cd6aa663
|
|
| BLAKE2b-256 |
c1dc03f72921af53f85aebcf28ed9c473644fd28bdc3b4f50dd109af2a6b10da
|
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
-
Statement:
-
Statement type:
https://in-toto.io/Statement/v1 -
Predicate type:
https://docs.pypi.org/attestations/publish/v1 -
Subject name:
entangle_jax-0.1.0-cp313-cp313-win_amd64.whl -
Subject digest:
ee8725f91f1568de4dc237108bfe6ea8f2ed01cd0257a7e5c430cf76fb575501 - Sigstore transparency entry: 2818322305
- Sigstore integration time:
-
Permalink:
nardi/entangle-jax@948eea4273f2332b131361b964c755cbf3f835b8 -
Branch / Tag:
refs/tags/v0.1.0 - Owner: https://github.com/nardi
-
Access:
public
-
Token Issuer:
https://token.actions.githubusercontent.com -
Runner Environment:
github-hosted -
Publication workflow:
publish.yml@948eea4273f2332b131361b964c755cbf3f835b8 -
Trigger Event:
release
-
Statement type:
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
- Download URL: entangle_jax-0.1.0-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl
- Upload date:
- Size: 69.6 kB
- Tags: CPython 3.13, manylinux: glibc 2.24+ x86-64, manylinux: glibc 2.28+ x86-64
- Uploaded using Trusted Publishing? Yes
- Uploaded via:
twine/7.0.0 CPython/3.13.14
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
b2d4607840db069fb8296a839d55a48dd664601d758fb59df6588ec45110e597
|
|
| MD5 |
a30ed376141be7aed869f9f22ecab5c3
|
|
| BLAKE2b-256 |
903e116d5382be56ee9d9516ae72f9a39f733f9876d933b8ce63e1fa07a93722
|
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
-
Statement:
-
Statement type:
https://in-toto.io/Statement/v1 -
Predicate type:
https://docs.pypi.org/attestations/publish/v1 -
Subject name:
entangle_jax-0.1.0-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl -
Subject digest:
b2d4607840db069fb8296a839d55a48dd664601d758fb59df6588ec45110e597 - Sigstore transparency entry: 2818322582
- Sigstore integration time:
-
Permalink:
nardi/entangle-jax@948eea4273f2332b131361b964c755cbf3f835b8 -
Branch / Tag:
refs/tags/v0.1.0 - Owner: https://github.com/nardi
-
Access:
public
-
Token Issuer:
https://token.actions.githubusercontent.com -
Runner Environment:
github-hosted -
Publication workflow:
publish.yml@948eea4273f2332b131361b964c755cbf3f835b8 -
Trigger Event:
release
-
Statement type:
File details
Details for the file entangle_jax-0.1.0-cp313-cp313-macosx_11_0_arm64.whl.
File metadata
- Download URL: entangle_jax-0.1.0-cp313-cp313-macosx_11_0_arm64.whl
- Upload date:
- Size: 43.4 kB
- Tags: CPython 3.13, macOS 11.0+ ARM64
- Uploaded using Trusted Publishing? Yes
- Uploaded via:
twine/7.0.0 CPython/3.13.14
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
ceaa7be28d3e69ae4cc355ff51c6605ea749c38cf1bb825c9bfe418dc3079f95
|
|
| MD5 |
4b4a87655843ec86104ccd8ec0b45fb2
|
|
| BLAKE2b-256 |
2d5c79cb0d16b8499ad799a631d57f9b915f070ea641ba2c1f8ca40e64864a73
|
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
-
Statement:
-
Statement type:
https://in-toto.io/Statement/v1 -
Predicate type:
https://docs.pypi.org/attestations/publish/v1 -
Subject name:
entangle_jax-0.1.0-cp313-cp313-macosx_11_0_arm64.whl -
Subject digest:
ceaa7be28d3e69ae4cc355ff51c6605ea749c38cf1bb825c9bfe418dc3079f95 - Sigstore transparency entry: 2818322478
- Sigstore integration time:
-
Permalink:
nardi/entangle-jax@948eea4273f2332b131361b964c755cbf3f835b8 -
Branch / Tag:
refs/tags/v0.1.0 - Owner: https://github.com/nardi
-
Access:
public
-
Token Issuer:
https://token.actions.githubusercontent.com -
Runner Environment:
github-hosted -
Publication workflow:
publish.yml@948eea4273f2332b131361b964c755cbf3f835b8 -
Trigger Event:
release
-
Statement type:
File details
Details for the file entangle_jax-0.1.0-cp312-cp312-win_amd64.whl.
File metadata
- Download URL: entangle_jax-0.1.0-cp312-cp312-win_amd64.whl
- Upload date:
- Size: 228.8 kB
- Tags: CPython 3.12, Windows x86-64
- Uploaded using Trusted Publishing? Yes
- Uploaded via:
twine/7.0.0 CPython/3.13.14
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
e9b80e0c5a35d63b1734bebcb5148942834f64988534992c700d97e5791378a1
|
|
| MD5 |
0c4fcb5dfd33ac14256f7a4a028155d2
|
|
| BLAKE2b-256 |
bd6a8eb75f6d05df1e584bf1626eab191ff10579b9bd7be8621ba3c30b9b7498
|
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
-
Statement:
-
Statement type:
https://in-toto.io/Statement/v1 -
Predicate type:
https://docs.pypi.org/attestations/publish/v1 -
Subject name:
entangle_jax-0.1.0-cp312-cp312-win_amd64.whl -
Subject digest:
e9b80e0c5a35d63b1734bebcb5148942834f64988534992c700d97e5791378a1 - Sigstore transparency entry: 2818322638
- Sigstore integration time:
-
Permalink:
nardi/entangle-jax@948eea4273f2332b131361b964c755cbf3f835b8 -
Branch / Tag:
refs/tags/v0.1.0 - Owner: https://github.com/nardi
-
Access:
public
-
Token Issuer:
https://token.actions.githubusercontent.com -
Runner Environment:
github-hosted -
Publication workflow:
publish.yml@948eea4273f2332b131361b964c755cbf3f835b8 -
Trigger Event:
release
-
Statement type:
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
- Download URL: entangle_jax-0.1.0-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl
- Upload date:
- Size: 70.3 kB
- Tags: CPython 3.12, manylinux: glibc 2.24+ x86-64, manylinux: glibc 2.28+ x86-64
- Uploaded using Trusted Publishing? Yes
- Uploaded via:
twine/7.0.0 CPython/3.13.14
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
f00544883d02f761b4d9bcb9715ade2e06c5cae3ac457fa0e8ca2b7dfa239531
|
|
| MD5 |
f29df56fc77351bd4bf4fe3512fbe244
|
|
| BLAKE2b-256 |
48ffc4eb88e28acaec415cf2c2d034a6111ba7c1762b0384a094d614148e0f6f
|
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
-
Statement:
-
Statement type:
https://in-toto.io/Statement/v1 -
Predicate type:
https://docs.pypi.org/attestations/publish/v1 -
Subject name:
entangle_jax-0.1.0-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl -
Subject digest:
f00544883d02f761b4d9bcb9715ade2e06c5cae3ac457fa0e8ca2b7dfa239531 - Sigstore transparency entry: 2818322428
- Sigstore integration time:
-
Permalink:
nardi/entangle-jax@948eea4273f2332b131361b964c755cbf3f835b8 -
Branch / Tag:
refs/tags/v0.1.0 - Owner: https://github.com/nardi
-
Access:
public
-
Token Issuer:
https://token.actions.githubusercontent.com -
Runner Environment:
github-hosted -
Publication workflow:
publish.yml@948eea4273f2332b131361b964c755cbf3f835b8 -
Trigger Event:
release
-
Statement type:
File details
Details for the file entangle_jax-0.1.0-cp312-cp312-macosx_11_0_arm64.whl.
File metadata
- Download URL: entangle_jax-0.1.0-cp312-cp312-macosx_11_0_arm64.whl
- Upload date:
- Size: 44.1 kB
- Tags: CPython 3.12, macOS 11.0+ ARM64
- Uploaded using Trusted Publishing? Yes
- Uploaded via:
twine/7.0.0 CPython/3.13.14
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
d62052d8589ec8c340b93bc9bea91fd2a38bbc8830bb61507d6440b9eb0bf1af
|
|
| MD5 |
06e5e439b30b897080a619ac1bf3a851
|
|
| BLAKE2b-256 |
01fbb0df97e3ce2401ebef02ddddd5732dd00962717be2f2635b27020aae54f4
|
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
-
Statement:
-
Statement type:
https://in-toto.io/Statement/v1 -
Predicate type:
https://docs.pypi.org/attestations/publish/v1 -
Subject name:
entangle_jax-0.1.0-cp312-cp312-macosx_11_0_arm64.whl -
Subject digest:
d62052d8589ec8c340b93bc9bea91fd2a38bbc8830bb61507d6440b9eb0bf1af - Sigstore transparency entry: 2818322524
- Sigstore integration time:
-
Permalink:
nardi/entangle-jax@948eea4273f2332b131361b964c755cbf3f835b8 -
Branch / Tag:
refs/tags/v0.1.0 - Owner: https://github.com/nardi
-
Access:
public
-
Token Issuer:
https://token.actions.githubusercontent.com -
Runner Environment:
github-hosted -
Publication workflow:
publish.yml@948eea4273f2332b131361b964c755cbf3f835b8 -
Trigger Event:
release
-
Statement type:
File details
Details for the file entangle_jax-0.1.0-cp311-cp311-win_amd64.whl.
File metadata
- Download URL: entangle_jax-0.1.0-cp311-cp311-win_amd64.whl
- Upload date:
- Size: 228.1 kB
- Tags: CPython 3.11, Windows x86-64
- Uploaded using Trusted Publishing? Yes
- Uploaded via:
twine/7.0.0 CPython/3.13.14
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
6faa3383f178a8b4625ad4191886e46d90d4759519ce1dc0defb8ee4e7716a5a
|
|
| MD5 |
891e714208a0b3e8e2328a356c79ecd6
|
|
| BLAKE2b-256 |
222a032c3feb5b39f297072dbec0750fadd5d4348bbbd990d6d0ea0664e3a2d4
|
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
-
Statement:
-
Statement type:
https://in-toto.io/Statement/v1 -
Predicate type:
https://docs.pypi.org/attestations/publish/v1 -
Subject name:
entangle_jax-0.1.0-cp311-cp311-win_amd64.whl -
Subject digest:
6faa3383f178a8b4625ad4191886e46d90d4759519ce1dc0defb8ee4e7716a5a - Sigstore transparency entry: 2818322236
- Sigstore integration time:
-
Permalink:
nardi/entangle-jax@948eea4273f2332b131361b964c755cbf3f835b8 -
Branch / Tag:
refs/tags/v0.1.0 - Owner: https://github.com/nardi
-
Access:
public
-
Token Issuer:
https://token.actions.githubusercontent.com -
Runner Environment:
github-hosted -
Publication workflow:
publish.yml@948eea4273f2332b131361b964c755cbf3f835b8 -
Trigger Event:
release
-
Statement type:
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
- Download URL: entangle_jax-0.1.0-cp311-cp311-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl
- Upload date:
- Size: 64.8 kB
- Tags: CPython 3.11, manylinux: glibc 2.24+ x86-64, manylinux: glibc 2.28+ x86-64
- Uploaded using Trusted Publishing? Yes
- Uploaded via:
twine/7.0.0 CPython/3.13.14
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
07c10c8413e0c53957832ed907b7ff0ea976f6179ab75b311b34dabf98fc412e
|
|
| MD5 |
82db69f0f02529d1acf01a9817ad6221
|
|
| BLAKE2b-256 |
a6c72c0a03709e6b34c78276df046fedc06f7faeb17deec5994ff15f252c44e2
|
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
-
Statement:
-
Statement type:
https://in-toto.io/Statement/v1 -
Predicate type:
https://docs.pypi.org/attestations/publish/v1 -
Subject name:
entangle_jax-0.1.0-cp311-cp311-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl -
Subject digest:
07c10c8413e0c53957832ed907b7ff0ea976f6179ab75b311b34dabf98fc412e - Sigstore transparency entry: 2818322366
- Sigstore integration time:
-
Permalink:
nardi/entangle-jax@948eea4273f2332b131361b964c755cbf3f835b8 -
Branch / Tag:
refs/tags/v0.1.0 - Owner: https://github.com/nardi
-
Access:
public
-
Token Issuer:
https://token.actions.githubusercontent.com -
Runner Environment:
github-hosted -
Publication workflow:
publish.yml@948eea4273f2332b131361b964c755cbf3f835b8 -
Trigger Event:
release
-
Statement type:
File details
Details for the file entangle_jax-0.1.0-cp311-cp311-macosx_11_0_arm64.whl.
File metadata
- Download URL: entangle_jax-0.1.0-cp311-cp311-macosx_11_0_arm64.whl
- Upload date:
- Size: 43.2 kB
- Tags: CPython 3.11, macOS 11.0+ ARM64
- Uploaded using Trusted Publishing? Yes
- Uploaded via:
twine/7.0.0 CPython/3.13.14
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
fe7e781a0f649d0a5c5b6d1623cb22b690daa1855f74225f7d018b681fa761bc
|
|
| MD5 |
1533958c61ad566d5274a9951739527e
|
|
| BLAKE2b-256 |
a9259b55ccfe153982a9961e5011267f2e53ad3ceb33f52c5e78179499bf6582
|
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
-
Statement:
-
Statement type:
https://in-toto.io/Statement/v1 -
Predicate type:
https://docs.pypi.org/attestations/publish/v1 -
Subject name:
entangle_jax-0.1.0-cp311-cp311-macosx_11_0_arm64.whl -
Subject digest:
fe7e781a0f649d0a5c5b6d1623cb22b690daa1855f74225f7d018b681fa761bc - Sigstore transparency entry: 2818322185
- Sigstore integration time:
-
Permalink:
nardi/entangle-jax@948eea4273f2332b131361b964c755cbf3f835b8 -
Branch / Tag:
refs/tags/v0.1.0 - Owner: https://github.com/nardi
-
Access:
public
-
Token Issuer:
https://token.actions.githubusercontent.com -
Runner Environment:
github-hosted -
Publication workflow:
publish.yml@948eea4273f2332b131361b964c755cbf3f835b8 -
Trigger Event:
release
-
Statement type: