prompt: We are currently solving the following issue within our repository. Here is the issue text:
  --- BEGIN ISSUE ---
  Pooling: passing "sequence of `n` `(low, high)` integer pairs" resulting in TypeError
  Trying to pass a tuple or list of tuples to a pool operation's padding parameter gives out the following errors: 
  `TypeError: Unknown padding type: (1, 1).`
  `TypeError : unhashable type: 'list' `
  
  
  Sample code for reproducing the bug:
  ```python3
  from flax import nn
  from jax import random
  
  class FlaxModel(nn.Module):
      def apply(self, x):
          x = nn.max_pool(x, (3, 3), strides=(2, 2), padding=[(1, 1), (1, 1)])
          return x
  
  rng = random.PRNGKey(0)
  model, _ = FlaxModel.init_by_shape(rng, [(1, 100, 100, 1)])
  ```
  
  
  --- 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/nn/pooling.py`
  Content:
  ```
  1 # Copyright 2020 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 onp
  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   strides = (1,) + strides + (1,)
  46   dims = (1,) + window_shape + (1,)
  47   return lax.reduce_window(inputs, init, reduce_fn, dims, strides, padding)
  48 
  49 
  50 def avg_pool(inputs, window_shape, strides=None, padding="VALID"):
  51   """Pools the input by taking the average over a window.
  52 
  53   Args:
  54     inputs: input data with dimensions (batch, window dims..., features).
  55     window_shape: a shape tuple defining the window to reduce over.
  56     strides: a sequence of `n` integers, representing the inter-window
  57         strides (default: `(1, ..., 1)`).
  58     padding: either the string `'SAME'`, the string `'VALID'`, or a sequence
  59       of `n` `(low, high)` integer pairs that give the padding to apply before
  60       and after each spatial dimension (default: `'VALID'`).
  61   Returns:
  62     The average for each window slice.
  63   """
  64   y = pool(inputs, 0., lax.add, window_shape, strides, padding)
  65   y = y / onp.prod(window_shape)
  66   return y
  67 
  68 
  69 def max_pool(inputs, window_shape, strides=None, padding="VALID"):
  70   """Pools the input by taking the maximum of a window slice.
  71 
  72   Args:
  73     inputs: input data with dimensions (batch, window dims..., features).
  74     window_shape: a shape tuple defining the window to reduce over.
  75     strides: a sequence of `n` integers, representing the inter-window
  76         strides (default: `(1, ..., 1)`).
  77     padding: either the string `'SAME'`, the string `'VALID'`, or a sequence
  78       of `n` `(low, high)` integer pairs that give the padding to apply before
  79       and after each spatial dimension (default: `'VALID'`).
  80   Returns:
  81     The maximum for each window slice.
  82   """
  83   y = pool(inputs, -jnp.inf, lax.max, window_shape, strides, padding)
  84   return y
  85
  ```
  --- 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/nn/pooling.py b/flax/nn/pooling.py
  --- a/flax/nn/pooling.py
  +++ b/flax/nn/pooling.py
  @@ -44,6 +44,14 @@
     strides = strides or (1,) * len(window_shape)
     strides = (1,) + strides + (1,)
     dims = (1,) + window_shape + (1,)
  +  if not isinstance(padding, str):
  +    padding = tuple(map(tuple, padding))
  +    assert(len(padding) == len(window_shape)), (
  +      f"padding {padding} must specify pads for same number of dims as "
  +      f"window_shape {window_shape}")
  +    assert(all([len(x) == 2 for x in padding])), (
  +      f"each entry in padding {padding} must be length 2")
  +    padding = ((0,0),) + padding + ((0,0),)
     return lax.reduce_window(inputs, init, reduce_fn, dims, strides, padding)
  
verification_info: {"golden_diff": "diff --git a/flax/nn/pooling.py b/flax/nn/pooling.py\n--- a/flax/nn/pooling.py\n+++ b/flax/nn/pooling.py\n@@ -44,6 +44,14 @@\n   strides = strides or (1,) * len(window_shape)\n   strides = (1,) + strides + (1,)\n   dims = (1,) + window_shape + (1,)\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   return lax.reduce_window(inputs, init, reduce_fn, dims, strides, padding)\n", "issue": "Pooling: passing \"sequence of `n` `(low, high)` integer pairs\" resulting in TypeError\nTrying to pass a tuple or list of tuples to a pool operation's padding parameter gives out the following errors: \r\n`TypeError: Unknown padding type: (1, 1).`\r\n`TypeError : unhashable type: 'list' `\r\n\r\n\r\nSample code for reproducing the bug:\r\n```python3\r\nfrom flax import nn\r\nfrom jax import random\r\n\r\nclass FlaxModel(nn.Module):\r\n    def apply(self, x):\r\n        x = nn.max_pool(x, (3, 3), strides=(2, 2), padding=[(1, 1), (1, 1)])\r\n        return x\r\n\r\nrng = random.PRNGKey(0)\r\nmodel, _ = FlaxModel.init_by_shape(rng, [(1, 100, 100, 1)])\r\n```\r\n\n", "before_files": [{"content": "# Copyright 2020 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 onp\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  strides = (1,) + strides + (1,)\n  dims = (1,) + window_shape + (1,)\n  return lax.reduce_window(inputs, init, reduce_fn, dims, strides, padding)\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 / onp.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", "path": "flax/nn/pooling.py"}], "after_files": [{"content": "# Copyright 2020 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 onp\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  strides = (1,) + strides + (1,)\n  dims = (1,) + window_shape + (1,)\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  return lax.reduce_window(inputs, init, reduce_fn, dims, strides, padding)\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 / onp.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", "path": "flax/nn/pooling.py"}]}