Implementation:NVIDIA TransformerEngine JAX Quantized Tensor
| Field | Value |
|---|---|
| Sources | TransformerEngine |
| Domains | Deep_Learning, JAX, Quantization |
| Last Updated | 2026-02-07 14:00 GMT |
Overview
Defines the tensor type hierarchy for quantized tensors, providing JAX pytree-compatible containers that carry quantized data alongside scale factors and metadata.
Description
AbstractBaseTensor defines the interface (dequantize(), get_tensor(), apply_sharding_constraint_by_logical_axes()). NoScaleTensor wraps higher-precision data with a no-op dequantize. ScaledTensor1x holds rowwise or colwise quantized data with a single scale. ScaledTensor2x combines both rowwise and colwise quantized views of the same tensor (needed for delayed scaling where both layouts are pre-computed). GroupedScaledTensor1x handles grouped quantization for MoE. ScaledTensorFactory creates the appropriate tensor type based on scaling mode and layout. All classes are registered as JAX pytree nodes via @register_pytree_node_class.
This is the fundamental data structure for FP8/FP4 tensors throughout the system -- every quantized tensor is represented as one of these types, enabling transparent passing through JAX transformations (jit, grad, vmap) while carrying quantization metadata.
Usage
Use ScaledTensorFactory.create() to construct quantized tensor wrappers. Use the dequantize() method to convert back to higher precision. These types are returned by quantizer classes and consumed by GEMM and other operations.
Code Reference
Source Location
- Repository
NVIDIA/TransformerEngine- File
transformer_engine/jax/quantize/tensor.py- Lines
- 1--843
Signature
class AbstractBaseTensor(ABC):
@abstractmethod
def dequantize(self) -> jnp.ndarray: ...
@abstractmethod
def get_tensor(self) -> jnp.ndarray: ...
class AbstractBaseTensor1x(AbstractBaseTensor): ...
class NoScaleTensor(AbstractBaseTensor1x):
"""Wraps higher-precision data with no-op dequantize."""
...
class ScaledTensor(ABC): ...
class ScaledTensor1x(AbstractBaseTensor1x, ScaledTensor):
"""Single quantization layout (rowwise or colwise) with one scale."""
data: jnp.ndarray
scale_inv: jnp.ndarray
...
class GroupedScaledTensor1x(ScaledTensor1x):
"""Grouped quantization for MoE workloads."""
...
class ScaledTensor2x(AbstractBaseTensor, ScaledTensor):
"""Both rowwise and colwise quantized views for delayed scaling."""
rowwise: ScaledTensor1x
colwise: ScaledTensor1x
...
class ScaledTensorFactory:
@staticmethod
def create(data, scale_inv, scaling_mode, ...) -> ScaledTensor: ...
Import
from transformer_engine.jax.quantize.tensor import ScaledTensor1x, ScaledTensor2x, ScaledTensorFactory, NoScaleTensor
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) |
| scaling_mode | ScalingMode |
Yes | The scaling mode used for quantization |
| dequantizer | Dequantizer |
Yes | Dequantization strategy for this tensor |
Outputs
| Name | Type | Description |
|---|---|---|
| tensor | ScaledTensor |
Quantized tensor container with metadata |
| dequantized | jnp.ndarray |
Higher-precision tensor (via dequantize() method) |
Usage Examples
from transformer_engine.jax.quantize.tensor import ScaledTensor1x, ScaledTensorFactory
# Create a quantized tensor via factory
qt = ScaledTensorFactory.create(
data=fp8_data, scale_inv=scale_inv,
scaling_mode=ScalingMode.CURRENT_TENSOR_SCALING,
dequantizer=tensor_scale_dequantizer,
)
# Dequantize back to higher precision
original = qt.dequantize()
# Access raw quantized data
raw_fp8 = qt.get_tensor()
# Apply sharding constraint
qt = qt.apply_sharding_constraint_by_logical_axes(("batch", "hidden"))