# dm_pix: JAX-native image processing from Google DeepMind

> dm_pix ports the everyday image operations (resize, flip, crop, colour conversion) into pure JAX, so they can be jitted, vmapped and pmapped. The library is small, the API is familiar, and the trade-off is that you must manage your own JAX installation.

**google-deepmind/dm_pix** — PIX is an image processing library in JAX, for JAX.

- Repository: https://github.com/google-deepmind/dm_pix
- Website: https://dm-pix.readthedocs.io
- Stars: 448 · Forks: 30
- Language: Python
- License: Apache-2.0
- Published: 2026-09-10 · Updated: 2026-09-10 · Language: en
- Canonical page: https://hysenlabs.com/projects/google-deepmind-dm-pix

## The gap dm_pix fills between NumPy image code and JAX transforms

Most image preprocessing is written against NumPy or a CPU-bound library. That is fine until the preprocessing has to live inside a jitted training step, at which point every host-side resize or colour conversion becomes a synchronization point. dm_pix exists to remove that boundary. Its stated goal is to provide image processing functions and tools to JAX in a way that they can be optimised and parallelised through jax.jit, jax.vmap and jax.pmap.

The intended audience is narrow and identifiable: researchers and engineers who already train or evaluate models in JAX and want augmentation and image manipulation to be part of the same compiled computation rather than a separate Python stage that hands arrays back and forth. If your stack is PyTorch, dm_pix offers nothing you cannot get elsewhere, and the README makes no attempt to reach you.

## How dm_pix works: pure Python functions over JAX arrays

The README describes PIX as written in pure Python, depending on C++ code only through JAX itself. There is no custom kernel, no compiled extension shipped by this repository. Each operation is a function that takes and returns JAX arrays, which is precisely why jax.jit can trace it, jax.vmap can batch it and jax.pmap can spread it across devices.

The repository layout matches that description. The dm_pix/ package holds the library, examples/ holds two worked examples (an image augmentation notebook and its .py counterpart), and docs/ backs the Read the Docs site. The only runtime dependency listed in pyproject.toml is chex, with jax and jaxlib present but commented out in the dependencies list and instead exposed through an extras group.

That last detail is the design decision worth noticing. Because JAX installation differs by CUDA version, dm_pix refuses to pin it. The consequence is that dm_pix inherits whatever JAX you installed, including its numerical behaviour. The README is explicit that results across the plain, jitted, vmapped and pmapped versions are the same only up to accelerator floating point accuracy.

## Installing dm_pix and running a first flip

The README is unambiguous that JAX comes first. Follow the JAX installation instructions to get JAX with the accelerator support you need, then install dm_pix from PyPI.

```bash
pip install dm-pix
```

After that, the import convention is a short alias. The README's quickstart loads an image into a NumPy array with whatever library you prefer and flips it left to right.

```python
import dm_pix as pix

# Load an image into a NumPy array with your preferred library.
image = load_image()

flip_left_right_image = pix.flip_left_right(image)
```

The same function can be wrapped in the JAX transforms. Note the leading axis the README adds before vmapping or pmapping, since both expect a batch dimension.

```python
import dm_pix as pix
import jax

image = load_image()
flip_left_right_image = jax.jit(pix.flip_left_right)(image)

image = image[np.newaxis, ...]
flip_left_right_image = jax.vmap(pix.flip_left_right)(image)
flip_left_right_image = jax.pmap(pix.flip_left_right)(image)
```

If you want to confirm the environment before writing your own code, the README points at the test suite. Every test file carries a _test suffix and runs under pytest.

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

The repository also ships test.sh, which the README describes as a utility for an isolated virtual environment. For a more involved starting point, examples/image_augmentation.py and the matching notebook are the closest thing to an end-to-end demonstration in the repository.

## Where dm_pix stops: I/O, Python versions and the JAX coupling

The quickstart's load_image() is a stand-in, not a function dm_pix provides. There is no decoder, no file reader, no dataset abstraction in the library. You bring your own loader, which in practice means Pillow (listed under the examples extra) or TensorFlow (listed under test). The library processes arrays; it does not acquire them.

The Python floor is 3.11, per requires-python in pyproject.toml. That is a real constraint for teams still on 3.9 or 3.10, and it is stricter than many JAX-adjacent packages. The project also classifies itself as Development Status :: 4 - Beta, so API stability should not be assumed across minor versions.

