prompt: We are currently solving the following issue within our repository. Here is the issue text:
  --- BEGIN ISSUE ---
  Document default stride for pooling functions
  ### Discussed in https://github.com/google/flax/discussions/2023
  
  <div type='discussions-op-text'>
  
  <sup>Originally posted by **dogeplusplus** April  3, 2022</sup>
  A bit of a nitpick but I was wondering why the default behavior of pooling functions is to have stride 1 instead of the `window_shape`? I feel that for most use cases the stride would be the dimension of the kernel size as in other frameworks.</div>
  
  --- END ISSUE ---
  
  Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
  --- BEGIN FILES ---
  Path: `flax/linen/pooling.py`
  Content:
  ```
  1 # Copyright 2022 The Flax Authors.
  2 #
  3 # Licensed under the Apache License, Version 2.0 (the "License");
  4 # you may not use this file except in compliance with the License.
  5 # You may obtain a copy of the License at
  6 #
  7 #     http://www.apache.org/licenses/LICENSE-2.0
  8 #
  9 # Unless required by applicable law or agreed to in writing, software
  10 # distributed under the License is distributed on an "AS IS" BASIS,
  11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12 # See the License for the specific language governing permissions and
  13 # limitations under the License.
  14 
  15 """Pooling modules."""
  16 
  17 from jax import lax
  18 import jax.numpy as jnp
  19 
  20 import numpy as np
  21 
  22 
  23 def pool(inputs, init, reduce_fn, window_shape, strides, padding):
  24   """Helper function to define pooling functions.
  25 
  26   Pooling functions are implemented using the ReduceWindow XLA op.
  27   NOTE: Be aware that pooling is not generally differentiable.
  28   That means providing a reduce_fn that is differentiable does not imply
  29   that pool is differentiable.
  30 
  31   Args:
  32     inputs: input data with dimensions (batch, window dims..., features).
  33     init: the initial value for the reduction
  34     reduce_fn: a reduce function of the form `(T, T) -> T`.
  35     window_shape: a shape tuple defining the window to reduce over.
  36     strides: a sequence of `n` integers, representing the inter-window
  37         strides.
  38     padding: either the string `'SAME'`, the string `'VALID'`, or a sequence
  39       of `n` `(low, high)` integer pairs that give the padding to apply before
  40       and after each spatial dimension.
  41   Returns:
  42     The output of the reduction for each window slice.
  43   """
  44   strides = strides or (1,) * len(window_shape)
  45   assert len(window_shape) == len(strides), (
  46       f"len({window_shape}) must equal len({strides})")
  47   strides = (1,) + strides + (1,)
  48   dims = (1,) + window_shape + (1,)
  49 
  50   is_single_input = False
  51   if inputs.ndim == len(dims) - 1:
  52     # add singleton batch dimension because lax.reduce_window always
  53     # needs a batch dimension.
  54     inputs = inputs[None]
  55     is_single_input = True
  56 
  57   assert inputs.ndim == len(dims), f"len({inputs.shape}) != len({dims})"
  58   if not isinstance(padding, str):
  59     padding = tuple(map(tuple, padding))
  60     assert len(padding) == len(window_shape), (
  61         f"padding {padding} must specify pads for same number of dims as "
  62         f"window_shape {window_shape}")
  63     assert all([len(x) == 2 for x in padding]), (
  64         f"each entry in padding {padding} must be length 2")
  65     padding = ((0, 0),) + padding + ((0, 0),)
  66   y = lax.reduce_window(inputs, init, reduce_fn, dims, strides, padding)
  67   if is_single_input:
  68     y = jnp.squeeze(y, axis=0)
  69   return y
  70 
  71 
  72 def avg_pool(inputs, window_shape, strides=None, padding="VALID"):
  73   """Pools the input by taking the average over a window.
  74 
  75   Args:
  76     inputs: input data with dimensions (batch, window dims..., features).
  77     window_shape: a shape tuple defining the window to reduce over.
  78     strides: a sequence of `n` integers, representing the inter-window
  79         strides (default: `(1, ..., 1)`).
  80     padding: either the string `'SAME'`, the string `'VALID'`, or a sequence
  81       of `n` `(low, high)` integer pairs that give the padding to apply before
  82       and after each spatial dimension (default: `'VALID'`).
  83   Returns:
  84     The average for each window slice.
  85   """
  86   y = pool(inputs, 0., lax.add, window_shape, strides, padding)
  87   y = y / np.prod(window_shape)
  88   return y
  89 
  90 
  91 def max_pool(inputs, window_shape, strides=None, padding="VALID"):
  92   """Pools the input by taking the maximum of a window slice.
  93 
  94   Args:
  95     inputs: input data with dimensions (batch, window dims..., features).
  96     window_shape: a shape tuple defining the window to reduce over.
  97     strides: a sequence of `n` integers, representing the inter-window
  98         strides (default: `(1, ..., 1)`).
  99     padding: either the string `'SAME'`, the string `'VALID'`, or a sequence
  100       of `n` `(low, high)` integer pairs that give the padding to apply before
  101       and after each spatial dimension (default: `'VALID'`).
  102   Returns:
  103     The maximum for each window slice.
  104   """
  105   y = pool(inputs, -jnp.inf, lax.max, window_shape, strides, padding)
  106   return y
  107 
  108 
  109 def min_pool(inputs, window_shape, strides=None, padding="VALID"):
  110   """Pools the input by taking the minimum of a window slice.
  111 
  112   Args:
  113     inputs: Input data with dimensions (batch, window dims..., features).
  114     window_shape: A shape tuple defining the window to reduce over.
  115     strides: A sequence of `n` integers, representing the inter-window strides
  116         (default: `(1, ..., 1)`).
  117     padding: Either the string `'SAME'`, the string `'VALID'`, or a sequence of
  118       `n` `(low, high)` integer pairs that give the padding to apply before and
  119       after each spatial dimension (default: `'VALID'`).
  120 
  121   Returns:
  122     The minimum for each window slice.
  123   """
  124   return pool(inputs, jnp.inf, lax.min, window_shape, strides, padding)
  125
  ```
  --- END FILES ---
  
  Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.
  
  Here is an example:
  
  ```diff
  diff --git a/examples/server_async.py b/examples/server_async.py
  --- a/examples/server_async.py
  +++ b/examples/server_async.py
  @@ -313,4 +313,4 @@
   
   
   if __name__ == "__main__":
  -    asyncio.run(run_async_server("."), debug=True)
  +    asyncio.run(run_async_server(), debug=True)
  diff --git a/examples/server_sync.py b/examples/server_sync.py
  --- a/examples/server_sync.py
  +++ b/examples/server_sync.py
  @@ -313,5 +313,5 @@
   
   
   if __name__ == "__main__":
  -    server = run_sync_server(".")
  +    server = run_sync_server()
       server.shutdown()
  
  ```
  
