Add *kwargs to get_epsilon

Summary: Add **kwargs as argument to the get_epsilon() function in accountants. It allows for building custom PrivacyEngines that take in additional arguments.

Differential Revision: D67514874

opacus/accountants/gdp.py

@@ -44,7 +44,7 @@ def step(self, *, noise_multiplier: float, sample_rate: float):
         else:
             self.history = [(noise_multiplier, sample_rate, 1)]
 
-    def get_epsilon(self, delta: float, poisson: bool = True) -> float:
+    def get_epsilon(self, delta: float, poisson: bool = True, **kwargs) -> float:
         """
         Return privacy budget (epsilon) expended so far.
 

opacus/accountants/prv.py

@@ -81,7 +81,12 @@ def step(self, *, noise_multiplier: float, sample_rate: float):
             self.history.append((noise_multiplier, sample_rate, 1))
 
     def get_epsilon(
-        self, delta: float, *, eps_error: float = 0.01, delta_error: float = None
+        self,
+        delta: float,
+        *,
+        eps_error: float = 0.01,
+        delta_error: float = None,
+        **kwargs,
     ) -> float:
         """
         Return privacy budget (epsilon) expended so far.

opacus/accountants/rdp.py

@@ -68,7 +68,10 @@ def get_privacy_spent(
         return float(eps), float(best_alpha)
 
     def get_epsilon(
-        self, delta: float, alphas: Optional[List[Union[float, int]]] = None
+        self,
+        delta: float,
+        alphas: Optional[List[Union[float, int]]] = None,
+        **kwargs,
     ):
         """
         Return privacy budget (epsilon) expended so far.