Skip to main content

/

Docs

Nightly
v24.5
Log in

Mojo function

avg_pool

avg_pool(input: Symbol, filter_shape: Tuple[Int, Int], stride: Tuple[Int, Int] = Tuple(VariadicPack(<store_to_mem({1}), store_to_mem({1})>, 1)), dilation: Tuple[Int, Int] = Tuple(VariadicPack(<store_to_mem({1}), store_to_mem({1})>, 1)), padding: Tuple[Int, Int, Int, Int] = Tuple(VariadicPack(<store_to_mem({0}), store_to_mem({0}), store_to_mem({0}), store_to_mem({0})>, 1)), count_boundary: Bool = 1) -> Symbol

Computes average pooling with the given filter shape, strides, and dilations.

The op supports 2D avg pooling (so input and filter must be 4D), with the following layout assumption:

  • input has layout NHWC, i.e., (batch_size, height, width, in_channels)

Individual elements in the hyperparameters applies to corresponding dimensions of the input (after ignoring the batch and channel dimensions), with padding representing a before/after pair for each axis. The padding values are expected to take the form (pad_dim1_before, pad_dim1_after, pad_dim2_before, pad_dim2_after...). In 2D Convolution, dim1 here repesents H and dim2 represents W.

This op currently only supports strides and dilations on the filter.

Args:

  • ​input (Symbol): The rank-4 input tensor to perform the pooling upon, in NHWC layout.
  • ​filter_shape (Tuple[Int, Int]): The shape of the pooling filter.
  • ​stride (Tuple[Int, Int]): The stride of the pooling operation.
  • ​dilation (Tuple[Int, Int]): The spacing between the kernel points.
  • ​padding (Tuple[Int, Int, Int, Int]): The amount of padding applied to the input.
  • ​count_boundary (Bool): Whether to include the zero-padding in the averaging calculation.

Returns:

A symbolic tensor value with the pooling applied.