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]. 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.

The output tensor will be in TILE layout and have the same dtype as the input_tensor

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}")