We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
 threading.Condition.notifyAll has been deprecated in favour of notify_all in Python 3.10
### Problem you have encountered:

 `threading.Condition.notifyAll` has been deprecated in favour of `notify_all` in Python 3.10. Ref : python/cpython#25174

### What you expected to happen:

use `notify_all` in below places.

```
rg -t py -w 'currentThread|notifyAll|activeCount|isDaemon|setDaemon' 
flax/training/prefetch_iterator.py
58:        self._cond.notifyAll()
68:      self._cond.notifyAll()
80:          self._cond.notifyAll()
88:          self._cond.notifyAll()
```

--- 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/training/prefetch_iterator.py`
Content:
```
1 # Copyright 2021 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 """Utility for constructing an iterator which prefetches data asynchronously.
16 """
17 
18 import threading
19 import warnings
20 
21 
22 class PrefetchIterator:
23   """Wraps an iterator to provide async prefetching.
24 
25   DEPRECATION WARNING:
26   TensorFlow datasets no longer require manual prefetching.
27 
28   Previously this class was used to make data loading using TensorFlow datasets
29   more efficient. Now TF data handles prefetching with NumPy iterators
30   correctly.
31 
32   Example::
33 
34     tf_iter = dataset.as_numpy_iterator()  # only loads data while calling next
35     tf_iter = PrefetchIterator(tf_iter)  # prefetches data in the background
36 
37   """
38 
39   def __init__(self, data_iter, buffer_size=1):
40     """Construct a PrefetchIterator.
41 
42     Args:
43       data_iter: the Iterator that should be prefetched.
44       buffer_size: how many items to prefetch (default: 1).
45     """
46     warnings.warn('PrefetchIterator is deprecated. Use the standard `tf.data`'
47                   ' prefetch method instead', DeprecationWarning)
48 
49     self._data_iter = data_iter
50     self.buffer_size = buffer_size
51     self._cond = threading.Condition()
52     self._buffer = []
53     self._active = True
54     self._thread = threading.Thread(target=self._prefetch_loop, daemon=True)
55     self._thread.start()
56     self._error = None
57 
58   def __iter__(self):
59     return self
60 
61   def __next__(self):
62     with self._cond:
63       self._cond.wait_for(lambda: self._buffer or not self._active)
64       if self._buffer:
65         item = self._buffer.pop(0)
66         self._cond.notifyAll()
67         return item
68       if self._error:
69         raise self._error  # pylint: disable=raising-bad-type
70       assert not self._active
71       raise StopIteration()
72 
73   def close(self):
74     with self._cond:
75       self._active = False
76       self._cond.notifyAll()
77 
78   def _prefetch_loop(self):
79     """Prefetch loop that prefetches a tf dataset."""
80     def _predicate():
81       return len(self._buffer) < self.buffer_size or not self._active
82 
83     while True:
84       try:
85         item = next(self._data_iter)
86         with self._cond:
87           self._buffer.append(item)
88           self._cond.notifyAll()
89           self._cond.wait_for(_predicate)
90           if not self._active:
91             return
92       except Exception as e:  # pylint: disable=broad-except
93         with self._cond:
94           self._error = e
95           self._active = False
96           self._cond.notifyAll()
97           return
98
```
--- 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/training/prefetch_iterator.py b/flax/training/prefetch_iterator.py
--- a/flax/training/prefetch_iterator.py
+++ b/flax/training/prefetch_iterator.py
@@ -55,7 +55,7 @@
       self._cond.wait_for(lambda: self._buffer or not self._active)
       if self._buffer:
         item = self._buffer.pop(0)
-        self._cond.notifyAll()
+        self._cond.notify_all()
         return item
       if self._error:
         raise self._error  # pylint: disable=raising-bad-type
@@ -65,7 +65,7 @@
   def close(self):
     with self._cond:
       self._active = False
-      self._cond.notifyAll()
+      self._cond.notify_all()
 
   def _prefetch_loop(self):
     """Prefetch loop that prefetches a tf dataset."""
@@ -77,7 +77,7 @@
         item = next(self._data_iter)
         with self._cond:
           self._buffer.append(item)
-          self._cond.notifyAll()
+          self._cond.notify_all()
           self._cond.wait_for(_predicate)
           if not self._active:
             return
@@ -85,5 +85,5 @@
         with self._cond:
           self._error = e
           self._active = False
-          self._cond.notifyAll()
+          self._cond.notify_all()
           return

{"golden_diff": "diff --git a/flax/training/prefetch_iterator.py b/flax/training/prefetch_iterator.py\n--- a/flax/training/prefetch_iterator.py\n+++ b/flax/training/prefetch_iterator.py\n@@ -55,7 +55,7 @@\n       self._cond.wait_for(lambda: self._buffer or not self._active)\n       if self._buffer:\n         item = self._buffer.pop(0)\n-        self._cond.notifyAll()\n+        self._cond.notify_all()\n         return item\n       if self._error:\n         raise self._error  # pylint: disable=raising-bad-type\n@@ -65,7 +65,7 @@\n   def close(self):\n     with self._cond:\n       self._active = False\n-      self._cond.notifyAll()\n+      self._cond.notify_all()\n \n   def _prefetch_loop(self):\n     \"\"\"Prefetch loop that prefetches a tf dataset.\"\"\"\n@@ -77,7 +77,7 @@\n         item = next(self._data_iter)\n         with self._cond:\n           self._buffer.append(item)\n-          self._cond.notifyAll()\n+          self._cond.notify_all()\n           self._cond.wait_for(_predicate)\n           if not self._active:\n             return\n@@ -85,5 +85,5 @@\n         with self._cond:\n           self._error = e\n           self._active = False\n-          self._cond.notifyAll()\n+          self._cond.notify_all()\n           return\n", "issue": " threading.Condition.notifyAll has been deprecated in favour of notify_all in Python 3.10\n### Problem you have encountered:\r\n\r\n `threading.Condition.notifyAll` has been deprecated in favour of `notify_all` in Python 3.10. Ref : python/cpython#25174\r\n\r\n### What you expected to happen:\r\n\r\nuse `notify_all` in below places.\r\n\r\n```\r\nrg -t py -w 'currentThread|notifyAll|activeCount|isDaemon|setDaemon' \r\nflax/training/prefetch_iterator.py\r\n58:        self._cond.notifyAll()\r\n68:      self._cond.notifyAll()\r\n80:          self._cond.notifyAll()\r\n88:          self._cond.notifyAll()\r\n```\n", "before_files": [{"content": "# Copyright 2021 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\"\"\"Utility for constructing an iterator which prefetches data asynchronously.\n\"\"\"\n\nimport threading\nimport warnings\n\n\nclass PrefetchIterator:\n  \"\"\"Wraps an iterator to provide async prefetching.\n\n  DEPRECATION WARNING:\n  TensorFlow datasets no longer require manual prefetching.\n\n  Previously this class was used to make data loading using TensorFlow datasets\n  more efficient. Now TF data handles prefetching with NumPy iterators\n  correctly.\n\n  Example::\n\n    tf_iter = dataset.as_numpy_iterator()  # only loads data while calling next\n    tf_iter = PrefetchIterator(tf_iter)  # prefetches data in the background\n\n  \"\"\"\n\n  def __init__(self, data_iter, buffer_size=1):\n    \"\"\"Construct a PrefetchIterator.\n\n    Args:\n      data_iter: the Iterator that should be prefetched.\n      buffer_size: how many items to prefetch (default: 1).\n    \"\"\"\n    warnings.warn('PrefetchIterator is deprecated. Use the standard `tf.data`'\n                  ' prefetch method instead', DeprecationWarning)\n\n    self._data_iter = data_iter\n    self.buffer_size = buffer_size\n    self._cond = threading.Condition()\n    self._buffer = []\n    self._active = True\n    self._thread = threading.Thread(target=self._prefetch_loop, daemon=True)\n    self._thread.start()\n    self._error = None\n\n  def __iter__(self):\n    return self\n\n  def __next__(self):\n    with self._cond:\n      self._cond.wait_for(lambda: self._buffer or not self._active)\n      if self._buffer:\n        item = self._buffer.pop(0)\n        self._cond.notifyAll()\n        return item\n      if self._error:\n        raise self._error  # pylint: disable=raising-bad-type\n      assert not self._active\n      raise StopIteration()\n\n  def close(self):\n    with self._cond:\n      self._active = False\n      self._cond.notifyAll()\n\n  def _prefetch_loop(self):\n    \"\"\"Prefetch loop that prefetches a tf dataset.\"\"\"\n    def _predicate():\n      return len(self._buffer) < self.buffer_size or not self._active\n\n    while True:\n      try:\n        item = next(self._data_iter)\n        with self._cond:\n          self._buffer.append(item)\n          self._cond.notifyAll()\n          self._cond.wait_for(_predicate)\n          if not self._active:\n            return\n      except Exception as e:  # pylint: disable=broad-except\n        with self._cond:\n          self._error = e\n          self._active = False\n          self._cond.notifyAll()\n          return\n", "path": "flax/training/prefetch_iterator.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# Lint as: python3\n\"\"\"Utility for constructing an iterator which prefetches data asynchronously.\n\"\"\"\n\nimport threading\n\n\nclass PrefetchIterator:\n  \"\"\"Wraps an iterator to provide async prefetching.\n\n  This class is particularly useful for making data loading using TensorFlow\n  more efficient. Currently, the numpy iterator in TensorFlow will not\n  automatically prefetch data in the background::\n\n    tf_iter = dataset.as_numpy_iterator()  # only loads data while calling next\n    tf_iter = PrefetchIterator(tf_iter)  # prefetches data in the background\n\n  \"\"\"\n\n  def __init__(self, data_iter, buffer_size=1):\n    \"\"\"Construct a PrefetchIterator.\n\n    Args:\n      data_iter: the Iterator that should be prefetched.\n      buffer_size: how many items to prefetch (default: 1).\n    \"\"\"\n    self._data_iter = data_iter\n    self.buffer_size = buffer_size\n    self._cond = threading.Condition()\n    self._buffer = []\n    self._active = True\n    self._thread = threading.Thread(target=self._prefetch_loop, daemon=True)\n    self._thread.start()\n    self._error = None\n\n  def __iter__(self):\n    return self\n\n  def __next__(self):\n    with self._cond:\n      self._cond.wait_for(lambda: self._buffer or not self._active)\n      if self._buffer:\n        item = self._buffer.pop(0)\n        self._cond.notify_all()\n        return item\n      if self._error:\n        raise self._error  # pylint: disable=raising-bad-type\n      assert not self._active\n      raise StopIteration()\n\n  def close(self):\n    with self._cond:\n      self._active = False\n      self._cond.notify_all()\n\n  def _prefetch_loop(self):\n    \"\"\"Prefetch loop that prefetches a tf dataset.\"\"\"\n    def _predicate():\n      return len(self._buffer) < self.buffer_size or not self._active\n\n    while True:\n      try:\n        item = next(self._data_iter)\n        with self._cond:\n          self._buffer.append(item)\n          self._cond.notify_all()\n          self._cond.wait_for(_predicate)\n          if not self._active:\n            return\n      except Exception as e:  # pylint: disable=broad-except\n        with self._cond:\n          self._error = e\n          self._active = False\n          self._cond.notify_all()\n          return\n", "path": "flax/training/prefetch_iterator.py"}]}