Implementation:NVIDIA TransformerEngine JAX Dequantizer
| Field | Value |
|---|---|
| Sources | TransformerEngine |
| Domains | Deep_Learning, JAX, Quantization |
| Last Updated | 2026-02-07 14:00 GMT |
Overview
Provides dequantization strategies for converting quantized tensors back to higher-precision formats, with implementations for each scaling mode.
Description
Dequantizer is an abstract base class defining the interface. NoopDequantizer returns data unchanged (for higher-precision tensors). TensorScaleDequantizer multiplies quantized data by the inverse scale factor (for delayed/current tensor scaling). BlockScaleDequantizer applies per-block E8M0 scale factors by reshaping data into blocks, broadcasting scales, and using 2^(scale-127) exponentiation. NVFP4Dequantizer handles FP4 dequantization with optional inverse Randomized Hadamard Transform. ScalingModeToDequantizerMap maps scaling modes to their appropriate dequantizer.
This module enables seamless conversion from FP8/MXFP8/NVFP4 quantized tensors back to standard precision, which is needed for operations that do not support low-precision inputs and for debugging/validation.
Usage
Use this module when you need to convert quantized tensors back to higher precision (e.g., for operations that do not support FP8/FP4 inputs, or for debugging). Dequantizers are typically accessed through the ScaledTensor.dequantize() method.
Code Reference
Source Location
- Repository
NVIDIA/TransformerEngine- File
transformer_engine/jax/quantize/dequantizer.py- Lines
- 1--341
Signature
class Dequantizer(ABC):
@abstractmethod
def dequantize(self, data, scale_inv, dtype) -> jnp.ndarray: ...
class NoopDequantizer(Dequantizer): ...
class TensorScaleDequantizer(Dequantizer): ...
class BlockScaleDequantizer(Dequantizer): ...
class NVFP4Dequantizer(Dequantizer): ...
Import
from transformer_engine.jax.quantize.dequantizer import Dequantizer, TensorScaleDequantizer, BlockScaleDequantizer
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| data | jnp.ndarray |
Yes | Quantized data tensor (FP8/FP4) |
| scale_inv | jnp.ndarray |
Yes | Inverse scale factor(s) |
| dtype | jnp.dtype |
Yes | Target output dtype (e.g., bfloat16, float32) |
Outputs
| Name | Type | Description |
|---|---|---|
| output | jnp.ndarray |
Dequantized tensor in the specified dtype |
Usage Examples
from transformer_engine.jax.quantize.dequantizer import TensorScaleDequantizer
import jax.numpy as jnp
# Dequantize an FP8 tensor with tensor scaling
dequantizer = TensorScaleDequantizer()
output = dequantizer.dequantize(fp8_data, scale_inv, jnp.bfloat16)
# More commonly, use the ScaledTensor interface:
# higher_precision = scaled_tensor.dequantize()