golden_diff: diff --git a/flax/linen/pooling.py b/flax/linen/pooling.py
  --- a/flax/linen/pooling.py
  +++ b/flax/linen/pooling.py
  @@ -25,8 +25,8 @@
   
     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).
  @@ -34,7 +34,7 @@
       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.
  @@ -76,7 +76,7 @@
       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'`).
  @@ -95,7 +95,7 @@
       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'`).
  @@ -113,7 +113,7 @@
       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'`).
  
verification_info: {"golden_diff": "diff --git a/flax/linen/pooling.py b/flax/linen/pooling.py\n--- a/flax/linen/pooling.py\n+++ b/flax/linen/pooling.py\n@@ -25,8 +25,8 @@\n \n   Pooling functions are implemented using the ReduceWindow XLA op.\n   NOTE: Be aware that pooling is not generally differentiable.\n-  That means providing a reduce_fn that is differentiable does not imply\n-  that pool is differentiable.\n+  That means providing a reduce_fn that is differentiable does not imply that\n+  pool is differentiable.\n \n   Args:\n     inputs: input data with dimensions (batch, window dims..., features).\n@@ -34,7 +34,7 @@\n     reduce_fn: a reduce function of the form `(T, T) -> T`.\n     window_shape: a shape tuple defining the window to reduce over.\n     strides: a sequence of `n` integers, representing the inter-window\n-        strides.\n+      strides (default: `(1, ..., 1)`).\n     padding: either the string `'SAME'`, the string `'VALID'`, or a sequence\n       of `n` `(low, high)` integer pairs that give the padding to apply before\n       and after each spatial dimension.\n@@ -76,7 +76,7 @@\n     inputs: input data with dimensions (batch, window dims..., features).\n     window_shape: a shape tuple defining the window to reduce over.\n     strides: a sequence of `n` integers, representing the inter-window\n-        strides (default: `(1, ..., 1)`).\n+      strides (default: `(1, ..., 1)`).\n     padding: either the string `'SAME'`, the string `'VALID'`, or a sequence\n       of `n` `(low, high)` integer pairs that give the padding to apply before\n       and after each spatial dimension (default: `'VALID'`).\n@@ -95,7 +95,7 @@\n     inputs: input data with dimensions (batch, window dims..., features).\n     window_shape: a shape tuple defining the window to reduce over.\n     strides: a sequence of `n` integers, representing the inter-window\n-        strides (default: `(1, ..., 1)`).\n+      strides (default: `(1, ..., 1)`).\n     padding: either the string `'SAME'`, the string `'VALID'`, or a sequence\n       of `n` `(low, high)` integer pairs that give the padding to apply before\n       and after each spatial dimension (default: `'VALID'`).\n@@ -113,7 +113,7 @@\n     inputs: Input data with dimensions (batch, window dims..., features).\n     window_shape: A shape tuple defining the window to reduce over.\n     strides: A sequence of `n` integers, representing the inter-window strides\n-        (default: `(1, ..., 1)`).\n+      (default: `(1, ..., 1)`).\n     padding: Either the string `'SAME'`, the string `'VALID'`, or a sequence of\n       `n` `(low, high)` integer pairs that give the padding to apply before and\n       after each spatial dimension (default: `'VALID'`).\n", "issue": "Document default stride for pooling functions\n### Discussed in https://github.com/google/flax/discussions/2023\r\n\r\n<div type='discussions-op-text'>\r\n\r\n<sup>Originally posted by **dogeplusplus** April  3, 2022</sup>\r\nA bit of a nitpick but I was wondering why the default behavior of pooling functions is to have stride 1 instead of the `window_shape`? I feel that for most use cases the stride would be the dimension of the kernel size as in other frameworks.</div>\n", "before_files": [{"content": "# Copyright 2022 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Pooling modules.\"\"\"\n\nfrom jax import lax\nimport jax.numpy as jnp\n\nimport numpy as np\n\n\ndef pool(inputs, init, reduce_fn, window_shape, strides, padding):\n  \"\"\"Helper function to define pooling functions.\n\n  Pooling functions are implemented using the ReduceWindow XLA op.\n  NOTE: Be aware that pooling is not generally differentiable.\n  That means providing a reduce_fn that is differentiable does not imply\n  that pool is differentiable.\n\n  Args:\n    inputs: input data with dimensions (batch, window dims..., features).\n    init: the initial value for the reduction\n    reduce_fn: a reduce function of the form `(T, T) -> T`.\n    window_shape: a shape tuple defining the window to reduce over.\n    strides: a sequence of `n` integers, representing the inter-window\n        strides.\n    padding: either the string `'SAME'`, the string `'VALID'`, or a sequence\n      of `n` `(low, high)` integer pairs that give the padding to apply before\n      and after each spatial dimension.\n  Returns:\n    The output of the reduction for each window slice.\n  \"\"\"\n  strides = strides or (1,) * len(window_shape)\n  assert len(window_shape) == len(strides), (\n      f\"len({window_shape}) must equal len({strides})\")\n  strides = (1,) + strides + (1,)\n  dims = (1,) + window_shape + (1,)\n\n  is_single_input = False\n  if inputs.ndim == len(dims) - 1:\n    # add singleton batch dimension because lax.reduce_window always\n    # needs a batch dimension.\n    inputs = inputs[None]\n    is_single_input = True\n\n  assert inputs.ndim == len(dims), f\"len({inputs.shape}) != len({dims})\"\n  if not isinstance(padding, str):\n    padding = tuple(map(tuple, padding))\n    assert len(padding) == len(window_shape), (\n        f\"padding {padding} must specify pads for same number of dims as \"\n        f\"window_shape {window_shape}\")\n    assert all([len(x) == 2 for x in padding]), (\n        f\"each entry in padding {padding} must be length 2\")\n    padding = ((0, 0),) + padding + ((0, 0),)\n  y = lax.reduce_window(inputs, init, reduce_fn, dims, strides, padding)\n  if is_single_input:\n    y = jnp.squeeze(y, axis=0)\n  return y\n\n\ndef avg_pool(inputs, window_shape, strides=None, padding=\"VALID\"):\n  \"\"\"Pools the input by taking the average over a window.\n\n  Args:\n    inputs: input data with dimensions (batch, window dims..., features).\n    window_shape: a shape tuple defining the window to reduce over.\n    strides: a sequence of `n` integers, representing the inter-window\n        strides (default: `(1, ..., 1)`).\n    padding: either the string `'SAME'`, the string `'VALID'`, or a sequence\n      of `n` `(low, high)` integer pairs that give the padding to apply before\n      and after each spatial dimension (default: `'VALID'`).\n  Returns:\n    The average for each window slice.\n  \"\"\"\n  y = pool(inputs, 0., lax.add, window_shape, strides, padding)\n  y = y / np.prod(window_shape)\n  return y\n\n\ndef max_pool(inputs, window_shape, strides=None, padding=\"VALID\"):\n  \"\"\"Pools the input by taking the maximum of a window slice.\n\n  Args:\n    inputs: input data with dimensions (batch, window dims..., features).\n    window_shape: a shape tuple defining the window to reduce over.\n    strides: a sequence of `n` integers, representing the inter-window\n        strides (default: `(1, ..., 1)`).\n    padding: either the string `'SAME'`, the string `'VALID'`, or a sequence\n      of `n` `(low, high)` integer pairs that give the padding to apply before\n      and after each spatial dimension (default: `'VALID'`).\n  Returns:\n    The maximum for each window slice.\n  \"\"\"\n  y = pool(inputs, -jnp.inf, lax.max, window_shape, strides, padding)\n  return y\n\n\ndef min_pool(inputs, window_shape, strides=None, padding=\"VALID\"):\n  \"\"\"Pools the input by taking the minimum of a window slice.\n\n  Args:\n    inputs: Input data with dimensions (batch, window dims..., features).\n    window_shape: A shape tuple defining the window to reduce over.\n    strides: A sequence of `n` integers, representing the inter-window strides\n        (default: `(1, ..., 1)`).\n    padding: Either the string `'SAME'`, the string `'VALID'`, or a sequence of\n      `n` `(low, high)` integer pairs that give the padding to apply before and\n      after each spatial dimension (default: `'VALID'`).\n\n  Returns:\n    The minimum for each window slice.\n  \"\"\"\n  return pool(inputs, jnp.inf, lax.min, window_shape, strides, padding)\n", "path": "flax/linen/pooling.py"}], "after_files": [{"content": "# Copyright 2022 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Pooling modules.\"\"\"\n\nfrom jax import lax\nimport jax.numpy as jnp\n\nimport numpy as np\n\n\ndef pool(inputs, init, reduce_fn, window_shape, strides, padding):\n  \"\"\"Helper function to define pooling functions.\n\n  Pooling functions are implemented using the ReduceWindow XLA op.\n  NOTE: Be aware that pooling is not generally differentiable.\n  That means providing a reduce_fn that is differentiable does not imply that\n  pool is differentiable.\n\n  Args:\n    inputs: input data with dimensions (batch, window dims..., features).\n    init: the initial value for the reduction\n    reduce_fn: a reduce function of the form `(T, T) -> T`.\n    window_shape: a shape tuple defining the window to reduce over.\n    strides: a sequence of `n` integers, representing the inter-window\n      strides (default: `(1, ..., 1)`).\n    padding: either the string `'SAME'`, the string `'VALID'`, or a sequence\n      of `n` `(low, high)` integer pairs that give the padding to apply before\n      and after each spatial dimension.\n  Returns:\n    The output of the reduction for each window slice.\n  \"\"\"\n  strides = strides or (1,) * len(window_shape)\n  assert len(window_shape) == len(strides), (\n      f\"len({window_shape}) must equal len({strides})\")\n  strides = (1,) + strides + (1,)\n  dims = (1,) + window_shape + (1,)\n\n  is_single_input = False\n  if inputs.ndim == len(dims) - 1:\n    # add singleton batch dimension because lax.reduce_window always\n    # needs a batch dimension.\n    inputs = inputs[None]\n    is_single_input = True\n\n  assert inputs.ndim == len(dims), f\"len({inputs.shape}) != len({dims})\"\n  if not isinstance(padding, str):\n    padding = tuple(map(tuple, padding))\n    assert len(padding) == len(window_shape), (\n        f\"padding {padding} must specify pads for same number of dims as \"\n        f\"window_shape {window_shape}\")\n    assert all([len(x) == 2 for x in padding]), (\n        f\"each entry in padding {padding} must be length 2\")\n    padding = ((0, 0),) + padding + ((0, 0),)\n  y = lax.reduce_window(inputs, init, reduce_fn, dims, strides, padding)\n  if is_single_input:\n    y = jnp.squeeze(y, axis=0)\n  return y\n\n\ndef avg_pool(inputs, window_shape, strides=None, padding=\"VALID\"):\n  \"\"\"Pools the input by taking the average over a window.\n\n  Args:\n    inputs: input data with dimensions (batch, window dims..., features).\n    window_shape: a shape tuple defining the window to reduce over.\n    strides: a sequence of `n` integers, representing the inter-window\n      strides (default: `(1, ..., 1)`).\n    padding: either the string `'SAME'`, the string `'VALID'`, or a sequence\n      of `n` `(low, high)` integer pairs that give the padding to apply before\n      and after each spatial dimension (default: `'VALID'`).\n  Returns:\n    The average for each window slice.\n  \"\"\"\n  y = pool(inputs, 0., lax.add, window_shape, strides, padding)\n  y = y / np.prod(window_shape)\n  return y\n\n\ndef max_pool(inputs, window_shape, strides=None, padding=\"VALID\"):\n  \"\"\"Pools the input by taking the maximum of a window slice.\n\n  Args:\n    inputs: input data with dimensions (batch, window dims..., features).\n    window_shape: a shape tuple defining the window to reduce over.\n    strides: a sequence of `n` integers, representing the inter-window\n      strides (default: `(1, ..., 1)`).\n    padding: either the string `'SAME'`, the string `'VALID'`, or a sequence\n      of `n` `(low, high)` integer pairs that give the padding to apply before\n      and after each spatial dimension (default: `'VALID'`).\n  Returns:\n    The maximum for each window slice.\n  \"\"\"\n  y = pool(inputs, -jnp.inf, lax.max, window_shape, strides, padding)\n  return y\n\n\ndef min_pool(inputs, window_shape, strides=None, padding=\"VALID\"):\n  \"\"\"Pools the input by taking the minimum of a window slice.\n\n  Args:\n    inputs: Input data with dimensions (batch, window dims..., features).\n    window_shape: A shape tuple defining the window to reduce over.\n    strides: A sequence of `n` integers, representing the inter-window strides\n      (default: `(1, ..., 1)`).\n    padding: Either the string `'SAME'`, the string `'VALID'`, or a sequence of\n      `n` `(low, high)` integer pairs that give the padding to apply before and\n      after each spatial dimension (default: `'VALID'`).\n\n  Returns:\n    The minimum for each window slice.\n  \"\"\"\n  return pool(inputs, jnp.inf, lax.min, window_shape, strides, padding)\n", "path": "flax/linen/pooling.py"}]}