We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Constant folding on `hk.avg_pool` with `padding='SAME'` on large arrays
Hi, lately my model has increased its compilation time due to constant folding and I discovered that the issue comes from using `padding='SAME'` in `hk.avg_pool`.

I tried to create a dummy code to showcase the issue (at the bottom), but the main point is that when `hk.avg_pool` uses `padding='VALID'` I get these output.

```
Wall time [Initialisation]: 0.112s
Wall time [Compilation]: 1.33s
Wall time [Train step]: 0.0024s
```

but when I use `padding='SAME'`
```
Wall time [Initialisation]: 0.119s
2022-10-12 16:01:47.590034: E external/org_tensorflow/tensorflow/compiler/xla/service/slow_operation_alarm.cc:65] Constant folding an instruction is taking > 1s:

  reduce-window.16 (displaying the full instruction incurs a runtime overhead. Raise your logging level to 4 or above).

This isn't necessarily a bug; constant-folding is inherently a trade-off between compilation time and speed at runtime.  XLA has some guards that attempt to keep constant folding from taking too long, but fundamentally you'll always be able to come up with an input program that takes a long time.

If you'd like to file a bug, run with envvar XLA_FLAGS=--xla_dump_to=/tmp/foo and attach the results.
2022-10-12 16:02:05.867327: E external/org_tensorflow/tensorflow/compiler/xla/service/slow_operation_alarm.cc:133] The operation took 19.277362298s
Constant folding an instruction is taking > 1s:

  reduce-window.16 (displaying the full instruction incurs a runtime overhead. Raise your logging level to 4 or above).

This isn't necessarily a bug; constant-folding is inherently a trade-off between compilation time and speed at runtime.  XLA has some guards that attempt to keep constant folding from taking too long, but fundamentally you'll always be able to come up with an input program that takes a long time.

If you'd like to file a bug, run with envvar XLA_FLAGS=--xla_dump_to=/tmp/foo and attach the results.
Wall time [Compilation]: 21.8s
Wall time [Train step]: 0.00221s
```
with `21.8s` of compilation!

