From 6497d8255a9be35ef6cb852a4ea3639a626c1c30 Mon Sep 17 00:00:00 2001
From: Iden Kalemaj <ikalemaj@meta.com>
Date: Fri, 20 Dec 2024 04:03:36 -0800
Subject: [PATCH] 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 | 2 +-
 opacus/accountants/prv.py | 7 ++++++-
 opacus/accountants/rdp.py | 5 ++++-
 3 files changed, 11 insertions(+), 3 deletions(-)

diff --git a/opacus/accountants/gdp.py b/opacus/accountants/gdp.py
index c39dc765..8a147ae0 100644
--- a/opacus/accountants/gdp.py
+++ b/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.
 
diff --git a/opacus/accountants/prv.py b/opacus/accountants/prv.py
index ee0e0d6d..b1d09ebf 100644
--- a/opacus/accountants/prv.py
+++ b/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.
diff --git a/opacus/accountants/rdp.py b/opacus/accountants/rdp.py
index a5b9cb03..1697e352 100644
--- a/opacus/accountants/rdp.py
+++ b/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.