prompt: 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()
  
  ```
  
golden_diff: 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
  
verification_info: {"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"}]}