The code:
```python
import haiku as hk
import jax
import contextlib
import timeit

@contextlib.contextmanager
def time_eval(task):
  start = timeit.default_timer()
  try:
    yield
  finally:
    end = timeit.default_timer()
    print(f'Wall time [{task}]: {(end - start):.3}s')

class Model(hk.Module):
    def __call__(self, x):
        x = hk.avg_pool(x, window_shape=(1,3,3,1), strides=(1,2,2,1), padding='SAME')
        x = hk.Conv2D(32, 4, 2)(x)
        x = jax.nn.relu(x)
        return x

def forward(x):
    return Model()(x)

forward = hk.without_apply_rng(hk.transform(forward))

rng = hk.PRNGSequence(jax.random.PRNGKey(42))
x = jax.random.uniform(next(rng), ([128, 512, 512, 3]))

with time_eval('Initialisation'):
    params = jax.jit(forward.init)(next(rng), x)

forward_apply = jax.jit(forward.apply)
with time_eval('Compilation'):
    logits = forward_apply(params, x).block_until_ready()

with time_eval('Train step'):
    logits = forward_apply(params, x).block_until_ready()

```

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `haiku/_src/pool.py`
Content:
```
1 # Copyright 2019 DeepMind Technologies Limited. All Rights Reserved.
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 Haiku modules."""
16 
17 import types
18 from typing import Optional, Sequence, Tuple, Union
19 import warnings
20 
21 from haiku._src import module
22 from jax import lax
23 import jax.numpy as jnp
24 import numpy as np
25 
26 # If you are forking replace this block with `import haiku as hk`.
27 hk = types.ModuleType("haiku")
28 hk.Module = module.Module
29 del module
30 
31 
32 def _infer_shape(
33     x: jnp.ndarray,
34     size: Union[int, Sequence[int]],
35     channel_axis: Optional[int] = -1,
36 ) -> Tuple[int, ...]:
37   """Infer shape for pooling window or strides."""
38   if isinstance(size, int):
39     if channel_axis and not 0 <= abs(channel_axis) < x.ndim:
40       raise ValueError(f"Invalid channel axis {channel_axis} for {x.shape}")
41     if channel_axis and channel_axis < 0:
42       channel_axis = x.ndim + channel_axis
43     return (1,) + tuple(size if d != channel_axis else 1
44                         for d in range(1, x.ndim))
45   elif len(size) < x.ndim:
46     # Assume additional dimensions are batch dimensions.
47     return (1,) * (x.ndim - len(size)) + tuple(size)
48   else:
49     assert x.ndim == len(size)
50     return tuple(size)
51 
52 
53 _VMAP_SHAPE_INFERENCE_WARNING = (
54     "When running under vmap, passing an `int` (except for `1`) for "
55     "`window_shape` or `strides` will result in the wrong shape being inferred "
56     "because the batch dimension is not visible to Haiku. Please update your "
57     "code to specify a full unbatched size. "
58     ""
59     "For example if you had `pool(x, window_shape=3, strides=1)` before, you "
60     "should now pass `pool(x, window_shape=(3, 3, 1), strides=1)`. "
61     ""
62     "Haiku will assume that any additional dimensions in your input are "
63     "batch dimensions, and will pad `window_shape` and `strides` accordingly "
64     "making your module support both batched and per-example inputs."
65 )
66 
67 
68 def _warn_if_unsafe(window_shape, strides):
69   unsafe = lambda size: isinstance(size, int) and size != 1
70   if unsafe(window_shape) or unsafe(strides):
71     warnings.warn(_VMAP_SHAPE_INFERENCE_WARNING, DeprecationWarning)
72 
73 
74 def max_pool(
75     value: jnp.ndarray,
76     window_shape: Union[int, Sequence[int]],
77     strides: Union[int, Sequence[int]],
78     padding: str,
79     channel_axis: Optional[int] = -1,
80 ) -> jnp.ndarray:
81   """Max pool.
82 
83   Args:
84     value: Value to pool.
85     window_shape: Shape of the pooling window, an int or same rank as value.
86     strides: Strides of the pooling window, an int or same rank as value.
87     padding: Padding algorithm. Either ``VALID`` or ``SAME``.
88     channel_axis: Axis of the spatial channels for which pooling is skipped,
89       used to infer ``window_shape`` or ``strides`` if they are an integer.
90 
91   Returns:
92     Pooled result. Same rank as value.
93   """
94   if padding not in ("SAME", "VALID"):
95     raise ValueError(f"Invalid padding '{padding}', must be 'SAME' or 'VALID'.")
96 
97   _warn_if_unsafe(window_shape, strides)
98   window_shape = _infer_shape(value, window_shape, channel_axis)
99   strides = _infer_shape(value, strides, channel_axis)
100 
101   return lax.reduce_window(value, -jnp.inf, lax.max, window_shape, strides,
102                            padding)
103 
104 
105 def avg_pool(
106     value: jnp.ndarray,
107     window_shape: Union[int, Sequence[int]],
108     strides: Union[int, Sequence[int]],
109     padding: str,
110     channel_axis: Optional[int] = -1,
111 ) -> jnp.ndarray:
112   """Average pool.
113 
114   Args:
115     value: Value to pool.
116     window_shape: Shape of the pooling window, an int or same rank as value.
117     strides: Strides of the pooling window, an int or same rank as value.
118     padding: Padding algorithm. Either ``VALID`` or ``SAME``.
119     channel_axis: Axis of the spatial channels for which pooling is skipped,
120       used to infer ``window_shape`` or ``strides`` if they are an integer.
121 
122   Returns:
123     Pooled result. Same rank as value.
124 
125   Raises:
126     ValueError: If the padding is not valid.
127   """
128   if padding not in ("SAME", "VALID"):
129     raise ValueError(f"Invalid padding '{padding}', must be 'SAME' or 'VALID'.")
130 
131   _warn_if_unsafe(window_shape, strides)
132   window_shape = _infer_shape(value, window_shape, channel_axis)
133   strides = _infer_shape(value, strides, channel_axis)
134 
135   reduce_window_args = (0., lax.add, window_shape, strides, padding)
136   pooled = lax.reduce_window(value, *reduce_window_args)
137   if padding == "VALID":
138     # Avoid the extra reduce_window.
139     return pooled / np.prod(window_shape)
140   else:
141     # Count the number of valid entries at each input point, then use that for
142     # computing average. Assumes that any two arrays of same shape will be
143     # padded the same.
144     window_counts = lax.reduce_window(jnp.ones_like(value), *reduce_window_args)
145     assert pooled.shape == window_counts.shape
146     return pooled / window_counts
147 
148 
149 class MaxPool(hk.Module):
150   """Max pool.
151 
152   Equivalent to partial application of :func:`max_pool`.
153   """
154 
155   def __init__(
156       self,
157       window_shape: Union[int, Sequence[int]],
158       strides: Union[int, Sequence[int]],
159       padding: str,
160       channel_axis: Optional[int] = -1,
161       name: Optional[str] = None,
162   ):
163     """Max pool.
164 
165     Args:
166       window_shape: Shape of window to pool over. Same rank as value or ``int``.
167       strides: Strides for the window. Same rank as value or ``int``.
168       padding: Padding algorithm. Either ``VALID`` or ``SAME``.
169       channel_axis: Axis of the spatial channels for which pooling is skipped.
170       name: String name for the module.
171     """
172     super().__init__(name=name)
173     self.window_shape = window_shape
174     self.strides = strides
175     self.padding = padding
176     self.channel_axis = channel_axis
177 
178   def __call__(self, value: jnp.ndarray) -> jnp.ndarray:
179     return max_pool(value, self.window_shape, self.strides,
180                     self.padding, self.channel_axis)
181 
182 
183 class AvgPool(hk.Module):
184   """Average pool.
185 
186   Equivalent to partial application of :func:`avg_pool`.
187   """
188 
189   def __init__(
190       self,
191       window_shape: Union[int, Sequence[int]],
192       strides: Union[int, Sequence[int]],
193       padding: str,
194       channel_axis: Optional[int] = -1,
195       name: Optional[str] = None,
196   ):
197     """Average pool.
198 
199     Args:
200       window_shape: Shape of window to pool over. Same rank as value or ``int``.
201       strides: Strides for the window. Same rank as value or ``int``.
202       padding: Padding algorithm. Either ``VALID`` or ``SAME``.
203       channel_axis: Axis of the spatial channels for which pooling is skipped.
204       name: String name for the module.
205     """
206     super().__init__(name=name)
207     self.window_shape = window_shape
208     self.strides = strides
209     self.padding = padding
210     self.channel_axis = channel_axis
211 
212   def __call__(self, value: jnp.ndarray) -> jnp.ndarray:
213     return avg_pool(value, self.window_shape, self.strides,
214                     self.padding, self.channel_axis)
215
```
--- 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/haiku/_src/pool.py b/haiku/_src/pool.py
--- a/haiku/_src/pool.py
+++ b/haiku/_src/pool.py
@@ -140,9 +140,9 @@
   else:
     # Count the number of valid entries at each input point, then use that for
     # computing average. Assumes that any two arrays of same shape will be
