We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Improve Early Stopping API
See discussion in #2090

--- 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/early_stopping.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 """Early stopping."""
16 
17 import math
18 from flax import struct
19 
20 
21 class EarlyStopping(struct.PyTreeNode):
22   """Early stopping to avoid overfitting during training.
23 
24   The following example stops training early if the difference between losses
25   recorded in the current epoch and previous epoch is less than 1e-3
26   consecutively for 2 times::
27 
28     early_stop = EarlyStopping(min_delta=1e-3, patience=2)
29     for epoch in range(1, num_epochs+1):
30       rng, input_rng = jax.random.split(rng)
31       optimizer, train_metrics = train_epoch(
32           optimizer, train_ds, config.batch_size, epoch, input_rng)
33       _, early_stop = early_stop.update(train_metrics['loss'])
34       if early_stop.should_stop:
35         print('Met early stopping criteria, breaking...')
36         break
37 
38   Attributes:
39     min_delta: Minimum delta between updates to be considered an
40         improvement.
41     patience: Number of steps of no improvement before stopping.
42     best_metric: Current best metric value.
43     patience_count: Number of steps since last improving update.
44     should_stop: Whether the training loop should stop to avoid
45         overfitting.
46   """
47 
48   min_delta: float = 0
49   patience: int = 0
50   best_metric: float = float('inf')
51   patience_count: int = 0
52   should_stop: bool = False
53 
54   def reset(self):
55     return self.replace(
56         best_metric=float('inf'), patience_count=0, should_stop=False
57     )
58 
59   def update(self, metric):
60     """Update the state based on metric.
61 
62     Returns:
63       A pair (has_improved, early_stop), where `has_improved` is True when there
64       was an improvement greater than `min_delta` from the previous
65       `best_metric` and `early_stop` is the updated `EarlyStop` object.
66     """
67 
68     if (
69         math.isinf(self.best_metric)
70         or self.best_metric - metric > self.min_delta
71     ):
72       return True, self.replace(best_metric=metric, patience_count=0)
73     else:
74       should_stop = self.patience_count >= self.patience or self.should_stop
75       return False, self.replace(
76           patience_count=self.patience_count + 1, should_stop=should_stop
77       )
78
```
--- 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/early_stopping.py b/flax/training/early_stopping.py
--- a/flax/training/early_stopping.py
+++ b/flax/training/early_stopping.py
@@ -30,7 +30,7 @@
       rng, input_rng = jax.random.split(rng)
       optimizer, train_metrics = train_epoch(
           optimizer, train_ds, config.batch_size, epoch, input_rng)
-      _, early_stop = early_stop.update(train_metrics['loss'])
+      early_stop = early_stop.update(train_metrics['loss'])
       if early_stop.should_stop:
         print('Met early stopping criteria, breaking...')
         break
@@ -43,6 +43,8 @@
     patience_count: Number of steps since last improving update.
     should_stop: Whether the training loop should stop to avoid
         overfitting.
+    has_improved: Whether the metric has improved greater or
+      equal to the min_delta in the last `.update` call.
   """
 
   min_delta: float = 0
@@ -50,28 +52,29 @@
   best_metric: float = float('inf')
   patience_count: int = 0
   should_stop: bool = False
+  has_improved: bool = False
 
   def reset(self):
     return self.replace(
-        best_metric=float('inf'), patience_count=0, should_stop=False
+        best_metric=float('inf'), patience_count=0, should_stop=False, has_improved=False
     )
 
   def update(self, metric):
     """Update the state based on metric.
 
     Returns:
-      A pair (has_improved, early_stop), where `has_improved` is True when there
-      was an improvement greater than `min_delta` from the previous
-      `best_metric` and `early_stop` is the updated `EarlyStop` object.
+      The updated EarlyStopping class. The `.has_improved` attribute is True
+      when there was an improvement greater than `min_delta` from the previous
+      `best_metric`.
     """
 
     if (
         math.isinf(self.best_metric)
         or self.best_metric - metric > self.min_delta
     ):
-      return True, self.replace(best_metric=metric, patience_count=0)
+      return self.replace(best_metric=metric, patience_count=0, has_improved=True)
     else:
       should_stop = self.patience_count >= self.patience or self.should_stop
-      return False, self.replace(
-          patience_count=self.patience_count + 1, should_stop=should_stop
+      return self.replace(
+          patience_count=self.patience_count + 1, should_stop=should_stop, has_improved=False
       )

{"golden_diff": "diff --git a/flax/training/early_stopping.py b/flax/training/early_stopping.py\n--- a/flax/training/early_stopping.py\n+++ b/flax/training/early_stopping.py\n@@ -30,7 +30,7 @@\n       rng, input_rng = jax.random.split(rng)\n       optimizer, train_metrics = train_epoch(\n           optimizer, train_ds, config.batch_size, epoch, input_rng)\n-      _, early_stop = early_stop.update(train_metrics['loss'])\n+      early_stop = early_stop.update(train_metrics['loss'])\n       if early_stop.should_stop:\n         print('Met early stopping criteria, breaking...')\n         break\n@@ -43,6 +43,8 @@\n     patience_count: Number of steps since last improving update.\n     should_stop: Whether the training loop should stop to avoid\n         overfitting.\n+    has_improved: Whether the metric has improved greater or\n+      equal to the min_delta in the last `.update` call.\n   \"\"\"\n \n   min_delta: float = 0\n@@ -50,28 +52,29 @@\n   best_metric: float = float('inf')\n   patience_count: int = 0\n   should_stop: bool = False\n+  has_improved: bool = False\n \n   def reset(self):\n     return self.replace(\n-        best_metric=float('inf'), patience_count=0, should_stop=False\n+        best_metric=float('inf'), patience_count=0, should_stop=False, has_improved=False\n     )\n \n   def update(self, metric):\n     \"\"\"Update the state based on metric.\n \n     Returns:\n-      A pair (has_improved, early_stop), where `has_improved` is True when there\n-      was an improvement greater than `min_delta` from the previous\n-      `best_metric` and `early_stop` is the updated `EarlyStop` object.\n+      The updated EarlyStopping class. The `.has_improved` attribute is True\n+      when there was an improvement greater than `min_delta` from the previous\n+      `best_metric`.\n     \"\"\"\n \n     if (\n         math.isinf(self.best_metric)\n         or self.best_metric - metric > self.min_delta\n     ):\n-      return True, self.replace(best_metric=metric, patience_count=0)\n+      return self.replace(best_metric=metric, patience_count=0, has_improved=True)\n     else:\n       should_stop = self.patience_count >= self.patience or self.should_stop\n-      return False, self.replace(\n-          patience_count=self.patience_count + 1, should_stop=should_stop\n+      return self.replace(\n+          patience_count=self.patience_count + 1, should_stop=should_stop, has_improved=False\n       )\n", "issue": "Improve Early Stopping API\nSee discussion in #2090\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\"\"\"Early stopping.\"\"\"\n\nimport math\nfrom flax import struct\n\n\nclass EarlyStopping(struct.PyTreeNode):\n  \"\"\"Early stopping to avoid overfitting during training.\n\n  The following example stops training early if the difference between losses\n  recorded in the current epoch and previous epoch is less than 1e-3\n  consecutively for 2 times::\n\n    early_stop = EarlyStopping(min_delta=1e-3, patience=2)\n    for epoch in range(1, num_epochs+1):\n      rng, input_rng = jax.random.split(rng)\n      optimizer, train_metrics = train_epoch(\n          optimizer, train_ds, config.batch_size, epoch, input_rng)\n      _, early_stop = early_stop.update(train_metrics['loss'])\n      if early_stop.should_stop:\n        print('Met early stopping criteria, breaking...')\n        break\n\n  Attributes:\n    min_delta: Minimum delta between updates to be considered an\n        improvement.\n    patience: Number of steps of no improvement before stopping.\n    best_metric: Current best metric value.\n    patience_count: Number of steps since last improving update.\n    should_stop: Whether the training loop should stop to avoid\n        overfitting.\n  \"\"\"\n\n  min_delta: float = 0\n  patience: int = 0\n  best_metric: float = float('inf')\n  patience_count: int = 0\n  should_stop: bool = False\n\n  def reset(self):\n    return self.replace(\n        best_metric=float('inf'), patience_count=0, should_stop=False\n    )\n\n  def update(self, metric):\n    \"\"\"Update the state based on metric.\n\n    Returns:\n      A pair (has_improved, early_stop), where `has_improved` is True when there\n      was an improvement greater than `min_delta` from the previous\n      `best_metric` and `early_stop` is the updated `EarlyStop` object.\n    \"\"\"\n\n    if (\n        math.isinf(self.best_metric)\n        or self.best_metric - metric > self.min_delta\n    ):\n      return True, self.replace(best_metric=metric, patience_count=0)\n    else:\n      should_stop = self.patience_count >= self.patience or self.should_stop\n      return False, self.replace(\n          patience_count=self.patience_count + 1, should_stop=should_stop\n      )\n", "path": "flax/training/early_stopping.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\"\"\"Early stopping.\"\"\"\n\nimport math\nfrom flax import struct\n\n\nclass EarlyStopping(struct.PyTreeNode):\n  \"\"\"Early stopping to avoid overfitting during training.\n\n  The following example stops training early if the difference between losses\n  recorded in the current epoch and previous epoch is less than 1e-3\n  consecutively for 2 times::\n\n    early_stop = EarlyStopping(min_delta=1e-3, patience=2)\n    for epoch in range(1, num_epochs+1):\n      rng, input_rng = jax.random.split(rng)\n      optimizer, train_metrics = train_epoch(\n          optimizer, train_ds, config.batch_size, epoch, input_rng)\n      early_stop = early_stop.update(train_metrics['loss'])\n      if early_stop.should_stop:\n        print('Met early stopping criteria, breaking...')\n        break\n\n  Attributes:\n    min_delta: Minimum delta between updates to be considered an\n        improvement.\n    patience: Number of steps of no improvement before stopping.\n    best_metric: Current best metric value.\n    patience_count: Number of steps since last improving update.\n    should_stop: Whether the training loop should stop to avoid\n        overfitting.\n    has_improved: Whether the metric has improved greater or\n      equal to the min_delta in the last `.update` call.\n  \"\"\"\n\n  min_delta: float = 0\n  patience: int = 0\n  best_metric: float = float('inf')\n  patience_count: int = 0\n  should_stop: bool = False\n  has_improved: bool = False\n\n  def reset(self):\n    return self.replace(\n        best_metric=float('inf'), patience_count=0, should_stop=False, has_improved=False\n    )\n\n  def update(self, metric):\n    \"\"\"Update the state based on metric.\n\n    Returns:\n      The updated EarlyStopping class. The `.has_improved` attribute is True\n      when there was an improvement greater than `min_delta` from the previous\n      `best_metric`.\n    \"\"\"\n\n    if (\n        math.isinf(self.best_metric)\n        or self.best_metric - metric > self.min_delta\n    ):\n      return self.replace(best_metric=metric, patience_count=0, has_improved=True)\n    else:\n      should_stop = self.patience_count >= self.patience or self.should_stop\n      return self.replace(\n          patience_count=self.patience_count + 1, should_stop=should_stop, has_improved=False\n      )\n", "path": "flax/training/early_stopping.py"}]}