AvgPool

Average pooling layer that averages over the penultimate dimension of an input tensor.

AvgPool()

Forward

(x: jaxtyping.Float[Tensor, '... n d']) -> jaxtyping.Float[Tensor, '... d']