dm_pix: image ops that survive jax.jit, jax.vmap and jax.pmap

The problem is not image processing, it is image processing inside a transform

Most Python image libraries assume you are working eagerly on host arrays. You load a picture, call a function, get an array back. That model breaks the moment the operation sits in the middle of a JAX computation. A resize between two network layers, a random crop inside a training step, a colour conversion applied per sample in a batch: if the operation forces a transfer to the host, you lose the compiled graph and pay the round trip on every call.

PIX targets exactly that gap. The README frames the goal as providing image processing functions and tools to JAX in a way that can be optimised and parallelised through jax.jit, jax.vmap and jax.pmap. The audience is therefore narrow and identifiable: people writing JAX code who currently hand-roll flips, crops and colour maths, or who reach for a NumPy-based library and then fight it inside a jitted function. If your pipeline is PyTorch or plain NumPy, PIX is not aimed at you.

What the transform-first design commits the library to

The README makes a claim that is easy to skim past: all the functions in PIX can be jax.jit-ed, jax.vmap-ed and jax.pmap-ed. That is a constraint on how each function is written, not a marketing line. It rules out anything that needs Python-level control flow over array values, dynamic shapes that XLA cannot specialise, or host-side state.

The quickstart demonstrates the pattern on a single operation, flip_left_right. The same function is called four ways: directly, wrapped in jax.jit, wrapped in jax.vmap after adding a leading axis with image[np.newaxis, ...], and wrapped in jax.pmap for multi-device use. The README states that the four results agree up to accelerator floating point accuracy. That is the whole architecture in miniature. PIX supplies the primitive; JAX supplies tracing, batching and device placement. There is no PIX runtime, no scheduler, no session object.

The consequence cuts both ways. You get composition for free, because a PIX function is just a traced function like any other. You also inherit every JAX constraint, including the ones that have nothing to do with images. A vmap over a function that internally reshapes in a data-dependent way will fail, and it will fail as a JAX tracing error, not as a friendly PIX message.

Installation: JAX first, and the dependency PIX refuses to declare

This is the part of the README worth reading twice. PIX is written in pure Python but depends on C++ code through JAX. Because JAX installation differs by CUDA version, PIX does not list JAX as a dependency in pyproject.toml. The README says it is technically listed for reference but commented out.

So the order is fixed. Follow the JAX installation instructions to install JAX with the accelerator support you need, then run:

pip install dm-pix

That commented-out dependency is a deliberate trade. It prevents pip from pulling a CPU-only JAX wheel over your CUDA build, which is the failure everyone hits once. It also means pip cannot tell you that JAX is missing. If you install dm-pix into a clean environment and import it without JAX present, you will get an import error from JAX, not from PIX. The README does not describe a version compatibility matrix between PIX releases and JAX releases, and none is stated in the material available, so treat JAX pinning as your responsibility rather than something the package resolves for you.

Once installed, usage is one import: import dm_pix as pix. The README's example loads an image into a NumPy array with your preferred library, so PIX does not ship a loader and does not care how the array arrived.

Running the test suite as an environment check

The README treats the test suite as a way to validate your setup, not only to validate the library. Test files carry a _test suffix and run under pytest. From an installed checkout:

pip install -e ".[test]" python -m pytest [-n <NUMCPUS>] dm_pix

The -n flag is optional and maps to pytest-xdist for parallel execution, which matters on a machine with many cores or a slow accelerator. There is also a wrapper script for an isolated virtual environment:

./test.sh

Running the suite is the cheapest way to confirm that your JAX build, your accelerator and your PIX version agree before you build anything on top. If the tests fail at import or on device placement, the problem is in the JAX layer, and no amount of reading PIX source will fix it. The README does not state expected test duration or a minimum supported Python version, so those are things you find out by running it.

Where PIX stops being the right tool

The README's own examples section is the honest signal here. It says the examples in the examples/ folder are not much more involved than the flip_left_right snippet and may be a good starting point. That is a fair description of a library whose documented surface in the README is one operation plus pointers to Read the Docs.

Three cases where you should look elsewhere. First, classical computer vision with heavy geometry: feature detection, contour work, homographies, camera calibration. Nothing in the supplied material suggests PIX covers that ground, and a library built around JAX tracing is a poor fit for algorithms that branch on pixel values. Second, interactive or exploratory work. If you are poking at an image in a notebook and want immediate feedback, the JAX install and compilation overhead buys you nothing. Third, pipelines that must run outside JAX. Adopting PIX to call one resize function from a NumPy script adds a JAX dependency to your project in exchange for very little.

There is also a subtler limitation in the transform-first promise. jit, vmap and pmap are not free. A function that is jittable may still recompile when input shapes change, and image pipelines frequently change shape between training and evaluation. The README does not discuss compilation caching or shape polymorphism, so budget for profiling rather than assuming the transforms are costless.

How this differs from doing it in NumPy or SciPy

The obvious alternative is the NumPy and SciPy stack, which the README itself names as part of what JAX unifies. NumPy gives you a far wider set of array operations and SciPy adds image routines on top. The difference is not the operation list. It is where the operation executes.

A NumPy call inside a jitted function forces the array out of the compiled computation. You either restructure the pipeline to keep NumPy work outside the jit boundary, or you accept a host transfer per call. PIX exists so that the image operation stays inside the graph and inherits batching from vmap and device sharding from pmap. The flip_left_right example is trivial on purpose: the point is that the same call site works in all four execution modes without rewriting.

So the decision is architectural, not feature-by-feature. If your image ops are already batch-transformed JAX code, PIX removes hand-written primitives. If they are not, NumPy and SciPy remain the broader and better-documented choice, and adding PIX would put a JAX dependency in front of a problem that does not need one.

Licence, releases and what maintenance looks like from the outside

PIX is Apache-2.0, the same permissive licence used across the DeepMind JAX ecosystem, and the README points to the ecosystem citation file for academic use. Apache-2.0 includes an explicit patent grant and requires attribution and notice retention. That is a summary of the licence text, not legal advice; if you are redistributing PIX inside a product, read the licence and your own policy.

The release cadence visible in the material is uneven: v0.4.3 in July 2024, v0.4.4 in February 2025, v0.4.5 in June 2026. The repository is not archived and the last push postdates the latest release. Interpretation of that spacing is speculative, but the practical read is that you should not expect frequent breaking changes, and equally should not expect a rapid response to gaps in the operation list. Contributions are invited through the contributing guidelines and pull requests, which is the realistic path if you need a function that is not there.

Upgrade cost is dominated by JAX, not PIX. Because JAX is not a declared dependency, a PIX version bump will not pull a JAX bump with it, and a JAX bump will not be validated against PIX by your package manager. Pin both, and re-run the pytest suite after either changes. That suite is the only compatibility check the material describes.

Editorial conclusion

Adopt dm_pix if your image operations already live inside JAX-transformed code and you want them to compose with jax.jit, jax.vmap and jax.pmap without a host round trip. Do not adopt it if you need a broad classical CV toolbox, or if you are not prepared to pin and manage JAX yourself, since the project deliberately leaves that dependency to you. Before committing, verify two things in your own environment: that your JAX accelerator build matches the CUDA version you actually run, and that the specific pix function you need exists and behaves as you expect, because the README shows only flip_left_right and points to the examples/ folder and the Read the Docs site for the rest.

Official sources

1. google-deepmind/dm_pix on GitHub
2. License: Apache-2.0
3. Project website
4. README
5. Releases