jmp
JMP is a Mixed Precision library for JAX.
Decision gist · record as of 2026-08-14
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.AI-flagged interpretation of the facts on this page — verify before relying
Before you install
- JAX must be installed separately first with the correct accelerator support before installing JMP.
- Low friction: pure Python package with a wheel distribution.
- Maintenance is active with recent commits (last on 2026-07-08).
License · maintenance · safety
Apache 2.0 (permissive) — 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.
last release 2023-01-30 (1292 days) · last repo commit 2026-07-08 · 213 stars
0 known vulnerabilities (OSV.dev, 2026-08-14) · 385,143 downloads/mo, #7,064 on PyPI
Alternatives
Verify before relying
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)- 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.
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 on it.
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
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.
JAX must be installed separately first with the correct accelerator support before installing JMP.
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)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 | Not specified |
| Install friction | Low. Pure-Python wheel |
| Runtime dependencies | 2 packagesnumpydataclasses |
| Maintenance | Actively maintained 1,292 days since the last release |
| Last repo commit | |
| First released | |
| Downloads | 385,143 / month, #7,064 on PyPI 30-day window, as of 2026-08-14 |
| Known vulnerabilities | None known OSV.dev, checked 2026-08-14 |
| Classifiers | Development Status :: 4 - BetaIntended Audience :: DevelopersIntended Audience :: EducationIntended Audience :: Science/ResearchLicense :: OSI Approved :: Apache Software LicenseProgramming Language :: Python :: 3Programming Language :: Python :: 3.8Programming Language :: Python :: 3.9Topic :: Scientific/Engineering :: MathematicsTopic :: Software Development :: LibrariesTopic :: Software Development :: Libraries :: Python Modules |
Evidence: jmp-0.0.4-py3-none-any.whl
Tags
Let your AI agent find packages like this
Example. Real query, live index.
You found this page by searching. An agent finds it by wishing: SkillFed indexes 14,416 PyPI packages by what they can do, searchable in plain language.
wish › “mixed precision training”
- jmpJMP provides mixed precision training support for JAX by managing…
- trainerTrainer is a PyTorch model training framework that handles the…
- coqui-tts-trainerA PyTorch model trainer framework that handles training loops,…
Give your agent the search over MCP, or paste the wish link into any chat.
More Libraries packages
urllib3 is an HTTP client library that provides thread-safe connection pooling, SSL/TLS verification, multipart file uploads, request retries, compression support, and proxy handling for Python applications.
Requests is a Python HTTP library that simplifies sending HTTP/1.1 requests with automatic handling of headers, authentication, cookies, and response parsing.
Pluggy provides a plugin system that lets you define hook specifications and register implementations to be called in sequence, enabling extensible Python applications without tight coupling.
Install it if you're building an extensible application or framework.
Provides parsing, arithmetic, and recurrence rule computation for dates and times, with timezone support and iCalendar RFC compliance.
Install it if you need to parse flexible date strings, compute relative dates, handle timezones, or work with recurrence rules—it's the de facto choice for these tasks.
Six provides utility functions to write Python code that runs on both Python 2.7 and Python 3.3+, smoothing over language differences between the two versions.
pytest is a testing framework that lets you write test functions using plain assert statements and automatically discovers and runs them, with detailed failure reporting.
See also optax · dm-haiku · jax · jax-cuda12-pjrt · jax-cuda13-pjrt · jraph · jax-cuda12-plugin · transformer-engine-cu13 · jaxlib · transformer-engine-cu12