JAX serialize_executable.deserialize_and_load — Unsafe Deserialization

Bug

jax.experimental.serialize_executable.deserialize_and_load is a public, documented JAX API for loading ahead-of-time-compiled (AOT) JAX executables. Internally it instantiates a _JaxPjrtUnpickler over attacker-supplied bytes and calls .load(). The unpickler subclasses Python's standard deserializer but does NOT install a find_class allowlist. Standard reduce-protocol payload resolves any importable callable.

Files

exploit.py — Self-contained PoC. Builds a 170-byte payload and feeds it to deserialize_and_load. Drops sentinel /tmp/JAX_PWN_EXEC.txt.
repro.sh — One-command reproducer (sets up venv, runs exploit.py, asserts sentinel)

Reproduce

bash repro.sh

Expected: /tmp/JAX_PWN_EXEC.txt is created and repro.sh exits 0.

Affected version

jax==0.10.0 (latest PyPI release as of 2026-05-12). Tag: jax-v0.10.0.

Affected files in JAX source

jax/experimental/serialize_executable.py:48-78 — deserialize_and_load body (constructs unpickler over attacker bytes)
jax/experimental/serialize_executable.py:96-122 — _JaxPjrtUnpickler class (no find_class override)
jax/_src/pickle_util.py:63 — adjacent sink (cloudpickle.loads for host callbacks)
jax/experimental/colocated_python/serialization.py:279 — adjacent sink (RPC cloudpickle.loads)

Severity

CVSS 3.1 8.8 HIGH (AV:N/AC:L/PR:N/UI:R/S:U/C:H/I:H/A:H). ML-multiplier eligible: AOT-compiled JAX executables ARE the model artifact in federated learning, model marketplaces, distributed training rollouts, CI inference farms. Model integrity directly compromised.

Reporter

Tony Davis (lendtrain on huntr) — 2026-05-12.

Downloads last month: -; Downloads are not tracked for this model. How to track

Inference Providers NEW

This model isn't deployed by any Inference Provider. 🙋 Ask for provider support