-    # padded the same.
-    window_counts = lax.reduce_window(jnp.ones_like(value), *reduce_window_args)
-    assert pooled.shape == window_counts.shape
+    # padded the same. Avoid broadcasting on axis where pooling is skipped. 
+    _shape = tuple(vd if wd!=1 else 1 for (vd, wd) in zip(value.shape, window_shape))
+    window_counts = lax.reduce_window(jnp.ones(_shape), *reduce_window_args)
     return pooled / window_counts

{"golden_diff": "diff --git a/haiku/_src/pool.py b/haiku/_src/pool.py\n--- a/haiku/_src/pool.py\n+++ b/haiku/_src/pool.py\n@@ -140,9 +140,9 @@\n   else:\n     # Count the number of valid entries at each input point, then use that for\n     # computing average. Assumes that any two arrays of same shape will be\n-    # padded the same.\n-    window_counts = lax.reduce_window(jnp.ones_like(value), *reduce_window_args)\n-    assert pooled.shape == window_counts.shape\n+    # padded the same. Avoid broadcasting on axis where pooling is skipped. \n+    _shape = tuple(vd if wd!=1 else 1 for (vd, wd) in zip(value.shape, window_shape))\n+    window_counts = lax.reduce_window(jnp.ones(_shape), *reduce_window_args)\n     return pooled / window_counts\n", "issue": "Constant folding on `hk.avg_pool` with `padding='SAME'` on large arrays\nHi, lately my model has increased its compilation time due to constant folding and I discovered that the issue comes from using `padding='SAME'` in `hk.avg_pool`.\r\n\r\nI tried to create a dummy code to showcase the issue (at the bottom), but the main point is that when `hk.avg_pool` uses `padding='VALID'` I get these output.\r\n\r\n```\r\nWall time [Initialisation]: 0.112s\r\nWall time [Compilation]: 1.33s\r\nWall time [Train step]: 0.0024s\r\n```\r\n\r\nbut when I use `padding='SAME'`\r\n```\r\nWall time [Initialisation]: 0.119s\r\n2022-10-12 16:01:47.590034: E external/org_tensorflow/tensorflow/compiler/xla/service/slow_operation_alarm.cc:65] Constant folding an instruction is taking > 1s:\r\n\r\n  reduce-window.16 (displaying the full instruction incurs a runtime overhead. Raise your logging level to 4 or above).\r\n\r\nThis isn't necessarily a bug; constant-folding is inherently a trade-off between compilation time and speed at runtime.  XLA has some guards that attempt to keep constant folding from taking too long, but fundamentally you'll always be able to come up with an input program that takes a long time.\r\n\r\nIf you'd like to file a bug, run with envvar XLA_FLAGS=--xla_dump_to=/tmp/foo and attach the results.\r\n2022-10-12 16:02:05.867327: E external/org_tensorflow/tensorflow/compiler/xla/service/slow_operation_alarm.cc:133] The operation took 19.277362298s\r\nConstant folding an instruction is taking > 1s:\r\n\r\n  reduce-window.16 (displaying the full instruction incurs a runtime overhead. Raise your logging level to 4 or above).\r\n\r\nThis isn't necessarily a bug; constant-folding is inherently a trade-off between compilation time and speed at runtime.  XLA has some guards that attempt to keep constant folding from taking too long, but fundamentally you'll always be able to come up with an input program that takes a long time.\r\n\r\nIf you'd like to file a bug, run with envvar XLA_FLAGS=--xla_dump_to=/tmp/foo and attach the results.\r\nWall time [Compilation]: 21.8s\r\nWall time [Train step]: 0.00221s\r\n```\r\nwith `21.8s` of compilation!\r\n\r\nThe code:\r\n```python\r\nimport haiku as hk\r\nimport jax\r\nimport contextlib\r\nimport timeit\r\n\r\n@contextlib.contextmanager\r\ndef time_eval(task):\r\n  start = timeit.default_timer()\r\n  try:\r\n    yield\r\n  finally:\r\n    end = timeit.default_timer()\r\n    print(f'Wall time [{task}]: {(end - start):.3}s')\r\n\r\nclass Model(hk.Module):\r\n    def __call__(self, x):\r\n        x = hk.avg_pool(x, window_shape=(1,3,3,1), strides=(1,2,2,1), padding='SAME')\r\n        x = hk.Conv2D(32, 4, 2)(x)\r\n        x = jax.nn.relu(x)\r\n        return x\r\n\r\ndef forward(x):\r\n    return Model()(x)\r\n\r\nforward = hk.without_apply_rng(hk.transform(forward))\r\n\r\nrng = hk.PRNGSequence(jax.random.PRNGKey(42))\r\nx = jax.random.uniform(next(rng), ([128, 512, 512, 3]))\r\n\r\nwith time_eval('Initialisation'):\r\n    params = jax.jit(forward.init)(next(rng), x)\r\n\r\nforward_apply = jax.jit(forward.apply)\r\nwith time_eval('Compilation'):\r\n    logits = forward_apply(params, x).block_until_ready()\r\n\r\nwith time_eval('Train step'):\r\n    logits = forward_apply(params, x).block_until_ready()\r\n\r\n```\n", "before_files": [{"content": "# Copyright 2019 DeepMind Technologies Limited. All Rights Reserved.\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 Haiku modules.\"\"\"\n\nimport types\nfrom typing import Optional, Sequence, Tuple, Union\nimport warnings\n\nfrom haiku._src import module\nfrom jax import lax\nimport jax.numpy as jnp\nimport numpy as np\n\n# If you are forking replace this block with `import haiku as hk`.\nhk = types.ModuleType(\"haiku\")\nhk.Module = module.Module\ndel module\n\n\ndef _infer_shape(\n    x: jnp.ndarray,\n    size: Union[int, Sequence[int]],\n    channel_axis: Optional[int] = -1,\n) -> Tuple[int, ...]:\n  \"\"\"Infer shape for pooling window or strides.\"\"\"\n  if isinstance(size, int):\n    if channel_axis and not 0 <= abs(channel_axis) < x.ndim:\n      raise ValueError(f\"Invalid channel axis {channel_axis} for {x.shape}\")\n    if channel_axis and channel_axis < 0:\n      channel_axis = x.ndim + channel_axis\n    return (1,) + tuple(size if d != channel_axis else 1\n                        for d in range(1, x.ndim))\n  elif len(size) < x.ndim:\n    # Assume additional dimensions are batch dimensions.\n    return (1,) * (x.ndim - len(size)) + tuple(size)\n  else:\n    assert x.ndim == len(size)\n    return tuple(size)\n\n\n_VMAP_SHAPE_INFERENCE_WARNING = (\n    \"When running under vmap, passing an `int` (except for `1`) for \"\n    \"`window_shape` or `strides` will result in the wrong shape being inferred \"\n    \"because the batch dimension is not visible to Haiku. Please update your \"\n    \"code to specify a full unbatched size. \"\n    \"\"\n    \"For example if you had `pool(x, window_shape=3, strides=1)` before, you \"\n    \"should now pass `pool(x, window_shape=(3, 3, 1), strides=1)`. \"\n    \"\"\n    \"Haiku will assume that any additional dimensions in your input are \"\n    \"batch dimensions, and will pad `window_shape` and `strides` accordingly \"\n    \"making your module support both batched and per-example inputs.\"\n)\n\n\ndef _warn_if_unsafe(window_shape, strides):\n  unsafe = lambda size: isinstance(size, int) and size != 1\n  if unsafe(window_shape) or unsafe(strides):\n    warnings.warn(_VMAP_SHAPE_INFERENCE_WARNING, DeprecationWarning)\n\n\ndef max_pool(\n    value: jnp.ndarray,\n    window_shape: Union[int, Sequence[int]],\n    strides: Union[int, Sequence[int]],\n    padding: str,\n    channel_axis: Optional[int] = -1,\n) -> jnp.ndarray:\n  \"\"\"Max pool.\n\n  Args:\n    value: Value to pool.\n    window_shape: Shape of the pooling window, an int or same rank as value.\n    strides: Strides of the pooling window, an int or same rank as value.\n    padding: Padding algorithm. Either ``VALID`` or ``SAME``.\n    channel_axis: Axis of the spatial channels for which pooling is skipped,\n      used to infer ``window_shape`` or ``strides`` if they are an integer.\n\n  Returns:\n    Pooled result. Same rank as value.\n  \"\"\"\n  if padding not in (\"SAME\", \"VALID\"):\n    raise ValueError(f\"Invalid padding '{padding}', must be 'SAME' or 'VALID'.\")\n\n  _warn_if_unsafe(window_shape, strides)\n  window_shape = _infer_shape(value, window_shape, channel_axis)\n  strides = _infer_shape(value, strides, channel_axis)\n\n  return lax.reduce_window(value, -jnp.inf, lax.max, window_shape, strides,\n                           padding)\n\n\ndef avg_pool(\n    value: jnp.ndarray,\n    window_shape: Union[int, Sequence[int]],\n    strides: Union[int, Sequence[int]],\n    padding: str,\n    channel_axis: Optional[int] = -1,\n) -> jnp.ndarray:\n  \"\"\"Average pool.\n\n  Args:\n    value: Value to pool.\n    window_shape: Shape of the pooling window, an int or same rank as value.\n    strides: Strides of the pooling window, an int or same rank as value.\n    padding: Padding algorithm. Either ``VALID`` or ``SAME``.\n    channel_axis: Axis of the spatial channels for which pooling is skipped,\n      used to infer ``window_shape`` or ``strides`` if they are an integer.\n\n  Returns:\n    Pooled result. Same rank as value.\n\n  Raises:\n    ValueError: If the padding is not valid.\n  \"\"\"\n  if padding not in (\"SAME\", \"VALID\"):\n    raise ValueError(f\"Invalid padding '{padding}', must be 'SAME' or 'VALID'.\")\n\n  _warn_if_unsafe(window_shape, strides)\n  window_shape = _infer_shape(value, window_shape, channel_axis)\n  strides = _infer_shape(value, strides, channel_axis)\n\n  reduce_window_args = (0., lax.add, window_shape, strides, padding)\n  pooled = lax.reduce_window(value, *reduce_window_args)\n  if padding == \"VALID\":\n    # Avoid the extra reduce_window.\n    return pooled / np.prod(window_shape)\n  else:\n    # Count the number of valid entries at each input point, then use that for\n    # computing average. Assumes that any two arrays of same shape will be\n    # padded the same.\n    window_counts = lax.reduce_window(jnp.ones_like(value), *reduce_window_args)\n    assert pooled.shape == window_counts.shape\n    return pooled / window_counts\n\n\nclass MaxPool(hk.Module):\n  \"\"\"Max pool.\n\n  Equivalent to partial application of :func:`max_pool`.\n  \"\"\"\n\n  def __init__(\n      self,\n      window_shape: Union[int, Sequence[int]],\n      strides: Union[int, Sequence[int]],\n      padding: str,\n      channel_axis: Optional[int] = -1,\n      name: Optional[str] = None,\n  ):\n    \"\"\"Max pool.\n\n    Args:\n      window_shape: Shape of window to pool over. Same rank as value or ``int``.\n      strides: Strides for the window. Same rank as value or ``int``.\n      padding: Padding algorithm. Either ``VALID`` or ``SAME``.\n      channel_axis: Axis of the spatial channels for which pooling is skipped.\n      name: String name for the module.\n    \"\"\"\n    super().__init__(name=name)\n    self.window_shape = window_shape\n    self.strides = strides\n    self.padding = padding\n    self.channel_axis = channel_axis\n\n  def __call__(self, value: jnp.ndarray) -> jnp.ndarray:\n    return max_pool(value, self.window_shape, self.strides,\n                    self.padding, self.channel_axis)\n\n\nclass AvgPool(hk.Module):\n  \"\"\"Average pool.\n\n  Equivalent to partial application of :func:`avg_pool`.\n  \"\"\"\n\n  def __init__(\n      self,\n      window_shape: Union[int, Sequence[int]],\n      strides: Union[int, Sequence[int]],\n      padding: str,\n      channel_axis: Optional[int] = -1,\n      name: Optional[str] = None,\n  ):\n    \"\"\"Average pool.\n\n    Args:\n      window_shape: Shape of window to pool over. Same rank as value or ``int``.\n      strides: Strides for the window. Same rank as value or ``int``.\n      padding: Padding algorithm. Either ``VALID`` or ``SAME``.\n      channel_axis: Axis of the spatial channels for which pooling is skipped.\n      name: String name for the module.\n    \"\"\"\n    super().__init__(name=name)\n    self.window_shape = window_shape\n    self.strides = strides\n    self.padding = padding\n    self.channel_axis = channel_axis\n\n  def __call__(self, value: jnp.ndarray) -> jnp.ndarray:\n    return avg_pool(value, self.window_shape, self.strides,\n                    self.padding, self.channel_axis)\n", "path": "haiku/_src/pool.py"}], "after_files": [{"content": "# Copyright 2019 DeepMind Technologies Limited. All Rights Reserved.\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 Haiku modules.\"\"\"\n\nimport types\nfrom typing import Optional, Sequence, Tuple, Union\nimport warnings\n\nfrom haiku._src import module\nfrom jax import lax\nimport jax.numpy as jnp\nimport numpy as np\n\n# If you are forking replace this block with `import haiku as hk`.\nhk = types.ModuleType(\"haiku\")\nhk.Module = module.Module\ndel module\n\n\ndef _infer_shape(\n    x: jnp.ndarray,\n    size: Union[int, Sequence[int]],\n    channel_axis: Optional[int] = -1,\n) -> Tuple[int, ...]:\n  \"\"\"Infer shape for pooling window or strides.\"\"\"\n  if isinstance(size, int):\n    if channel_axis and not 0 <= abs(channel_axis) < x.ndim:\n      raise ValueError(f\"Invalid channel axis {channel_axis} for {x.shape}\")\n    if channel_axis and channel_axis < 0:\n      channel_axis = x.ndim + channel_axis\n    return (1,) + tuple(size if d != channel_axis else 1\n                        for d in range(1, x.ndim))\n  elif len(size) < x.ndim:\n    # Assume additional dimensions are batch dimensions.\n    return (1,) * (x.ndim - len(size)) + tuple(size)\n  else:\n    assert x.ndim == len(size)\n    return tuple(size)\n\n\n_VMAP_SHAPE_INFERENCE_WARNING = (\n    \"When running under vmap, passing an `int` (except for `1`) for \"\n    \"`window_shape` or `strides` will result in the wrong shape being inferred \"\n    \"because the batch dimension is not visible to Haiku. Please update your \"\n    \"code to specify a full unbatched size. \"\n    \"\"\n    \"For example if you had `pool(x, window_shape=3, strides=1)` before, you \"\n    \"should now pass `pool(x, window_shape=(3, 3, 1), strides=1)`. \"\n    \"\"\n    \"Haiku will assume that any additional dimensions in your input are \"\n    \"batch dimensions, and will pad `window_shape` and `strides` accordingly \"\n    \"making your module support both batched and per-example inputs.\"\n)\n\n\ndef _warn_if_unsafe(window_shape, strides):\n  unsafe = lambda size: isinstance(size, int) and size != 1\n  if unsafe(window_shape) or unsafe(strides):\n    warnings.warn(_VMAP_SHAPE_INFERENCE_WARNING, DeprecationWarning)\n\n\ndef max_pool(\n    value: jnp.ndarray,\n    window_shape: Union[int, Sequence[int]],\n    strides: Union[int, Sequence[int]],\n    padding: str,\n    channel_axis: Optional[int] = -1,\n) -> jnp.ndarray:\n  \"\"\"Max pool.\n\n  Args:\n    value: Value to pool.\n    window_shape: Shape of the pooling window, an int or same rank as value.\n    strides: Strides of the pooling window, an int or same rank as value.\n    padding: Padding algorithm. Either ``VALID`` or ``SAME``.\n    channel_axis: Axis of the spatial channels for which pooling is skipped,\n      used to infer ``window_shape`` or ``strides`` if they are an integer.\n\n  Returns:\n    Pooled result. Same rank as value.\n  \"\"\"\n  if padding not in (\"SAME\", \"VALID\"):\n    raise ValueError(f\"Invalid padding '{padding}', must be 'SAME' or 'VALID'.\")\n\n  _warn_if_unsafe(window_shape, strides)\n  window_shape = _infer_shape(value, window_shape, channel_axis)\n  strides = _infer_shape(value, strides, channel_axis)\n\n  return lax.reduce_window(value, -jnp.inf, lax.max, window_shape, strides,\n                           padding)\n\n\ndef avg_pool(\n    value: jnp.ndarray,\n    window_shape: Union[int, Sequence[int]],\n    strides: Union[int, Sequence[int]],\n    padding: str,\n    channel_axis: Optional[int] = -1,\n) -> jnp.ndarray:\n  \"\"\"Average pool.\n\n  Args:\n    value: Value to pool.\n    window_shape: Shape of the pooling window, an int or same rank as value.\n    strides: Strides of the pooling window, an int or same rank as value.\n    padding: Padding algorithm. Either ``VALID`` or ``SAME``.\n    channel_axis: Axis of the spatial channels for which pooling is skipped,\n      used to infer ``window_shape`` or ``strides`` if they are an integer.\n\n  Returns:\n    Pooled result. Same rank as value.\n\n  Raises:\n    ValueError: If the padding is not valid.\n  \"\"\"\n  if padding not in (\"SAME\", \"VALID\"):\n    raise ValueError(f\"Invalid padding '{padding}', must be 'SAME' or 'VALID'.\")\n\n  _warn_if_unsafe(window_shape, strides)\n  window_shape = _infer_shape(value, window_shape, channel_axis)\n  strides = _infer_shape(value, strides, channel_axis)\n\n  reduce_window_args = (0., lax.add, window_shape, strides, padding)\n  pooled = lax.reduce_window(value, *reduce_window_args)\n  if padding == \"VALID\":\n    # Avoid the extra reduce_window.\n    return pooled / np.prod(window_shape)\n  else:\n    # Count the number of valid entries at each input point, then use that for\n    # computing average. Assumes that any two arrays of same shape will be\n    # padded the same. Avoid broadcasting on axis where pooling is skipped. \n    _shape = tuple(vd if wd!=1 else 1 for (vd, wd) in zip(value.shape, window_shape))\n    window_counts = lax.reduce_window(jnp.ones(_shape), *reduce_window_args)\n    return pooled / window_counts\n\n\nclass MaxPool(hk.Module):\n  \"\"\"Max pool.\n\n  Equivalent to partial application of :func:`max_pool`.\n  \"\"\"\n\n  def __init__(\n      self,\n      window_shape: Union[int, Sequence[int]],\n      strides: Union[int, Sequence[int]],\n      padding: str,\n      channel_axis: Optional[int] = -1,\n      name: Optional[str] = None,\n  ):\n    \"\"\"Max pool.\n\n    Args:\n      window_shape: Shape of window to pool over. Same rank as value or ``int``.\n      strides: Strides for the window. Same rank as value or ``int``.\n      padding: Padding algorithm. Either ``VALID`` or ``SAME``.\n      channel_axis: Axis of the spatial channels for which pooling is skipped.\n      name: String name for the module.\n    \"\"\"\n    super().__init__(name=name)\n    self.window_shape = window_shape\n    self.strides = strides\n    self.padding = padding\n    self.channel_axis = channel_axis\n\n  def __call__(self, value: jnp.ndarray) -> jnp.ndarray:\n    return max_pool(value, self.window_shape, self.strides,\n                    self.padding, self.channel_axis)\n\n\nclass AvgPool(hk.Module):\n  \"\"\"Average pool.\n\n  Equivalent to partial application of :func:`avg_pool`.\n  \"\"\"\n\n  def __init__(\n      self,\n      window_shape: Union[int, Sequence[int]],\n      strides: Union[int, Sequence[int]],\n      padding: str,\n      channel_axis: Optional[int] = -1,\n      name: Optional[str] = None,\n  ):\n    \"\"\"Average pool.\n\n    Args:\n      window_shape: Shape of window to pool over. Same rank as value or ``int``.\n      strides: Strides for the window. Same rank as value or ``int``.\n      padding: Padding algorithm. Either ``VALID`` or ``SAME``.\n      channel_axis: Axis of the spatial channels for which pooling is skipped.\n      name: String name for the module.\n    \"\"\"\n    super().__init__(name=name)\n    self.window_shape = window_shape\n    self.strides = strides\n    self.padding = padding\n    self.channel_axis = channel_axis\n\n  def __call__(self, value: jnp.ndarray) -> jnp.ndarray:\n    return avg_pool(value, self.window_shape, self.strides,\n                    self.padding, self.channel_axis)\n", "path": "haiku/_src/pool.py"}]}