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()

```

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)

{"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"}]}