We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
ResNetV1 implementation is actually ResNetV1.5
[Downsampling is performed by the 3x3 convolution](https://github.com/google/flax/blob/3ea6381139a340fa7b29c74a84205d8bad18b489/examples/imagenet/models.py#L74), which means it is actually [ResNetV1.5](https://catalog.ngc.nvidia.com/orgs/nvidia/resources/resnet_50_v1_5_for_pytorch), not ResNetV1 as claimed. (In the original version the first convolution in the block has stride 2: see [original paper](https://arxiv.org/pdf/1512.03385.pdf), the caption to Table 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: `examples/imagenet/models.py`
Content:
```
1 # Copyright 2023 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 """Flax implementation of ResNet V1."""
16 
17 # See issue #620.
18 # pytype: disable=wrong-arg-count
19 
20 from functools import partial
21 from typing import Any, Callable, Sequence, Tuple
22 
23 from flax import linen as nn
24 import jax.numpy as jnp
25 
26 ModuleDef = Any
27 
28 
29 class ResNetBlock(nn.Module):
30   """ResNet block."""
31 
32   filters: int
33   conv: ModuleDef
34   norm: ModuleDef
35   act: Callable
36   strides: Tuple[int, int] = (1, 1)
37 
38   @nn.compact
39   def __call__(
40       self,
41       x,
42   ):
43     residual = x
44     y = self.conv(self.filters, (3, 3), self.strides)(x)
45     y = self.norm()(y)
46     y = self.act(y)
47     y = self.conv(self.filters, (3, 3))(y)
48     y = self.norm(scale_init=nn.initializers.zeros_init())(y)
49 
50     if residual.shape != y.shape:
51       residual = self.conv(
52           self.filters, (1, 1), self.strides, name='conv_proj'
53       )(residual)
54       residual = self.norm(name='norm_proj')(residual)
55 
56     return self.act(residual + y)
57 
58 
59 class BottleneckResNetBlock(nn.Module):
60   """Bottleneck ResNet block."""
61 
62   filters: int
63   conv: ModuleDef
64   norm: ModuleDef
65   act: Callable
66   strides: Tuple[int, int] = (1, 1)
67 
68   @nn.compact
69   def __call__(self, x):
70     residual = x
71     y = self.conv(self.filters, (1, 1))(x)
72     y = self.norm()(y)
73     y = self.act(y)
74     y = self.conv(self.filters, (3, 3), self.strides)(y)
75     y = self.norm()(y)
76     y = self.act(y)
77     y = self.conv(self.filters * 4, (1, 1))(y)
78     y = self.norm(scale_init=nn.initializers.zeros_init())(y)
79 
80     if residual.shape != y.shape:
81       residual = self.conv(
82           self.filters * 4, (1, 1), self.strides, name='conv_proj'
83       )(residual)
84       residual = self.norm(name='norm_proj')(residual)
85 
86     return self.act(residual + y)
87 
88 
89 class ResNet(nn.Module):
90   """ResNetV1."""
91 
92   stage_sizes: Sequence[int]
93   block_cls: ModuleDef
94   num_classes: int
95   num_filters: int = 64
96   dtype: Any = jnp.float32
97   act: Callable = nn.relu
98   conv: ModuleDef = nn.Conv
99 
100   @nn.compact
101   def __call__(self, x, train: bool = True):
102     conv = partial(self.conv, use_bias=False, dtype=self.dtype)
103     norm = partial(
104         nn.BatchNorm,
105         use_running_average=not train,
106         momentum=0.9,
107         epsilon=1e-5,
108         dtype=self.dtype,
109         axis_name='batch',
110     )
111 
112     x = conv(
113         self.num_filters,
114         (7, 7),
115         (2, 2),
116         padding=[(3, 3), (3, 3)],
117         name='conv_init',
118     )(x)
119     x = norm(name='bn_init')(x)
120     x = nn.relu(x)
121     x = nn.max_pool(x, (3, 3), strides=(2, 2), padding='SAME')
122     for i, block_size in enumerate(self.stage_sizes):
123       for j in range(block_size):
124         strides = (2, 2) if i > 0 and j == 0 else (1, 1)
125         x = self.block_cls(
126             self.num_filters * 2**i,
127             strides=strides,
128             conv=conv,
129             norm=norm,
130             act=self.act,
131         )(x)
132     x = jnp.mean(x, axis=(1, 2))
133     x = nn.Dense(self.num_classes, dtype=self.dtype)(x)
134     x = jnp.asarray(x, self.dtype)
135     return x
136 
137 
138 ResNet18 = partial(ResNet, stage_sizes=[2, 2, 2, 2], block_cls=ResNetBlock)
139 ResNet34 = partial(ResNet, stage_sizes=[3, 4, 6, 3], block_cls=ResNetBlock)
140 ResNet50 = partial(
141     ResNet, stage_sizes=[3, 4, 6, 3], block_cls=BottleneckResNetBlock
142 )
143 ResNet101 = partial(
144     ResNet, stage_sizes=[3, 4, 23, 3], block_cls=BottleneckResNetBlock
145 )
146 ResNet152 = partial(
147     ResNet, stage_sizes=[3, 8, 36, 3], block_cls=BottleneckResNetBlock
148 )
149 ResNet200 = partial(
150     ResNet, stage_sizes=[3, 24, 36, 3], block_cls=BottleneckResNetBlock
151 )
152 
153 
154 ResNet18Local = partial(
155     ResNet, stage_sizes=[2, 2, 2, 2], block_cls=ResNetBlock, conv=nn.ConvLocal
156 )
157 
158 
159 # Used for testing only.
160 _ResNet1 = partial(ResNet, stage_sizes=[1], block_cls=ResNetBlock)
161 _ResNet1Local = partial(
162     ResNet, stage_sizes=[1], block_cls=ResNetBlock, conv=nn.ConvLocal
163 )
164
```
--- 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/examples/imagenet/models.py b/examples/imagenet/models.py
--- a/examples/imagenet/models.py
+++ b/examples/imagenet/models.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Flax implementation of ResNet V1."""
+"""Flax implementation of ResNet V1.5."""
 
 # See issue #620.
 # pytype: disable=wrong-arg-count
@@ -87,7 +87,7 @@
 
 
 class ResNet(nn.Module):
-  """ResNetV1."""
+  """ResNetV1.5."""
 
   stage_sizes: Sequence[int]
   block_cls: ModuleDef

{"golden_diff": "diff --git a/examples/imagenet/models.py b/examples/imagenet/models.py\n--- a/examples/imagenet/models.py\n+++ b/examples/imagenet/models.py\n@@ -12,7 +12,7 @@\n # See the License for the specific language governing permissions and\n # limitations under the License.\n \n-\"\"\"Flax implementation of ResNet V1.\"\"\"\n+\"\"\"Flax implementation of ResNet V1.5.\"\"\"\n \n # See issue #620.\n # pytype: disable=wrong-arg-count\n@@ -87,7 +87,7 @@\n \n \n class ResNet(nn.Module):\n-  \"\"\"ResNetV1.\"\"\"\n+  \"\"\"ResNetV1.5.\"\"\"\n \n   stage_sizes: Sequence[int]\n   block_cls: ModuleDef\n", "issue": "ResNetV1 implementation is actually ResNetV1.5\n[Downsampling is performed by the 3x3 convolution](https://github.com/google/flax/blob/3ea6381139a340fa7b29c74a84205d8bad18b489/examples/imagenet/models.py#L74), which means it is actually [ResNetV1.5](https://catalog.ngc.nvidia.com/orgs/nvidia/resources/resnet_50_v1_5_for_pytorch), not ResNetV1 as claimed. (In the original version the first convolution in the block has stride 2: see [original paper](https://arxiv.org/pdf/1512.03385.pdf), the caption to Table 1.)\n", "before_files": [{"content": "# Copyright 2023 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\"\"\"Flax implementation of ResNet V1.\"\"\"\n\n# See issue #620.\n# pytype: disable=wrong-arg-count\n\nfrom functools import partial\nfrom typing import Any, Callable, Sequence, Tuple\n\nfrom flax import linen as nn\nimport jax.numpy as jnp\n\nModuleDef = Any\n\n\nclass ResNetBlock(nn.Module):\n  \"\"\"ResNet block.\"\"\"\n\n  filters: int\n  conv: ModuleDef\n  norm: ModuleDef\n  act: Callable\n  strides: Tuple[int, int] = (1, 1)\n\n  @nn.compact\n  def __call__(\n      self,\n      x,\n  ):\n    residual = x\n    y = self.conv(self.filters, (3, 3), self.strides)(x)\n    y = self.norm()(y)\n    y = self.act(y)\n    y = self.conv(self.filters, (3, 3))(y)\n    y = self.norm(scale_init=nn.initializers.zeros_init())(y)\n\n    if residual.shape != y.shape:\n      residual = self.conv(\n          self.filters, (1, 1), self.strides, name='conv_proj'\n      )(residual)\n      residual = self.norm(name='norm_proj')(residual)\n\n    return self.act(residual + y)\n\n\nclass BottleneckResNetBlock(nn.Module):\n  \"\"\"Bottleneck ResNet block.\"\"\"\n\n  filters: int\n  conv: ModuleDef\n  norm: ModuleDef\n  act: Callable\n  strides: Tuple[int, int] = (1, 1)\n\n  @nn.compact\n  def __call__(self, x):\n    residual = x\n    y = self.conv(self.filters, (1, 1))(x)\n    y = self.norm()(y)\n    y = self.act(y)\n    y = self.conv(self.filters, (3, 3), self.strides)(y)\n    y = self.norm()(y)\n    y = self.act(y)\n    y = self.conv(self.filters * 4, (1, 1))(y)\n    y = self.norm(scale_init=nn.initializers.zeros_init())(y)\n\n    if residual.shape != y.shape:\n      residual = self.conv(\n          self.filters * 4, (1, 1), self.strides, name='conv_proj'\n      )(residual)\n      residual = self.norm(name='norm_proj')(residual)\n\n    return self.act(residual + y)\n\n\nclass ResNet(nn.Module):\n  \"\"\"ResNetV1.\"\"\"\n\n  stage_sizes: Sequence[int]\n  block_cls: ModuleDef\n  num_classes: int\n  num_filters: int = 64\n  dtype: Any = jnp.float32\n  act: Callable = nn.relu\n  conv: ModuleDef = nn.Conv\n\n  @nn.compact\n  def __call__(self, x, train: bool = True):\n    conv = partial(self.conv, use_bias=False, dtype=self.dtype)\n    norm = partial(\n        nn.BatchNorm,\n        use_running_average=not train,\n        momentum=0.9,\n        epsilon=1e-5,\n        dtype=self.dtype,\n        axis_name='batch',\n    )\n\n    x = conv(\n        self.num_filters,\n        (7, 7),\n        (2, 2),\n        padding=[(3, 3), (3, 3)],\n        name='conv_init',\n    )(x)\n    x = norm(name='bn_init')(x)\n    x = nn.relu(x)\n    x = nn.max_pool(x, (3, 3), strides=(2, 2), padding='SAME')\n    for i, block_size in enumerate(self.stage_sizes):\n      for j in range(block_size):\n        strides = (2, 2) if i > 0 and j == 0 else (1, 1)\n        x = self.block_cls(\n            self.num_filters * 2**i,\n            strides=strides,\n            conv=conv,\n            norm=norm,\n            act=self.act,\n        )(x)\n    x = jnp.mean(x, axis=(1, 2))\n    x = nn.Dense(self.num_classes, dtype=self.dtype)(x)\n    x = jnp.asarray(x, self.dtype)\n    return x\n\n\nResNet18 = partial(ResNet, stage_sizes=[2, 2, 2, 2], block_cls=ResNetBlock)\nResNet34 = partial(ResNet, stage_sizes=[3, 4, 6, 3], block_cls=ResNetBlock)\nResNet50 = partial(\n    ResNet, stage_sizes=[3, 4, 6, 3], block_cls=BottleneckResNetBlock\n)\nResNet101 = partial(\n    ResNet, stage_sizes=[3, 4, 23, 3], block_cls=BottleneckResNetBlock\n)\nResNet152 = partial(\n    ResNet, stage_sizes=[3, 8, 36, 3], block_cls=BottleneckResNetBlock\n)\nResNet200 = partial(\n    ResNet, stage_sizes=[3, 24, 36, 3], block_cls=BottleneckResNetBlock\n)\n\n\nResNet18Local = partial(\n    ResNet, stage_sizes=[2, 2, 2, 2], block_cls=ResNetBlock, conv=nn.ConvLocal\n)\n\n\n# Used for testing only.\n_ResNet1 = partial(ResNet, stage_sizes=[1], block_cls=ResNetBlock)\n_ResNet1Local = partial(\n    ResNet, stage_sizes=[1], block_cls=ResNetBlock, conv=nn.ConvLocal\n)\n", "path": "examples/imagenet/models.py"}], "after_files": [{"content": "# Copyright 2023 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\"\"\"Flax implementation of ResNet V1.5.\"\"\"\n\n# See issue #620.\n# pytype: disable=wrong-arg-count\n\nfrom functools import partial\nfrom typing import Any, Callable, Sequence, Tuple\n\nfrom flax import linen as nn\nimport jax.numpy as jnp\n\nModuleDef = Any\n\n\nclass ResNetBlock(nn.Module):\n  \"\"\"ResNet block.\"\"\"\n\n  filters: int\n  conv: ModuleDef\n  norm: ModuleDef\n  act: Callable\n  strides: Tuple[int, int] = (1, 1)\n\n  @nn.compact\n  def __call__(\n      self,\n      x,\n  ):\n    residual = x\n    y = self.conv(self.filters, (3, 3), self.strides)(x)\n    y = self.norm()(y)\n    y = self.act(y)\n    y = self.conv(self.filters, (3, 3))(y)\n    y = self.norm(scale_init=nn.initializers.zeros_init())(y)\n\n    if residual.shape != y.shape:\n      residual = self.conv(\n          self.filters, (1, 1), self.strides, name='conv_proj'\n      )(residual)\n      residual = self.norm(name='norm_proj')(residual)\n\n    return self.act(residual + y)\n\n\nclass BottleneckResNetBlock(nn.Module):\n  \"\"\"Bottleneck ResNet block.\"\"\"\n\n  filters: int\n  conv: ModuleDef\n  norm: ModuleDef\n  act: Callable\n  strides: Tuple[int, int] = (1, 1)\n\n  @nn.compact\n  def __call__(self, x):\n    residual = x\n    y = self.conv(self.filters, (1, 1))(x)\n    y = self.norm()(y)\n    y = self.act(y)\n    y = self.conv(self.filters, (3, 3), self.strides)(y)\n    y = self.norm()(y)\n    y = self.act(y)\n    y = self.conv(self.filters * 4, (1, 1))(y)\n    y = self.norm(scale_init=nn.initializers.zeros_init())(y)\n\n    if residual.shape != y.shape:\n      residual = self.conv(\n          self.filters * 4, (1, 1), self.strides, name='conv_proj'\n      )(residual)\n      residual = self.norm(name='norm_proj')(residual)\n\n    return self.act(residual + y)\n\n\nclass ResNet(nn.Module):\n  \"\"\"ResNetV1.5.\"\"\"\n\n  stage_sizes: Sequence[int]\n  block_cls: ModuleDef\n  num_classes: int\n  num_filters: int = 64\n  dtype: Any = jnp.float32\n  act: Callable = nn.relu\n  conv: ModuleDef = nn.Conv\n\n  @nn.compact\n  def __call__(self, x, train: bool = True):\n    conv = partial(self.conv, use_bias=False, dtype=self.dtype)\n    norm = partial(\n        nn.BatchNorm,\n        use_running_average=not train,\n        momentum=0.9,\n        epsilon=1e-5,\n        dtype=self.dtype,\n        axis_name='batch',\n    )\n\n    x = conv(\n        self.num_filters,\n        (7, 7),\n        (2, 2),\n        padding=[(3, 3), (3, 3)],\n        name='conv_init',\n    )(x)\n    x = norm(name='bn_init')(x)\n    x = nn.relu(x)\n    x = nn.max_pool(x, (3, 3), strides=(2, 2), padding='SAME')\n    for i, block_size in enumerate(self.stage_sizes):\n      for j in range(block_size):\n        strides = (2, 2) if i > 0 and j == 0 else (1, 1)\n        x = self.block_cls(\n            self.num_filters * 2**i,\n            strides=strides,\n            conv=conv,\n            norm=norm,\n            act=self.act,\n        )(x)\n    x = jnp.mean(x, axis=(1, 2))\n    x = nn.Dense(self.num_classes, dtype=self.dtype)(x)\n    x = jnp.asarray(x, self.dtype)\n    return x\n\n\nResNet18 = partial(ResNet, stage_sizes=[2, 2, 2, 2], block_cls=ResNetBlock)\nResNet34 = partial(ResNet, stage_sizes=[3, 4, 6, 3], block_cls=ResNetBlock)\nResNet50 = partial(\n    ResNet, stage_sizes=[3, 4, 6, 3], block_cls=BottleneckResNetBlock\n)\nResNet101 = partial(\n    ResNet, stage_sizes=[3, 4, 23, 3], block_cls=BottleneckResNetBlock\n)\nResNet152 = partial(\n    ResNet, stage_sizes=[3, 8, 36, 3], block_cls=BottleneckResNetBlock\n)\nResNet200 = partial(\n    ResNet, stage_sizes=[3, 24, 36, 3], block_cls=BottleneckResNetBlock\n)\n\n\nResNet18Local = partial(\n    ResNet, stage_sizes=[2, 2, 2, 2], block_cls=ResNetBlock, conv=nn.ConvLocal\n)\n\n\n# Used for testing only.\n_ResNet1 = partial(ResNet, stage_sizes=[1], block_cls=ResNetBlock)\n_ResNet1Local = partial(\n    ResNet, stage_sizes=[1], block_cls=ResNetBlock, conv=nn.ConvLocal\n)\n", "path": "examples/imagenet/models.py"}]}