Module Kaun.Pool
2-D max and average pooling.
2-D pooling.
Pure, parameter-free window reductions over the last two axes — the spatial axes of the NCHW layout [| batch; channels; height; width |] used by Conv. Every function is differentiable through Rune.
Windows are placed wherever they fit entirely inside the input (no padding): pooling height with a window of kh and a stride of sh yields (height - kh) / sh + 1 positions, and likewise for the width. All leading axes are preserved.
val max_pool2d :
kernel_size:(int * int) ->
?stride:(int * int) ->
(float, 'b) Nx.t ->
(float, 'b) Nx.tmax_pool2d ~kernel_size:(kh, kw) x is the maximum of each kh × kw window over x's last two axes.
stride is the (vertical, horizontal) step between windows and defaults to kernel_size (non-overlapping windows).
Raises Invalid_argument if a kernel or stride component is not positive, if x has rank below 2, or if the window does not fit (height < kh or width < kw).
val avg_pool2d :
kernel_size:(int * int) ->
?stride:(int * int) ->
(float, 'b) Nx.t ->
(float, 'b) Nx.tavg_pool2d ~kernel_size x is like max_pool2d except each window is reduced to its mean. Same defaults and errors as max_pool2d.