type hint fixes for adaptive/learner/average_learner.py

Author: basnijholt

adaptive/learner/average_learner.py

@@ -1,6 +1,5 @@
-from functools import partial
 from math import sqrt
-from typing import Callable, Dict, List, Optional, Tuple, Union
+from typing import Callable, Dict, List, Optional, Tuple
 
 import numpy as np
 
@@ -34,9 +33,9 @@ class AverageLearner(BaseLearner):
 
     def __init__(
         self,
-        function: Union[partial, Callable],
-        atol: None = None,
-        rtol: Optional[Union[int, float]] = None,
+        function: Callable,
+        atol: Optional[float] = None,
+        rtol: Optional[float] = None,
     ) -> None:
         if atol is None and rtol is None:
             raise Exception("At least one of `atol` and `rtol` should be set.")
@@ -75,7 +74,7 @@ def ask(self, n: int, tell_pending: bool = True) -> Tuple[List[int], List[float]
                 self.tell_pending(p)
         return points, loss_improvements
 
-    def tell(self, n: int, value: Union[int, float]) -> None:
+    def tell(self, n: int, value: float) -> None:
         if n in self.data:
             # The point has already been added before.
             return