ttnn.var_hw

ttnn.var_hw(input_tensor: ttnn.Tensor, *, memory_config: ttnn.MemoryConfig = None) → ttnn.Tensor

Computes the variance across the height (H) and width (W) dimensions for each batch and channel. The variance is calculated as \(\mathrm{Var}[X] = E[(X - \mu)^2]\) where \(\mu\) is the mean over H and W dimensions. Output shape: [N, C, 1, 1].

\[\mathrm{{output\_tensor}}_i = \verb|var_hw|(\mathrm{{input\_tensor}}_i)\]

Parameters:

input_tensor (ttnn.Tensor) – the input tensor.

Keyword Arguments:

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

Returns:

ttnn.Tensor – the output tensor.

Note

Supported dtypes, layouts, and ranks:

Dtypes	Layouts	Ranks
BFLOAT16	TILE	2, 3, 4

Example

# Create a 4D tensor
tensor = ttnn.from_torch(torch.randn(1, 2, 64, 64, dtype=torch.bfloat16), layout=ttnn.TILE_LAYOUT, device=device)

# Compute variance across height and width dimensions
output = ttnn.var_hw(tensor)
logger.info(f"Variance HW: {output}")