Skip to content

Updates pooling docstrings #2064

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Apr 25, 2022
Merged

Updates pooling docstrings #2064

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
96f5b6e
Adds default value to pooling.pool)
marcvanzee Apr 25, 2022
d4ff240
Fixs some docstrings
marcvanzee Apr 25, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 6 additions & 6 deletions flax/linen/pooling.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,16 +25,16 @@ def pool(inputs, init, reduce_fn, window_shape, strides, padding):

Pooling functions are implemented using the ReduceWindow XLA op.
NOTE: Be aware that pooling is not generally differentiable.
That means providing a reduce_fn that is differentiable does not imply
that pool is differentiable.
That means providing a reduce_fn that is differentiable does not imply that
pool is differentiable.

Args:
inputs: input data with dimensions (batch, window dims..., features).
init: the initial value for the reduction
reduce_fn: a reduce function of the form `(T, T) -> T`.
window_shape: a shape tuple defining the window to reduce over.
strides: a sequence of `n` integers, representing the inter-window
strides.
strides (default: `(1, ..., 1)`).
padding: either the string `'SAME'`, the string `'VALID'`, or a sequence
of `n` `(low, high)` integer pairs that give the padding to apply before
and after each spatial dimension.
Expand Down Expand Up @@ -76,7 +76,7 @@ def avg_pool(inputs, window_shape, strides=None, padding="VALID"):
inputs: input data with dimensions (batch, window dims..., features).
window_shape: a shape tuple defining the window to reduce over.
strides: a sequence of `n` integers, representing the inter-window
strides (default: `(1, ..., 1)`).
strides (default: `(1, ..., 1)`).
padding: either the string `'SAME'`, the string `'VALID'`, or a sequence
of `n` `(low, high)` integer pairs that give the padding to apply before
and after each spatial dimension (default: `'VALID'`).
Expand All @@ -95,7 +95,7 @@ def max_pool(inputs, window_shape, strides=None, padding="VALID"):
inputs: input data with dimensions (batch, window dims..., features).
window_shape: a shape tuple defining the window to reduce over.
strides: a sequence of `n` integers, representing the inter-window
strides (default: `(1, ..., 1)`).
strides (default: `(1, ..., 1)`).
padding: either the string `'SAME'`, the string `'VALID'`, or a sequence
of `n` `(low, high)` integer pairs that give the padding to apply before
and after each spatial dimension (default: `'VALID'`).
Expand All @@ -113,7 +113,7 @@ def min_pool(inputs, window_shape, strides=None, padding="VALID"):
inputs: Input data with dimensions (batch, window dims..., features).
window_shape: A shape tuple defining the window to reduce over.
strides: A sequence of `n` integers, representing the inter-window strides
(default: `(1, ..., 1)`).
(default: `(1, ..., 1)`).
padding: Either the string `'SAME'`, the string `'VALID'`, or a sequence of
`n` `(low, high)` integer pairs that give the padding to apply before and
after each spatial dimension (default: `'VALID'`).
Expand Down