ttnn.quantize

ttnn.quantize(input_tensor: ttnn.Tensor, scale: ttnn.Tensor or Number, zero_point: ttnn.Tensor or Number, *, axis: int = None, dtype: ttnn.DataType = None, memory_config: ttnn.MemoryConfig = None, output_tensor: ttnn.Tensor = None) → ttnn.Tensor

Quantizes a floating-point tensor into an integer tensor: q = input_tensor / scale + zero_point, per tensor by default or per channel along axis.

Parameters:

  • input_tensor (ttnn.Tensor)the input tensor.

  • scale (ttnn.Tensor or Number)the quantization scale.

  • zero_point (ttnn.Tensor or Number)the quantization zero point.

Keyword Arguments:

  • axis (int, optional)the axis of the quantization dimension of the input tensor. Defaults to None.

  • dtype (ttnn.DataType, optional)data type for the output tensor. Defaults to None.

  • memory_config (ttnn.MemoryConfig, optional)memory configuration for the operation. Defaults to None.

  • output_tensor (ttnn.Tensor, optional)preallocated output tensor. Defaults to None.

Returns:

ttnn.Tensorthe output tensor.

Note

Supported dtypes and layouts:

Dtypes

Layouts

BFLOAT16, BFLOAT8_B, BFLOAT4_B, FLOAT32

TILE

bfloat8_b/bfloat4_b supports only on TILE_LAYOUT

When scale and zero_point are tensors, they must be FLOAT32.

Example

# Create a float tensor to quantize
input_tensor = ttnn.from_torch(
    torch.tensor([[0.1, 0.2], [0.3, 0.4]], dtype=torch.bfloat16),
    dtype=ttnn.bfloat16,
    layout=ttnn.TILE_LAYOUT,
    device=device,
)

# Quantize the input onto the int8 range, giving [[-127, -42], [43, 127]]
scale = 0.001173
zero_point = -213
output = ttnn.quantize(input_tensor, scale, zero_point)
logger.info(f"Quantize result: {output}")