ttnn.global_avg_pool2d

ttnn.global_avg_pool2d(input_tensor, *, memory_config=None, dtype=None)

Global average pooling. Wrapper around avg_pool2d.

Mirrors PyTorch where global pooling has no dedicated op — it always routes through pool2d/adaptive_avg_pool2d. avg_pool2d’s pool2d() entry point detects the global-pool case (kernel == input spatial, no padding/dilation) and runs a single pool_sum reduction.

The caller does not need to flatten to (1, 1, N*H*W, C); the fast path inside pool2d() handles rank-4 NHWC directly via an explicit logical+padded reshape so it preserves pad-to-tile zero-padding from legacy callers.

Parameters:

input_tensor – Input tensor in NHWC format. Rank 2 [H,W], 3 [H,W,C], or 4 [N,H,W,C] accepted.

Keyword Arguments:

• memory_config – Optional output memory config.

• dtype – Optional output data type.

Returns:

ttnn.Tensor of shape (N, 1, 1, C) in TILE layout.

Example

# Create a random input tensor
tensor = ttnn.rand((10, 3, 32, 32), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=device)

# Perform global average pooling
output = ttnn.global_avg_pool2d(tensor)
logger.info(f"Output: {output}")