The coupling to JAX cuts both ways. Because JAX is not a declared dependency, pip will happily install dm_pix into an environment with an incompatible or CPU-only JAX, and the failure will surface at import or at first trace rather than at install time. Debugging that is your problem, not the package manager's. The README does not document a rollback path or a compatibility matrix between dm_pix releases and JAX versions, so pinning both sides is the only reliable approach.

## dm_pix compared with doing the same work in NumPy or TensorFlow

The obvious alternative is NumPy plus a CPU imaging library, which is what most people reach for first. The difference is not the arithmetic; a flip is a flip. It is where the operation executes. NumPy code runs on the host and forces a device-to-host transfer if your arrays live on an accelerator, and it cannot be traced into a jitted function at all. dm_pix keeps the operation inside the JAX computation, which is the entire reason to accept its constraints.

The second alternative is TensorFlow's image ops, or the preprocessing utilities that ship with a framework you already use. Those come with their own I/O and decoding, which dm_pix deliberately lacks, and they are tied to that framework's execution model. If you are not on JAX, they are the better fit. If you are on JAX, importing TensorFlow purely for a resize drags in a second numerical stack and a second set of device placement rules for one function call. The trade is scope for coherence: dm_pix does less, but everything it does composes with jax.jit, jax.vmap and jax.pmap by construction.

## Maintenance, releases and what the Apache-2.0 licence means here

The repository is not archived and the last push was on 2026-09-09, which is recent. Release cadence is slower than the commit history: v0.4.3 landed on 2024-07-03, v0.4.4 on 2025-02-17, and v0.4.5 on 2026-06-02. Between releases there is roughly a year, so if you depend on a fix that has been merged but not tagged, you are installing from git rather than from PyPI.

Upgrade cost is dominated by the JAX side. Since dm_pix does not pin JAX, a JAX upgrade can change numerical results or break tracing without any dm_pix release being involved, and the README's own caveat about accelerator floating point accuracy is a reminder that bit-exact reproducibility across backends is not promised. Budget for re-running your augmentation tests whenever either library moves.

The licence is Apache-2.0, declared in pyproject.toml and shipped as the LICENSE file at the repository root. That is a permissive licence with an explicit patent grant, which is generally the least friction option for commercial use. This is not legal advice; if the patent clause or notice requirements matter to your organisation, have counsel read the actual LICENSE text rather than a summary.

## Conclusion

Adopt dm_pix if your pipeline already runs on JAX and you need image operations that survive jax.jit, jax.vmap and jax.pmap without leaving the array abstraction. Do not adopt it if you need a standalone preprocessing stack with its own I/O, or if you are not prepared to install JAX yourself: the project deliberately leaves JAX out of its dependency list, so a plain pip install dm-pix will not pull an accelerator-enabled JAX with it. Before committing, check that the specific operation you need exists in the dm_pix namespace and that your Python is 3.11 or newer, then run the bundled test suite with pip install -e ".[test]" and python -m pytest dm_pix to confirm the install works on your hardware.

## FAQ

### What is a .pix file?

That question is about a file format, not about this library. dm_pix is a Python image processing library built on JAX, and it neither defines nor reads a .pix file format; it operates on arrays you load yourself.

### Does dm_pix install JAX for me when I run pip install dm-pix?

No. The README states that because JAX installation differs by CUDA version, PIX does not list JAX as a dependency in pyproject.toml, where it appears only as a commented reference. You install JAX first, then install dm-pix.

### Can dm_pix functions be used with jax.jit, jax.vmap and jax.pmap?

Yes. The README states that all functions in PIX can be jitted, vmapped and pmapped, and shows pix.flip_left_right used under all three transforms, with a leading batch axis added before vmapping or pmapping.

### Which Python versions does dm_pix support?

The pyproject.toml sets requires-python to ">=3.11", so anything older than Python 3.11 is outside the declared range. The project also classifies itself as Beta.

### How do I run the dm_pix test suite?

The README gives two paths: install the test extra with pip install -e ".[test]" and run python -m pytest dm_pix, optionally with a -n NUMCPUS argument, or run the bundled test.sh script for an isolated virtual environment.

## Sources

- [google-deepmind/dm_pix on GitHub](https://github.com/google-deepmind/dm_pix)
- [License: Apache-2.0](https://github.com/google-deepmind/dm_pix/blob/master/LICENSE)
- [Project website](https://dm-pix.readthedocs.io)
- [README](https://github.com/google-deepmind/dm_pix/blob/master/README.md)
- [Releases](https://github.com/google-deepmind/dm_pix/releases)

---

Hysen Labs editorial analysis, written from the project's own repository and release notes. Cite the canonical page: https://hysenlabs.com/projects/google-deepmind-dm-pix