From 07c248d86800938a6591fee85a0d852a93f77c11 Mon Sep 17 00:00:00 2001
From: Galen Andrew <galenandrew@google.com>
Date: Fri, 27 Aug 2021 17:33:11 -0700
Subject: [PATCH] Adds NeighboringRelation to Accountant and clarifies
 FixedBatchSample events to be with or without replacement.

PiperOrigin-RevId: 393459878
---
 .../privacy/analysis/dp_event.py              | 12 ++++++++++--
 .../privacy/analysis/privacy_accountant.py    | 19 ++++++++++++++++++-
 2 files changed, 28 insertions(+), 3 deletions(-)

diff --git a/tensorflow_privacy/privacy/analysis/dp_event.py b/tensorflow_privacy/privacy/analysis/dp_event.py
index da027e9..ac45a64 100644
--- a/tensorflow_privacy/privacy/analysis/dp_event.py
+++ b/tensorflow_privacy/privacy/analysis/dp_event.py
@@ -85,8 +85,16 @@ class PoissonSampledDpEvent(DpEvent):
 
 
 @attr.s(frozen=True, slots=True, auto_attribs=True)
-class EqualBatchSampledDpEvent(DpEvent):
-  """An application of sampling exactly `batch_size` records."""
+class FixedBatchSampledWrDpEvent(DpEvent):
+  """Sampling exactly `batch_size` records with replacement."""
+  dataset_size: int
+  batch_size: int
+  event: DpEvent
+
+
+@attr.s(frozen=True, slots=True, auto_attribs=True)
+class FixedBatchSampledWorDpEvent(DpEvent):
+  """Sampling exactly `batch_size` records without replacement."""
   dataset_size: int
   batch_size: int
   event: DpEvent
diff --git a/tensorflow_privacy/privacy/analysis/privacy_accountant.py b/tensorflow_privacy/privacy/analysis/privacy_accountant.py
index 2f1265f..9235156 100644
--- a/tensorflow_privacy/privacy/analysis/privacy_accountant.py
+++ b/tensorflow_privacy/privacy/analysis/privacy_accountant.py
@@ -14,17 +14,34 @@
 """PrivacyAccountant abstract base class."""
 
 import abc
+import enum
 
 from tensorflow_privacy.privacy.dp_event import dp_event
 from tensorflow_privacy.privacy.dp_event import dp_event_builder
 
 
+class NeighboringRelation(enum.Enum):
+  ADD_OR_REMOVE_ONE = 1
+  REPLACE_ONE = 2
+
+
 class PrivacyAccountant(metaclass=abc.ABCMeta):
   """Abstract base class for privacy accountants."""
 
-  def __init__(self):
+  def __init__(self, neighboring_relation: NeighboringRelation):
+    self._neighboring_relation = neighboring_relation
     self._ledger = dp_event_builder.DpEventBuilder()
 
+  @property
+  def neighboring_relation(self) -> NeighboringRelation:
+    """The neighboring relation used by the accountant.
+
+    The neighboring relation is expected to remain constant after
+    initialization. Subclasses should not override this property or change the
+    value of the private attribute.
+    """
+    return self._neighboring_relation
+
   @abc.abstractmethod
   def is_supported(self, event: dp_event.DpEvent) -> bool:
     """Checks whether the `DpEvent` can be processed by this accountant.