---
id: jmp
version: "0.0.4"
license: Apache 2.0
license_treatment: permissive
maintenance: active
---

# jmp — JMP is a Mixed Precision library for JAX.

License: permissive · Maintenance: active · Downloads: 385.1K/mo

## What it is and what it does

JMP is a mixed precision training library that lets you train neural networks using both full-precision and half-precision floating-point numbers in the same model. This reduces memory bandwidth requirements and speeds up computation, particularly on GPUs and TPUs. The library provides two main abstractions: policies that define which dtypes to use for parameters, computation, and output, and loss scaling mechanisms that prevent gradient underflow when using reduced precision.

You define a policy once and apply it to your parameters and activations throughout training. JMP also handles loss scaling—either statically (you pick a fixed scale factor) or dynamically (the library adjusts the scale during training to avoid numerical issues). The dynamic loss scaler is optimized to have minimal overhead. Neural network libraries can integrate JMP to provide automatic mixed precision support, simplifying adoption in existing codebases.

Use it for:

- Reduce GPU training time by mixing float32 and float16 precision on NVIDIA hardware.
- Reduce TPU training time by mixing float32 and bfloat16 precision.
- Prevent gradient underflow in low-precision training by applying dynamic loss scaling automatically.
- Define reusable mixed precision policies as command-line arguments or hyperparameters for experiments.
- Integrate mixed precision into custom training loops without modifying core model code.

## Worth the install?

AI-flagged interpretation of the facts above — verify before relying.

JMP provides mixed precision training support for JAX by managing dtype policies and loss scaling to reduce memory bandwidth and improve computational efficiency during neural network training.

Yes, if you train neural networks and want to reduce memory bandwidth or speed up training on accelerators. The library is actively maintained, has low install friction, and is permissively licensed under Apache 2.0. The main gotcha is that JAX must be installed separately first with the correct accelerator support; beyond that, JMP is straightforward to integrate into existing training loops.

## Install

pip install jmp
uv add jmp
poetry add jmp

## Installing jmp

Before you install: Low friction: pure Python package with a wheel distribution. Maintenance is active with recent commits (last on 2026-07-08). JAX must be installed separately first with appropriate accelerator support, as JMP does not list it as a dependency to accommodate different CUDA configurations.

License in practice: Apache 2.0 (permissive): you can use, modify, and distribute JMP freely in commercial and private projects, provided you include the license notice and state significant changes.

Quickstart:

pip install jmp

import jmp
import numpy

# Define a mixed precision policy
policy = jmp.Policy(compute_dtype='float16',
                    param_dtype='float32',
                    output_dtype='float16')

# Cast inputs for computation
params, x = policy.cast_to_compute((params, x))
y = policy.cast_to_output(y)

JAX must be installed separately first with the correct accelerator support before installing JMP.

Verify before relying:

- Whether dynamic loss scaling's claimed 1-5% performance overhead is typical across different hardware and model sizes.
- Compatibility with JAX versions beyond those tested in the repository.
- Whether the package is actively maintained or in maintenance-only mode given 1292 days since last release.

## Package facts

- License: Apache 2.0 (permissive)
- Python support: unspecified
- Install friction: low
- Maintenance: active
- Downloads: 385.1K/month (top 15,000 on PyPI)
- Known vulnerabilities: none known

## Tags

mixed precision training, loss scaling neural networks, dtype policy management, half precision float16 training, automatic mixed precision AMP, gradient scaling training, bfloat16 float32 casting, mixed-precision, training-optimization

[View on SkillFed](https://skillfed.io/packages/jmp) · [View on PyPI](https://pypi.org/project/jmp/)