jmp

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 on this page — 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

jmp on PyPI

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

Evidence: jmp-0.0.4-py3-none-any.whl

Tags

More Libraries packages