ttnn.batch_norm

ttnn.batch_norm(input_tensor: ttnn.Tensor, *, eps: float = 1e-05, momentum: float = 0.1, running_mean: ttnn.Tensor = None, running_var: ttnn.Tensor = None, weight: ttnn.Tensor = None, bias: ttnn.Tensor = None, training: bool = False, output: ttnn.Tensor = None, memory_config: ttnn.MemoryConfig = None, compute_kernel_config: ttnn.DeviceComputeKernelConfig = None) → ttnn.Tensor

Applies batch norm over each channel on input_tensor. See Spatial Batch Normalization for more details.

\[\text{batch_norm}(x, \gamma, \beta, \epsilon) = \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}} \cdot \gamma + \beta\]
Where:

  • \(\mu\) and \(\sigma^2\) are the mean and variance of the input tensor, respectively

  • \(\gamma\) and \(\beta\) are the learnable scale and shift parameters, respectively

  • \(\epsilon\) is a small constant.

Parameters:

input_tensor (ttnn.Tensor) – the input tensor of shape [N, C, H, W].

Keyword Arguments:

  • eps (float, optional) – Epsilon value. Defaults to 1e-05.

  • momentum (float, optional) – Momentum value. Defaults to 0.1.

  • running_mean (ttnn.Tensor, optional) – the running_mean of shape [1, C, 1, 1], required in inference mode. When in training mode, this tensor is optional and the updated running mean value is stored in-place based on the inputs provided. Defaults to None.

  • running_var (ttnn.Tensor, optional) – the running_var of shape [1, C, 1, 1], required in inference mode. When in training mode, this tensor is optional and the updated running variance value is stored in-place based on the inputs provided. Defaults to None.

  • weight (ttnn.Tensor, optional) – the weight or gamma value of shape [1, C, 1, 1]. Defaults to None.

  • bias (ttnn.Tensor, optional) – the bias or beta value of shape [1, C, 1, 1]. Defaults to None.

  • training (bool, optional) – Selection between training mode and inference (evaluation) mode. Defaults to False (Inference mode).

  • output (ttnn.Tensor, optional) – Preallocated output tensor to store batch norm result of shape [N, C, H, W]. Must have the same dtype as input_tensor. Defaults to None.

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

  • compute_kernel_config (ttnn.DeviceComputeKernelConfig, optional) – device compute kernel configuration for the operation. Defaults to None.

Returns:

ttnn.Tensor – the output tensor.

Note

Supported dtypes, layouts, and ranks:

Dtypes

Layouts

Ranks

BFLOAT16, FLOAT32

TILE

4

These apply for all the tensor inputs to this operation, including the optional output tensor.

Mixed precision is supported: input (and output, if provided) may use a different dtype from the parameter tensors (running_mean, running_var, weight, bias). The output dtype always matches the input dtype. All parameter tensors must share the same dtype with each other.

Memory Support:

  • Interleaved: DRAM and L1

Limitations:

  • All input tensors must be tilized, interleaved, rank 4, and on-device.

Example

# Setup input tensor and parameters
N, C, H, W = 2, 3, 4, 5

input_tensor = ttnn.rand([N, C, H, W], dtype=ttnn.DataType.BFLOAT16, layout=ttnn.TILE_LAYOUT, device=device)
running_mean = ttnn.rand([1, C, 1, 1], dtype=ttnn.DataType.BFLOAT16, layout=ttnn.TILE_LAYOUT, device=device)
running_var = ttnn.rand([1, C, 1, 1], dtype=ttnn.DataType.BFLOAT16, layout=ttnn.TILE_LAYOUT, device=device)
weight = ttnn.rand([1, C, 1, 1], dtype=ttnn.DataType.BFLOAT16, layout=ttnn.TILE_LAYOUT, device=device)
bias = ttnn.from_torch(torch.rand([1, C, 1, 1], dtype=torch.bfloat16), layout=ttnn.TILE_LAYOUT, device=device)

# Apply batch normalization
output = ttnn.batch_norm(
    input_tensor,
    running_mean=running_mean,
    running_var=running_var,
    weight=weight,
    bias=bias,
    eps=1e-05,
    momentum=0.1,
    training=True,
)
logger.info(f"Batch Norm result: {output}")