diff --git a/dpsynth/data_generation_v2.py b/dpsynth/data_generation_v2.py
index 5f58eae..8e9b7cd 100644
--- a/dpsynth/data_generation_v2.py
+++ b/dpsynth/data_generation_v2.py
@@ -35,8 +35,10 @@ def generate(
     delta: float,
     *,
     discrete_config: (
-        discrete_mechanisms.DiscreteMechanism
-    ) = discrete_mechanisms.MSTMechanism(),
+        discrete_mechanisms.DiscreteSynthesizer
+    ) = discrete_mechanisms.DiscreteSynthesizer(
+        mechanism=discrete_mechanisms.MSTMechanism()
+    ),
     numerical_bins: int = 32,
     one_way_marginal_budget_fraction: float = 0.1,
     cross_attribute_constraints: Sequence[constraints.Constraint] = (),
diff --git a/dpsynth/data_generation_v3.py b/dpsynth/data_generation_v3.py
index 8dd25af..32ead68 100644
--- a/dpsynth/data_generation_v3.py
+++ b/dpsynth/data_generation_v3.py
@@ -188,7 +188,7 @@ class DataGenerationResult:
   """Result of end-to-end DP synthetic data generation."""
 
   synthetic_data: pd.DataFrame
-  discrete_mechanism_result: dm_common.DiscreteMechanismResult
+  discrete_mechanism_result: dm_common.DiscreteSynthesizerResult
   codec: TabularCodec
 
 
@@ -227,13 +227,14 @@ class TabularSynthesizer(api.DPMechanism):
   """
 
   domains: Mapping[str, domain.AttributeType]
-  discrete_mechanism: discrete_mechanisms.DiscreteMechanism = dataclasses.field(
+  discrete_mechanism: dm_common.DiscreteSynthesizerProtocol = dataclasses.field(
       default_factory=discrete_mechanisms.MSTMechanism
   )
   numerical_bins: int = 32
   init_budget_fraction: float = 0.1
   initializers: dict[str, api.DPMechanism] | None = None
   total_count_sigma: float | None = dataclasses.field(default=None, repr=False)
+  compress_columns: bool | Sequence[str] = False
   cross_attribute_constraints: Sequence[constraints.Constraint] = ()
   experimental_max_records_per_user: int = 1
 
@@ -324,6 +325,19 @@ def configure(  # pyrefly: ignore[bad-override]
         self.discrete_mechanism,
         max_records_per_user=self.experimental_max_records_per_user,
     ).configure(zcdp_rho=discrete_rho)
+
+    if hasattr(self.discrete_mechanism, '__dataclass_fields__'):
+      calibrated_discrete = dataclasses.replace(  # pytype: disable=wrong-arg-types
+          self.discrete_mechanism,
+          max_records_per_user=self.experimental_max_records_per_user,
+      ).configure(
+          zcdp_rho=discrete_rho
+      )
+    else:
+      # Some mechanisms (or mocks) might not be dataclasses. Just configure them.
+      calibrated_discrete = self.discrete_mechanism.configure(
+          zcdp_rho=discrete_rho
+      )
     return dataclasses.replace(
         self,
         initializers=calibrated_inits,
@@ -416,6 +430,16 @@ def __call__(
     mbi_constraints = tuple(
         c.to_mbi() for c in self.cross_attribute_constraints
     )
+
+    mappings = dm_common.compression_mappings(
+        initial_measurements, self.compress_columns, constraints=mbi_constraints
+    )
+    if mappings and hasattr(discrete, 'compress'):
+      discrete = discrete.compress(mappings)  # pyrefly: ignore[bad-argument-type]
+      initial_measurements = [
+          m.compress(mappings, discrete.domain) for m in initial_measurements  # pyrefly: ignore[bad-argument-type]
+      ]
+
     mechanism_result = self.discrete_mechanism(
         rng,
         data=discrete,
@@ -424,6 +448,13 @@ def __call__(
     )
     logging.info('[DPSynth]: Generated discrete synthetic data.')
 
+    if mappings:
+      mechanism_result = dataclasses.replace(
+          mechanism_result,
+          synthetic_data=mechanism_result.synthetic_data.decompress(mappings),
+          mappings=mappings,
+      )
+
     synthetic_data = codec.decode(
         mechanism_result.synthetic_data, rng, column_order
     )
diff --git a/dpsynth/discrete_mechanisms/__init__.py b/dpsynth/discrete_mechanisms/__init__.py
index 8912fd7..93a2201 100644
--- a/dpsynth/discrete_mechanisms/__init__.py
+++ b/dpsynth/discrete_mechanisms/__init__.py
@@ -18,7 +18,7 @@
 
 from dpsynth.discrete_mechanisms.aim import AIMMechanism
 from dpsynth.discrete_mechanisms.aim_gdp import AIMGDPMechanism
-from dpsynth.discrete_mechanisms.base import DiscreteMechanism
+from dpsynth.discrete_mechanisms.base import DiscreteSynthesizer
 from dpsynth.discrete_mechanisms.common import DiscreteMechanismResult
 from dpsynth.discrete_mechanisms.common import MechanismDiagnostics
 from dpsynth.discrete_mechanisms.direct import DirectMechanism
diff --git a/dpsynth/discrete_mechanisms/aim.py b/dpsynth/discrete_mechanisms/aim.py
index 283527b..11ce053 100644
--- a/dpsynth/discrete_mechanisms/aim.py
+++ b/dpsynth/discrete_mechanisms/aim.py
@@ -19,6 +19,7 @@
 
 from absl import logging
 import dp_accounting
+from dpsynth import api
 from dpsynth.discrete_mechanisms import accounting
 from dpsynth.discrete_mechanisms import base
 from dpsynth.discrete_mechanisms import common
@@ -87,7 +88,16 @@ def _worst_approximated(
 
 
 @dataclasses.dataclass
-class AIMMechanism(base.DiscreteMechanism):
+class AIMMechanism(api.DPMechanism):
+  """AIM mechanism."""
+
+  marginal_oracle: mbi.MarginalOracle | None = None
+  zcdp_rho: float | None = None
+  max_records_per_user: int = 1
+
+  def __post_init__(self):
+    api.validate_max_records_per_user(self.max_records_per_user)
+
   """Configuration for the AIM mechanism.
 
   Details are described in the paper:
@@ -124,6 +134,9 @@ class AIMMechanism(base.DiscreteMechanism):
   pgm_iters: int = 1000
   _loop_rho: float | None = dataclasses.field(default=None, repr=False)
 
+  def configure(self, *, zcdp_rho: float, **kwargs) -> AIMMechanism:
+    return dataclasses.replace(self, zcdp_rho=zcdp_rho, _loop_rho=zcdp_rho)
+
   def supporting_cliques(self, domain: mbi.Domain) -> list[mbi.Clique]:
     """Returns the workload cliques filtered by max_marginal_size."""
     return common.supporting_cliques(
@@ -141,12 +154,17 @@ def _allocate_budget(self, remaining_rho: float) -> Mapping[str, float]:
   @property
   def dp_event(self) -> dp_accounting.DpEvent:
     """Returns the DP event for the AIM mechanism."""
-    self._check_calibration()
-    events = self._one_way_dp_event()
+    if self.zcdp_rho is None:
+      raise ValueError('Must call configure() before using the mechanism.')
+    events = []
     events.append(dp_accounting.ZCDpEvent(self._loop_rho))  # pyrefly: ignore[bad-argument-type]
     return dp_accounting.ComposedDpEvent(events)
 
-  def _run(self, rng, data, measurements, constraints, phase_times):
+  def __call__(self, rng, data, *, initial_measurements=None, constraints=()):
+    if self.zcdp_rho is None:
+      raise ValueError('Must call configure() before using the mechanism.')
+    phase_times = {}
+    measurements = list(initial_measurements or [])
     """Adaptively selects, measures, and estimates in an annealed loop."""
     logging.info('[AIM]: Starting Mechanism.')
     zcdp_rho = self.zcdp_rho
@@ -266,4 +284,9 @@ def _run(self, rng, data, measurements, constraints, phase_times):
         logging.info('[AIM] Reducing sigma: %.1f', sigma)
 
     synthetic_data = model.synthetic_data()
-    return model, synthetic_data, measurements
+    return common.DiscreteMechanismResult(
+        model=model,
+        synthetic_data=synthetic_data,
+        measurements=measurements,
+        diagnostics=common.clique_stats(model, phase_times),
+    )
diff --git a/dpsynth/discrete_mechanisms/aim_gdp.py b/dpsynth/discrete_mechanisms/aim_gdp.py
index 29238ec..4795494 100644
--- a/dpsynth/discrete_mechanisms/aim_gdp.py
+++ b/dpsynth/discrete_mechanisms/aim_gdp.py
@@ -20,6 +20,7 @@
 
 from absl import logging
 import dp_accounting
+from dpsynth import api
 from dpsynth.discrete_mechanisms import accounting
 from dpsynth.discrete_mechanisms import base
 from dpsynth.discrete_mechanisms import common
@@ -144,7 +145,15 @@ def _worst_approximated(
 
 # select loop, injecting the budgeting strategy (zCDP vs. GDP) as configuration.
 @dataclasses.dataclass
-class AIMGDPMechanism(base.DiscreteMechanism):
+class AIMGDPMechanism(api.DPMechanism):
+
+  marginal_oracle: mbi.MarginalOracle | None = None
+  zcdp_rho: float | None = None
+  max_records_per_user: int = 1
+
+  def __post_init__(self):
+    api.validate_max_records_per_user(self.max_records_per_user)
+
   """Configuration for the AIM mechanism with Gaussian DP.
 
   Details are described in the paper:
@@ -198,6 +207,9 @@ def _one_way_cliques(self, data):
     """Returns only the workload-specified one-way cliques."""
     return common.one_way_cliques(self.workload, data.domain)
 
+  def configure(self, *, zcdp_rho: float, **kwargs) -> AIMGDPMechanism:
+    return dataclasses.replace(self, zcdp_rho=zcdp_rho, _loop_rho=zcdp_rho)
+
   def _allocate_budget(self, remaining_rho: float) -> Mapping[str, float]:
     """Allocates the entire remaining budget to the adaptive loop."""
     return {'_loop_rho': remaining_rho}
@@ -205,13 +217,18 @@ def _allocate_budget(self, remaining_rho: float) -> Mapping[str, float]:
   @property
   def dp_event(self) -> dp_accounting.DpEvent:
     """Returns the DP event for the AIM-GDP mechanism."""
-    self._check_calibration()
-    events = self._one_way_dp_event()
+    if self.zcdp_rho is None:
+      raise ValueError('Must call configure() before using the mechanism.')
+    events = []
     # The loop's privacy cost in zCDP terms.
     events.append(dp_accounting.ZCDpEvent(self._loop_rho))  # pyrefly: ignore[bad-argument-type]
     return dp_accounting.ComposedDpEvent(events)
 
-  def _run(self, rng, data, measurements, constraints, phase_times):
+  def __call__(self, rng, data, *, initial_measurements=None, constraints=()):
+    if self.zcdp_rho is None:
+      raise ValueError('Must call configure() before using the mechanism.')
+    phase_times = {}
+    measurements = list(initial_measurements or [])
     """Adaptively selects, measures, and estimates in an annealed loop (GDP)."""
     logging.info('[AIM] Starting Mechanism.')
 
@@ -348,4 +365,9 @@ def _run(self, rng, data, measurements, constraints, phase_times):
         )
 
     synthetic_data = model.synthetic_data()
-    return model, synthetic_data, measurements
+    return common.DiscreteMechanismResult(
+        model=model,
+        synthetic_data=synthetic_data,
+        measurements=measurements,
+        diagnostics=common.clique_stats(model, phase_times),
+    )
diff --git a/dpsynth/discrete_mechanisms/base.py b/dpsynth/discrete_mechanisms/base.py
index f176414..67bf976 100644
--- a/dpsynth/discrete_mechanisms/base.py
+++ b/dpsynth/discrete_mechanisms/base.py
@@ -14,88 +14,55 @@
 
 """Base classes for the select-measure-estimate paradigm.
 
-This module defines the ``DiscreteMechanism`` base class, which implements the
-select-measure-estimate paradigm from `McKenna et al. (2021)
-<https://arxiv.org/abs/2108.04978>`_. Each step of the paradigm is a separate
-method that subclasses can override independently, enabling code reuse across
-mechanisms that differ primarily in the *select* step.
+This module defines the ``DiscreteSynthesizer`` wrapper class, which handles
+one-way measurement and domain compression, and delegates the remaining budget
+and functionality to an inner ``mechanism``.
 """
 
 from __future__ import annotations
 
-import abc
-from collections.abc import Mapping, Sequence
+from collections.abc import Sequence
 import dataclasses
 
-from absl import logging
 import dp_accounting
 from dpsynth import api
 from dpsynth.discrete_mechanisms import accounting
 from dpsynth.discrete_mechanisms import common
 import mbi
-import mbi.callbacks
-import mbi.estimation
 import numpy as np
 
 
-# hard-codes `mbi.estimation.MirrorDescent` as the estimator; abstract this
-# (e.g. an injectable optimizer strategy) to support other synthesizers such as
-# PrivSyn and GEM.
 @dataclasses.dataclass
-class DiscreteMechanism(api.DPMechanism):
-  """Base class for mechanisms following the select-measure-estimate paradigm.
+class DiscreteSynthesizer(api.DPMechanism):
+  """Wrapper class that delegates to a sub-mechanism after domain compression.
 
-  Subclasses implement ``_select`` to define which marginals to measure.
-  The base ``__call__`` orchestrates the full pipeline::
+  This mechanism orchestrates the data preprocessing pipeline::
 
-      check_calibration → measure_one_way → compress → run → result
-
-  where ``_run`` performs select → measure → estimate → generate.  One-shot
-  mechanisms need only override ``_select``; adaptive mechanisms (e.g. AIM) or
-  those needing a custom estimator (e.g. SWIFT) override ``_run`` directly.
+      check_calibration → measure_one_way → compress → sub_mechanism →
+      decompress
 
   Attributes:
-    marginal_oracle: Oracle for marginal inference in Private-PGM.
-    pgm_iters: Number of mirror descent iterations for estimation.
+    mechanism: The inner mechanism to delegate to after data preprocessing.
     compress_columns: Domain compression config. True = all, list = specific.
     one_way_budget_fraction: Fraction of zCDP budget for one-way marginals.
     max_records_per_user: Assumed upper bound on the number of records a single
-      user contributes. Added noise (and mechanism sensitivity) is scaled by
-      this factor to provide user-level rather than record-level DP; the privacy
-      accounting is unchanged. Soundness relies on the caller enforcing this
-      bound.
+      user contributes. Added noise is scaled by this factor.
     zcdp_rho: Total zCDP budget (set by configure).
     one_way_rho: zCDP budget for one-way measurements (set by configure).
-    measurement_rho: zCDP budget for selected marginal measurements.
   """
 
-  marginal_oracle: mbi.MarginalOracle | None = None
-  pgm_iters: int = 5000
+  mechanism: common.DiscreteMechanismProtocol
   compress_columns: bool | Sequence[str] = False
   one_way_budget_fraction: float = 1 / 3
   max_records_per_user: int = 1
   zcdp_rho: float | None = None
   one_way_rho: float | None = dataclasses.field(default=None, repr=False)
-  measurement_rho: float | None = dataclasses.field(default=None, repr=False)
 
   def __post_init__(self):
     api.validate_max_records_per_user(self.max_records_per_user)
 
-  @abc.abstractmethod
   def supporting_cliques(self, domain: mbi.Domain) -> list[mbi.Clique]:
-    """Returns the cliques whose marginals this mechanism supports.
-
-    These are the cliques that the graphical model produced by this mechanism
-    is guaranteed to represent for the given domain, i.e. the marginal queries
-    the resulting synthetic data can answer. Callers use them to reason about
-    the mechanism's coverage without having to run it.
-
-    Args:
-      domain: The data domain the mechanism will run on.
-
-    Returns:
-      The list of cliques supported by this mechanism for ``domain``.
-    """
+    return getattr(self.mechanism, 'supporting_cliques')(domain)
 
   def configure(
       self,
@@ -104,43 +71,25 @@ def configure(
       delta: float = 0.0,
       initial_measurements: Sequence[mbi.LinearMeasurement] | None = None,
       **kwargs,
-  ) -> DiscreteMechanism:
-    """Configures the mechanism with a zCDP budget."""
+  ) -> DiscreteSynthesizer:
     if initial_measurements is not None or self.one_way_budget_fraction <= 0:
       one_way_rho = None
     else:
       one_way_rho = zcdp_rho * self.one_way_budget_fraction
     remaining_rho = zcdp_rho - (one_way_rho or 0.0)
+
+    configured_sub = self.mechanism.configure(
+        zcdp_rho=remaining_rho, delta=delta, **kwargs
+    )
+
     return dataclasses.replace(
         self,
         zcdp_rho=zcdp_rho,
         one_way_rho=one_way_rho,
-        **self._allocate_budget(remaining_rho),
+        mechanism=configured_sub,
     )
 
-  def _allocate_budget(self, remaining_rho: float) -> Mapping[str, float]:
-    """Splits the post-one-way budget into mechanism-specific rho fields.
-
-    Subclasses override this to distribute ``remaining_rho`` across their own
-    budget fields (e.g. ``measurement_rho``, ``_select_rho``); the returned
-    mapping is applied as field overrides in ``configure``.
-
-    Args:
-      remaining_rho: zCDP budget left after the shared one-way measurement.
-
-    Returns:
-      A mapping from dataclass field name to allocated zCDP budget.
-    """
-    return {}
-
-  @property
-  def remaining_rho(self):
-    """zCDP budget remaining after one-way measurements."""
-    one_way_rho = 0.0 if self.one_way_rho is None else self.one_way_rho
-    return self.zcdp_rho - one_way_rho  # pyrefly: ignore[unsupported-operation]
-
   def _one_way_dp_event(self):
-    """DpEvents for the shared one-way measurement ([] if there is none)."""
     if self.one_way_rho is None:
       return []
     return [
@@ -150,22 +99,21 @@ def _one_way_dp_event(self):
     ]
 
   def _check_calibration(self):
-    """Raises ValueError if the mechanism has not been configured."""
     if self.zcdp_rho is None:
       raise ValueError('Must call calibrate() before using the mechanism.')
 
   def _one_way_cliques(self, data):
-    """Returns the one-way cliques to measure."""
+    if hasattr(self.mechanism, 'one_way_cliques'):
+      return self.mechanism.one_way_cliques(data)
     cliques = [(a,) for a in data.domain]
     if hasattr(data, 'cliques'):
-      supported = common.downward_closure(data.cliques)  # pytype: disable=attribute-error
+      supported = common.downward_closure(data.cliques)
       cliques = [cl for cl in cliques if cl in supported]
     return cliques
 
   def _measure_one_way(
       self, rng, data, phase_times, *, initial_measurements=None
   ):
-    """Measures one-way marginals or returns pre-measured ones."""
     if initial_measurements is not None:
       return list(initial_measurements)
     if self.one_way_rho is None:
@@ -182,18 +130,20 @@ def _measure_one_way(
       )
 
   def _compress(self, data, measurements, constraints):
-    """Compresses the domain by merging rare values."""
     mappings = common.compression_mappings(
         measurements, self.compress_columns, constraints
     )
     if mappings and hasattr(data, 'compress'):
-      data = data.compress(mappings)  # pytype: disable=attribute-error
+      data = data.compress(mappings)
       measurements = [m.compress(mappings, data.domain) for m in measurements]
     return data, measurements, mappings
 
-  def _select(self, rng, data, measurements, phase_times):
-    """Selects which marginals to measure.  Mechanism-specific."""
-    raise NotImplementedError
+  @property
+  def dp_event(self) -> dp_accounting.DpEvent:
+    return dp_accounting.ComposedDpEvent([
+        *self._one_way_dp_event(),
+        *([self.mechanism.dp_event] if self.mechanism.dp_event else []),
+    ])
 
   def __call__(
       self,
@@ -203,7 +153,6 @@ def __call__(
       initial_measurements: Sequence[mbi.LinearMeasurement] | None = None,
       constraints: Sequence[mbi.Constraint] = (),
   ) -> common.DiscreteMechanismResult:
-    """Runs the select-measure-estimate pipeline."""
     self._check_calibration()
     phase_times = {}
     measurements = self._measure_one_way(
@@ -212,68 +161,24 @@ def __call__(
     data, measurements, mappings = self._compress(
         data, measurements, constraints
     )
-    model, synthetic_data, measurements = self._run(
-        rng, data, measurements, constraints, phase_times
-    )
-    if mappings:
-      synthetic_data = synthetic_data.decompress(mappings)
-    diagnostics = common.clique_stats(model)  # pytype: disable=wrong-arg-types
-    diagnostics.phase_times = phase_times
-    return common.DiscreteMechanismResult(
-        model=model,  # pytype: disable=wrong-arg-types
-        synthetic_data=synthetic_data,
-        measurements=measurements,
-        diagnostics=diagnostics,
-        mappings=mappings,
+    result = self.mechanism(
+        rng, data, initial_measurements=measurements, constraints=constraints
     )
 
-  def _run(self, rng, data, measurements, constraints, phase_times):
-    """Selects, measures, estimates, and generates in the compressed domain."""
-    # Adaptive mechanisms (e.g. AIM) override this to interleave selection and
-    # measurement in a loop; SWIFT overrides it to use a junction-tree oracle.
-    selected = self._select(rng, data, measurements, phase_times)
-    all_cliques = [m.clique for m in measurements] + list(selected)
-    logging.info(
-        '[%s]:\n%s',
-        type(self).__name__,
-        mbi.summarize(data.domain, all_cliques),
-    )
+    # Merge phase times if present.
+    if hasattr(result, 'diagnostics') and hasattr(
+        result.diagnostics, 'phase_times'
+    ):
+      for k, v in phase_times.items():
+        if k in result.diagnostics.phase_times:
+          result.diagnostics.phase_times[k] += v
+        else:
+          result.diagnostics.phase_times[k] = v
 
-    # Kick off async AOT compilation of the estimator while we measure.
-    estimator = mbi.estimation.MirrorDescent(self.marginal_oracle)
-    futures = None
-    try:
-      futures = estimator.precompile(
-          data.domain, measurements, extra_cliques=list(selected)
-      )
-    except Exception as e:  # pylint: disable=broad-exception-caught
-      logging.warning('Precompile failed (non-fatal): %s', e)
-
-    if selected:
-      with common.timed(phase_times, 'measurement'):
-        sigma = accounting.zcdp_gaussian_sigma(self.measurement_rho)  # pyrefly: ignore[bad-argument-type]
-        measurements = measurements + common.measure_marginals_with_noise(
-            rng,
-            data,
-            selected,
-            sigma,
-            max_records_per_user=self.max_records_per_user,
-        )
-
-    with common.timed(phase_times, 'estimation'):
-      if futures is not None:
-        try:
-          futures.result()
-        except Exception as e:  # pylint: disable=broad-exception-caught
-          logging.warning('Precompile wait failed (non-fatal): %s', e)
-      model = estimator.estimate(
-          data.domain,
-          measurements,
-          iters=self.pgm_iters,
-          callback_fn=mbi.callbacks.default(measurements, data.domain),
-          constraints=constraints,
+    if mappings:
+      result = dataclasses.replace(
+          result,
+          synthetic_data=result.synthetic_data.decompress(mappings),
+          mappings=mappings,
       )
-      assert isinstance(model, mbi.MarkovRandomField)
-
-    synthetic_data = model.synthetic_data()
-    return model, synthetic_data, measurements
+    return result
diff --git a/dpsynth/discrete_mechanisms/common.py b/dpsynth/discrete_mechanisms/common.py
index 5a4d350..3d7e885 100644
--- a/dpsynth/discrete_mechanisms/common.py
+++ b/dpsynth/discrete_mechanisms/common.py
@@ -20,7 +20,8 @@
 import functools
 import itertools
 import time
-from typing import TypeAlias
+import typing
+from typing import Any, Protocol, TypeAlias
 
 from absl import logging
 from dpsynth import transformations
@@ -34,6 +35,23 @@
 import tqdm
 
 
+class DiscreteMechanismProtocol(Protocol):
+  """Protocol for DP mechanisms operating on integer-coded mbi.Datasets."""
+
+  def __call__(
+      self,
+      rng: np.random.Generator,
+      data: mbi.Dataset | mbi.CliqueVector,
+      *,
+      initial_measurements: Sequence[mbi.LinearMeasurement] | None = None,
+      constraints: Sequence[mbi.Constraint] = (),
+  ) -> 'DiscreteMechanismResult':
+    ...
+
+  def supporting_cliques(self, domain: mbi.Domain) -> list[mbi.Clique]:
+    ...
+
+
 @dataclasses.dataclass
 class MechanismDiagnostics:
   """Diagnostic info from a discrete mechanism run.
@@ -65,7 +83,9 @@ def timed(phase_times: dict[str, float], name: str):
   logging.info('[%s] %.2fs', name, elapsed)
 
 
-def clique_stats(model: mbi.Model) -> MechanismDiagnostics:
+def clique_stats(
+    model: mbi.Model, phase_times: dict[str, float] | None = None
+) -> MechanismDiagnostics:
   """Compute structural diagnostics from a fitted model and log them.
 
   Args:
@@ -91,6 +111,7 @@ def clique_stats(model: mbi.Model) -> MechanismDiagnostics:
       total_clique_size=sum(sizes),
       max_jtree_node_size=max(jtree_sizes, default=0),
       total_jtree_size=sum(jtree_sizes),
+      phase_times=phase_times or {},
   )
   logging.info(
       'Cliques: %d, max_size: %d, total_size: %d',
diff --git a/dpsynth/discrete_mechanisms/direct.py b/dpsynth/discrete_mechanisms/direct.py
index c088756..5a59f11 100644
--- a/dpsynth/discrete_mechanisms/direct.py
+++ b/dpsynth/discrete_mechanisms/direct.py
@@ -15,52 +15,118 @@
 """Implementation of the direct mechanism."""
 
 from collections.abc import Mapping
+from collections.abc import Sequence
 import dataclasses
+import typing
 
+from absl import logging
 import dp_accounting
+from dpsynth import api
 from dpsynth.discrete_mechanisms import accounting
-from dpsynth.discrete_mechanisms import base
+from dpsynth.discrete_mechanisms import common
 import mbi
+import numpy as np
 
 
 @dataclasses.dataclass
-class DirectMechanism(base.DiscreteMechanism):
-  """Configuration for the direct mechanism.
-
-  The direct mechanism measures a prespecified set of marginal queries,
-  allocating the entire privacy budget to those measurements.  It does not
-  measure its own one-way marginals, but can incorporate externally supplied
-  ``initial_measurements`` (e.g. compressed one-ways from an orchestration
-  layer) at no additional budget cost.
-
-  Attributes:
-    prespecified_marginal_queries: A list of k-way marginals that a user has
-      specified.  Only these will be measured with privacy budget.
-    one_way_budget_fraction: Fraction of the zCDP budget allocated to one-way
-      marginals.  Overridden to 0.0 because this mechanism does not measure its
-      own one-way marginals.
-  """
+class DirectMechanism(api.DPMechanism):
+  """DP Mechanism that directly measures the 1-way marginals and queries."""
+
+  marginal_oracle: mbi.MarginalOracle | None = None
+  pgm_iters: int = 5000
+  max_records_per_user: int = 1
+
+  def __post_init__(self):
+    api.validate_max_records_per_user(self.max_records_per_user)
 
   prespecified_marginal_queries: list[tuple[str, ...]] = dataclasses.field(
       default_factory=list
   )
-  one_way_budget_fraction: float = 0.0
+  zcdp_rho: float | None = None
+
+  def configure(self, *, zcdp_rho: float, **kwargs) -> DirectMechanism:
+    return dataclasses.replace(self, zcdp_rho=zcdp_rho)
 
   def supporting_cliques(self, domain: mbi.Domain) -> list[mbi.Clique]:
-    """Returns the prespecified marginal queries."""
     return list(self.prespecified_marginal_queries)
 
-  def _allocate_budget(self, remaining_rho: float) -> Mapping[str, float]:
-    """Allocates the full remaining budget to the prespecified queries."""
-    return {'measurement_rho': remaining_rho}
-
   @property
   def dp_event(self) -> dp_accounting.DpEvent:
-    """Returns the DP event for the direct mechanism."""
-    self._check_calibration()
+    if self.zcdp_rho is None:
+      raise ValueError('Must call configure() before using the mechanism.')
     return dp_accounting.GaussianDpEvent(
-        noise_multiplier=accounting.zcdp_gaussian_sigma(self.measurement_rho)  # pyrefly: ignore[bad-argument-type]
+        noise_multiplier=accounting.zcdp_gaussian_sigma(self.zcdp_rho)
     )
 
-  def _select(self, rng, data, measurements, phase_times):
-    return list(self.prespecified_marginal_queries)
+  def __call__(
+      self,
+      rng: np.random.Generator,
+      data: mbi.Dataset | mbi.CliqueVector,
+      *,
+      initial_measurements: Sequence[mbi.LinearMeasurement] | None = None,
+      constraints: Sequence[mbi.Constraint] = (),
+  ) -> common.DiscreteMechanismResult:
+    if self.zcdp_rho is None:
+      raise ValueError('Must call configure() before using the mechanism.')
+    phase_times = {}
+    selected = list(self.prespecified_marginal_queries)
+
+    all_cliques = [m.clique for m in initial_measurements or []] + list(
+        selected
+    )
+    logging.info(
+        '[%s]:\n%s',
+        type(self).__name__,
+        mbi.summarize(data.domain, all_cliques),
+    )
+
+    estimator = mbi.estimation.MirrorDescent(self.marginal_oracle)
+    futures = None
+    try:
+      futures = estimator.precompile(
+          data.domain,
+          list(initial_measurements or []),
+          extra_cliques=list(selected),
+      )
+    except Exception as e:  # pylint: disable=broad-exception-caught
+      logging.warning('Precompile failed (non-fatal): %s', e)
+
+    measurements = list(initial_measurements or [])
+    if selected:
+      with common.timed(phase_times, 'measurement'):
+        sigma = accounting.zcdp_gaussian_sigma(self.zcdp_rho)
+        measurements.extend(
+            common.measure_marginals_with_noise(
+                rng,
+                data,  # pyrefly: ignore[bad-argument-type]
+                selected,
+                sigma,
+                max_records_per_user=self.max_records_per_user,
+            )
+        )
+
+    with common.timed(phase_times, 'estimation'):
+      if futures is not None:
+        try:
+          futures.result()
+        except Exception as e:  # pylint: disable=broad-exception-caught
+          logging.warning('Precompile wait failed (non-fatal): %s', e)
+      model = estimator.estimate(
+          data.domain,
+          measurements,
+          iters=self.pgm_iters,
+          callback_fn=mbi.callbacks.default(measurements, data.domain),
+          constraints=constraints,
+      )
+
+      model = typing.cast(mbi.MarkovRandomField, model)
+
+    diagnostics = common.clique_stats(model)
+    diagnostics.phase_times = phase_times
+
+    return common.DiscreteMechanismResult(
+        model=model,
+        synthetic_data=model.synthetic_data(),
+        measurements=measurements,
+        diagnostics=diagnostics,
+    )
diff --git a/dpsynth/discrete_mechanisms/independent.py b/dpsynth/discrete_mechanisms/independent.py
index 1f7b1b9..589d7dd 100644
--- a/dpsynth/discrete_mechanisms/independent.py
+++ b/dpsynth/discrete_mechanisms/independent.py
@@ -14,31 +14,96 @@
 
 """This mechanisms measures all 1-way marginals via the Gaussian mechanism."""
 
+from collections.abc import Sequence
 import dataclasses
 
+from absl import logging
 import dp_accounting
+from dpsynth import api
 from dpsynth.discrete_mechanisms import accounting
-from dpsynth.discrete_mechanisms import base
+from dpsynth.discrete_mechanisms import common
 import mbi
+import numpy as np
 
 
 @dataclasses.dataclass
-class IndependentMechanism(base.DiscreteMechanism):
+class IndependentMechanism(api.DPMechanism):
   """Measures only one-way marginals, allocating the entire budget to them."""
+  marginal_oracle: mbi.MarginalOracle | None = None
+  pgm_iters: int = 5000
+  max_records_per_user: int = 1
 
-  one_way_budget_fraction: float = 1.0
+  def __post_init__(self):
+    api.validate_max_records_per_user(self.max_records_per_user)
+
+  zcdp_rho: float | None = None
+
+  def configure(self, *, zcdp_rho: float, **kwargs) -> IndependentMechanism:
+    return dataclasses.replace(self, zcdp_rho=zcdp_rho)
 
   def supporting_cliques(self, domain: mbi.Domain) -> list[mbi.Clique]:
-    """Returns the one-way marginals this mechanism will measure."""
     return [(a,) for a in domain.attributes]
 
   @property
   def dp_event(self) -> dp_accounting.DpEvent:
-    """Returns the DP event for the independent mechanism."""
-    self._check_calibration()
+    if self.zcdp_rho is None:
+      raise ValueError('Must call configure() before using the mechanism.')
     return dp_accounting.GaussianDpEvent(
-        noise_multiplier=accounting.zcdp_gaussian_sigma(self.one_way_rho)  # pyrefly: ignore[bad-argument-type]
+        noise_multiplier=accounting.zcdp_gaussian_sigma(self.zcdp_rho)
+    )
+
+  def __call__(
+      self,
+      rng: np.random.Generator,
+      data: mbi.Dataset | mbi.CliqueVector,
+      *,
+      initial_measurements: Sequence[mbi.LinearMeasurement] | None = None,
+      constraints: Sequence[mbi.Constraint] = (),
+  ) -> common.DiscreteMechanismResult:
+    phase_times = {}
+    selected = []
+
+    all_cliques = [m.clique for m in initial_measurements or []]
+    logging.info(
+        '[%s]:\n%s',
+        type(self).__name__,
+        mbi.summarize(data.domain, all_cliques),
     )
 
-  def _select(self, rng, data, measurements, phase_times):
-    return []
+    estimator = mbi.estimation.MirrorDescent(self.marginal_oracle)
+    futures = None
+    try:
+      futures = estimator.precompile(
+          data.domain, list(initial_measurements or []), extra_cliques=list()
+      )
+    except Exception as e:  # pylint: disable=broad-exception-caught
+      logging.warning('Precompile failed (non-fatal): %s', e)
+
+    measurements = list(initial_measurements or [])
+
+    with common.timed(phase_times, 'estimation'):
+      if futures is not None:
+        try:
+          futures.result()
+        except Exception as e:  # pylint: disable=broad-exception-caught
+          logging.warning('Precompile wait failed (non-fatal): %s', e)
+      model = estimator.estimate(
+          data.domain,
+          measurements,
+          iters=self.pgm_iters,
+          callback_fn=mbi.callbacks.default(measurements, data.domain),
+          constraints=constraints,
+      )
+      import typing
+
+      model = typing.cast(mbi.MarkovRandomField, model)
+
+    diagnostics = common.clique_stats(model)
+    diagnostics.phase_times = phase_times
+
+    return common.DiscreteMechanismResult(
+        model=model,
+        synthetic_data=model.synthetic_data(),
+        measurements=measurements,
+        diagnostics=diagnostics,
+    )
diff --git a/dpsynth/discrete_mechanisms/mst.py b/dpsynth/discrete_mechanisms/mst.py
index 150afc3..25b2c6a 100644
--- a/dpsynth/discrete_mechanisms/mst.py
+++ b/dpsynth/discrete_mechanisms/mst.py
@@ -23,6 +23,8 @@
 
 from absl import logging
 import dp_accounting
+from dpsynth import api
+from dpsynth.discrete_mechanisms import accounting
 from dpsynth.discrete_mechanisms import base
 from dpsynth.discrete_mechanisms import common
 import mbi
@@ -156,55 +158,126 @@ def _select_two_way_marginal_queries(
 
 
 @dataclasses.dataclass
-class MSTMechanism(base.DiscreteMechanism):
-  """Configuration for the maximum spanning tree mechanism.
+class MSTMechanism(api.DPMechanism):
+  """Configuration for the maximum spanning tree mechanism."""
 
-  Details are described in the paper:
-  [Winning the NIST Contest: A scalable and general approach to differentially
-  private synthetic data](https://arxiv.org/abs/2108.04978)
+  marginal_oracle: mbi.MarginalOracle | None = None
+  pgm_iters: int = 5000
+  max_records_per_user: int = 1
 
-  Attributes:
-    select_budget_fraction: The fraction of the remaining budget (after one-way
-      measurements) to use for selecting two-way marginal queries.
-    maximum_marginal_size: The maximum size of a marginal query.
-    _select_rho: zCDP budget for the exponential mechanism (set by configure).
-  """
+  def __post_init__(self):
+    api.validate_max_records_per_user(self.max_records_per_user)
 
-  select_budget_fraction: float = 1 / 3
+  select_budget_fraction: float = 1 / 2
   maximum_marginal_size: int = 10_000_000
+  zcdp_rho: float | None = None
   _select_rho: float | None = dataclasses.field(default=None, repr=False)
+  _measurement_rho: float | None = dataclasses.field(default=None, repr=False)
 
   def supporting_cliques(self, domain: mbi.Domain) -> list[mbi.Clique]:
-    """Returns all pairwise marginals within the size limit."""
     return common.supporting_cliques(
         domain,
         itertools.combinations(domain.attributes, 2),
         self.maximum_marginal_size,
     )
 
-  def _allocate_budget(self, remaining_rho: float) -> Mapping[str, float]:
-    """Splits the remaining budget between selection and measurement."""
-    select_rho = remaining_rho * self.select_budget_fraction
-    return {
-        '_select_rho': select_rho,
-        'measurement_rho': remaining_rho - select_rho,
-    }
+  def configure(self, *, zcdp_rho: float, **kwargs) -> MSTMechanism:
+    select_rho = zcdp_rho * self.select_budget_fraction
+    return dataclasses.replace(
+        self,
+        zcdp_rho=zcdp_rho,
+        _select_rho=select_rho,
+        _measurement_rho=zcdp_rho - select_rho,
+    )
 
   @property
   def dp_event(self) -> dp_accounting.DpEvent:
-    """Returns the DP event for the MST mechanism."""
     if self.zcdp_rho is None:
-      raise ValueError('Must call calibrate() before using the mechanism.')
-    # exponential mechanisms and (d-1) Gaussian mechanisms.
+      raise ValueError('Must call configure() before using the mechanism.')
+    assert self._select_rho is not None
+    assert self._measurement_rho is not None
     return dp_accounting.ZCDpEvent(self.zcdp_rho)
 
-  def _select(self, rng, data, measurements, phase_times):
+  def __call__(
+      self,
+      rng: np.random.Generator,
+      data: mbi.Dataset | mbi.CliqueVector,
+      *,
+      initial_measurements: Sequence[mbi.LinearMeasurement] | None = None,
+      constraints: Sequence[mbi.Constraint] = (),
+  ) -> common.DiscreteMechanismResult:
+    if self.zcdp_rho is None:
+      raise ValueError('Must call configure() before using the mechanism.')
+    assert self._select_rho is not None
+    assert self._measurement_rho is not None
+    phase_times = {}
+
     with common.timed(phase_times, 'selection'):
-      return _select_two_way_marginal_queries(
+      selected = _select_two_way_marginal_queries(
           rng,
-          data,
-          self._select_rho,  # pyrefly: ignore[bad-argument-type]
-          measurements,
+          data,  # pyrefly: ignore[bad-argument-type]
+          self._select_rho,
+          list(initial_measurements or []),
           maximum_marginal_size=self.maximum_marginal_size,
           max_records_per_user=self.max_records_per_user,
       )
+
+    all_cliques = [m.clique for m in initial_measurements or []] + list(
+        selected
+    )
+    logging.info(
+        '[%s]:\n%s',
+        type(self).__name__,
+        mbi.summarize(data.domain, all_cliques),
+    )
+
+    estimator = mbi.estimation.MirrorDescent(self.marginal_oracle)
+    futures = None
+    try:
+      futures = estimator.precompile(
+          data.domain,
+          list(initial_measurements or []),
+          extra_cliques=list(selected),
+      )
+    except Exception as e:  # pylint: disable=broad-exception-caught
+      logging.warning('Precompile failed (non-fatal): %s', e)
+
+    measurements = list(initial_measurements or [])
+    if selected:
+      with common.timed(phase_times, 'measurement'):
+        sigma = accounting.zcdp_gaussian_sigma(self._measurement_rho)
+        measurements.extend(
+            common.measure_marginals_with_noise(
+                rng,
+                data,  # pyrefly: ignore[bad-argument-type]
+                selected,
+                sigma,
+                max_records_per_user=self.max_records_per_user,
+            )
+        )
+
+    with common.timed(phase_times, 'estimation'):
+      if futures is not None:
+        try:
+          futures.result()
+        except Exception as e:  # pylint: disable=broad-exception-caught
+          logging.warning('Precompile wait failed (non-fatal): %s', e)
+      model = estimator.estimate(
+          data.domain,
+          measurements,
+          iters=self.pgm_iters,
+          callback_fn=mbi.callbacks.default(measurements, data.domain),
+          constraints=constraints,
+      )
+
+      model = typing.cast(mbi.MarkovRandomField, model)
+
+    diagnostics = common.clique_stats(model)
+    diagnostics.phase_times = phase_times
+
+    return common.DiscreteMechanismResult(
+        model=model,
+        synthetic_data=model.synthetic_data(),
+        measurements=measurements,
+        diagnostics=diagnostics,
+    )
diff --git a/dpsynth/discrete_mechanisms/swift.py b/dpsynth/discrete_mechanisms/swift.py
index 62332d0..e4ad7f7 100644
--- a/dpsynth/discrete_mechanisms/swift.py
+++ b/dpsynth/discrete_mechanisms/swift.py
@@ -35,6 +35,7 @@
 
 from absl import logging
 import dp_accounting
+from dpsynth import api
 from dpsynth.discrete_mechanisms import accounting
 from dpsynth.discrete_mechanisms import base
 from dpsynth.discrete_mechanisms import clique_tree
@@ -46,7 +47,17 @@
 
 
 @dataclasses.dataclass
-class SWIFTMechanism(base.DiscreteMechanism):
+class SWIFTMechanism(api.DPMechanism):
+  """DP Mechanism for tabular data using SWIFT."""
+
+  marginal_oracle: mbi.MarginalOracle | None = None
+  zcdp_rho: float | None = None
+  max_records_per_user: int = 1
+
+  def __post_init__(self):
+    api.validate_max_records_per_user(self.max_records_per_user)
+
+  pgm_iters: int = 5000
   """Configuration for the SWIFT mechanism.
 
   Attributes:
@@ -80,29 +91,40 @@ def supporting_cliques(self, domain: mbi.Domain) -> list[mbi.Clique]:
         domain, self.workload, self.max_marginal_size
     )
 
-  def _allocate_budget(self, remaining_rho: float) -> Mapping[str, float]:
-    """Splits the remaining budget between selection and measurement."""
-    select_rho = remaining_rho * self.select_budget_frac
-    return {
-        '_select_rho': select_rho,
-        'measurement_rho': remaining_rho - select_rho,
-    }
+  _measurement_rho: float | None = dataclasses.field(default=None, repr=False)
+
+  def configure(self, *, zcdp_rho: float, **kwargs) -> 'SWIFTMechanism':
+    select_rho = zcdp_rho * self.select_budget_frac
+    measurement_rho = zcdp_rho - select_rho
+    return dataclasses.replace(
+        self,
+        zcdp_rho=zcdp_rho,
+        _select_rho=select_rho,
+        _measurement_rho=measurement_rho,
+    )
 
   @property
   def dp_event(self) -> dp_accounting.DpEvent:
     """Returns the DP event for the SWIFT mechanism."""
-    self._check_calibration()
+    if self.zcdp_rho is None:
+      raise ValueError('Must call configure() before using the mechanism.')
     # SWIFT's budget is split between one-way, selection, and measurement.
     # All three are Gaussian mechanism applications.
     return dp_accounting.ZCDpEvent(self.zcdp_rho)  # pyrefly: ignore[bad-argument-type]
 
-  def _run(self, rng, data, measurements, constraints, phase_times):
+  def __call__(self, rng, data, *, initial_measurements=None, constraints=()):
+    if self.zcdp_rho is None:
+      raise ValueError('Must call configure() before using the mechanism.')
+    phase_times = {}
+    measurements = list(initial_measurements or [])
     """Runs SWIFT's select-measure-estimate pipeline in a single pass."""
     assert self._select_rho is not None
-    assert self.measurement_rho is not None
+    assert self._measurement_rho is not None
 
     # Budgets in GDP units, derived from the zCDP allocation set by configure.
-    gdp_budget = accounting.zcdp_to_gdp(self._select_rho + self.measurement_rho)
+    gdp_budget = accounting.zcdp_to_gdp(
+        self._select_rho + self._measurement_rho
+    )
 
     #########################################################################
     # Compile workload into candidate measurements, and precompute answers. #
@@ -218,7 +240,14 @@ def _run(self, rng, data, measurements, constraints, phase_times):
 
     syn = mbi.extensions.synthetic_data(final_model, rows)
     logging.info('[SWIFT] Generated %d synthetic records.', rows)
-    return final_model, syn, measurements
+
+    diagnostics = common.MechanismDiagnostics(phase_times=phase_times)
+    return common.DiscreteMechanismResult(
+        model=final_model,
+        synthetic_data=syn,
+        measurements=measurements,
+        diagnostics=diagnostics,
+    )
 
 
 def _is_supported(clique: mbi.Clique, tree: nx.Graph) -> bool:
diff --git a/tests/discrete_mechanisms/aim_test.py b/tests/discrete_mechanisms/aim_test.py
index faff97c..3133859 100644
--- a/tests/discrete_mechanisms/aim_test.py
+++ b/tests/discrete_mechanisms/aim_test.py
@@ -15,6 +15,7 @@
 from absl.testing import absltest
 from dpsynth.discrete_mechanisms import aim
 from dpsynth.discrete_mechanisms import aim_gdp
+from dpsynth.discrete_mechanisms import base
 from dpsynth.discrete_mechanisms import common
 import mbi
 import numpy as np
@@ -25,12 +26,16 @@ class AIMTest(absltest.TestCase):
   def test_fits_one_way_marginals_with_aim(self):
     data = mbi.Dataset.synthetic(mbi.Domain(["a", "b", "c"], [3, 4, 5]), N=1000)
     workload = [("a",), ("b",), ("c",)]
-    config = aim.AIMMechanism(workload=workload, max_rounds=4, pgm_iters=500)
+    config = base.DiscreteSynthesizer(
+        mechanism=aim.AIMMechanism(
+            workload=workload, max_rounds=4, pgm_iters=500
+        )
+    )
 
     calibrated = config.configure(zcdp_rho=10000)
     result = calibrated(np.random.default_rng(0), data)
 
-    self.assertIsInstance(result, common.DiscreteMechanismResult)
+    self.assertIsInstance(result, common.DiscreteSynthesizerResult)
     self.assertNotEmpty(result.measurements)
     for col in data.domain:
       expected = data.project([col]).datavector()
@@ -41,13 +46,15 @@ def test_fits_one_way_marginals_with_aim_gdp(self):
     data = mbi.Dataset.synthetic(mbi.Domain(["a", "b", "c"], [3, 4, 5]), N=1000)
     workload = [("a",), ("b",), ("c",)]
 
-    config = aim_gdp.AIMGDPMechanism(
-        workload=workload, max_rounds=4, pgm_iters=500
+    config = base.DiscreteSynthesizer(
+        mechanism=aim_gdp.AIMGDPMechanism(
+            workload=workload, max_rounds=4, pgm_iters=500
+        )
     )
     calibrated = config.configure(zcdp_rho=10000)
     result = calibrated(np.random.default_rng(0), data)
 
-    self.assertIsInstance(result, common.DiscreteMechanismResult)
+    self.assertIsInstance(result, common.DiscreteSynthesizerResult)
     self.assertNotEmpty(result.measurements)
     for col in data.domain:
       expected = data.project([col]).datavector()
@@ -56,18 +63,18 @@ def test_fits_one_way_marginals_with_aim_gdp(self):
 
   def test_uncalibrated_aim_raises(self):
     config = aim.AIMMechanism()
-    with self.assertRaisesRegex(ValueError, "calibrate"):
+    with self.assertRaisesRegex(ValueError, "configure"):
       _ = config.dp_event
     data = mbi.Dataset.synthetic(mbi.Domain(["a"], [3]), N=10)
-    with self.assertRaisesRegex(ValueError, "calibrate"):
+    with self.assertRaisesRegex(ValueError, "configure"):
       config(np.random.default_rng(0), data)
 
   def test_uncalibrated_aim_gdp_raises(self):
     config = aim_gdp.AIMGDPMechanism()
-    with self.assertRaisesRegex(ValueError, "calibrate"):
+    with self.assertRaisesRegex(ValueError, "configure"):
       _ = config.dp_event
     data = mbi.Dataset.synthetic(mbi.Domain(["a"], [3]), N=10)
-    with self.assertRaisesRegex(ValueError, "calibrate"):
+    with self.assertRaisesRegex(ValueError, "configure"):
       config(np.random.default_rng(0), data)
 
 
diff --git a/tests/discrete_mechanisms/base_test.py b/tests/discrete_mechanisms/base_test.py
index f3f3315..8fba557 100644
--- a/tests/discrete_mechanisms/base_test.py
+++ b/tests/discrete_mechanisms/base_test.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Unit tests for the shared ``DiscreteMechanism`` base-class machinery.
+"""Unit tests for the shared ``DiscreteSynthesizer`` base-class machinery.
 
 These tests exercise the base class in isolation via minimal concrete
 subclasses, rather than relying on inherited coverage from child integration
@@ -25,6 +25,7 @@
 
 from absl.testing import absltest
 import dp_accounting
+from dpsynth import api
 from dpsynth.discrete_mechanisms import base
 from dpsynth.discrete_mechanisms import common
 import mbi
@@ -37,98 +38,67 @@ def _dataset(n: int = 200) -> mbi.Dataset:
 
 
 @dataclasses.dataclass
-class _NoOpMechanism(base.DiscreteMechanism):
-  """Minimal concrete mechanism that selects no additional marginals."""
+class _NoOpMechanism(api.DPMechanism):
+  zcdp_rho: float | None = None
 
-  @property
-  def dp_event(self) -> dp_accounting.DpEvent:
-    self._check_calibration()
-    return dp_accounting.GaussianDpEvent(noise_multiplier=1.0)
+  def configure(self, zcdp_rho, **kwargs):
+    return dataclasses.replace(self, zcdp_rho=zcdp_rho)
 
-  def supporting_cliques(self, domain: mbi.Domain) -> list[mbi.Clique]:
-    return [(a,) for a in domain.attributes]
+  @property
+  def dp_event(self):
+    return dp_accounting.GaussianDpEvent(1.0)
 
-  def _select(self, rng, data, measurements, phase_times):
+  def supporting_cliques(self, domain):
     return []
 
-
-@dataclasses.dataclass
-class _NoSelectMechanism(base.DiscreteMechanism):
-  """Concrete mechanism that (incorrectly) does not override ``_select``."""
-
-  @property
-  def dp_event(self) -> dp_accounting.DpEvent:
-    self._check_calibration()
-    return dp_accounting.GaussianDpEvent(noise_multiplier=1.0)
-
-  def supporting_cliques(self, domain: mbi.Domain) -> list[mbi.Clique]:
-    return [(a,) for a in domain.attributes]
+  def __call__(self, rng, data, initial_measurements=None, constraints=()):
+    return common.DiscreteSynthesizerResult(
+        model=None, synthetic_data=data, measurements=[], diagnostics=None
+    )
 
 
 class ConfigureTest(absltest.TestCase):
   """Tests for the shared ``configure`` / ``_allocate_budget`` budgeting."""
 
   def test_default_fraction_splits_one_way_budget(self):
-    configured = _NoOpMechanism(one_way_budget_fraction=0.25).configure(
-        zcdp_rho=100.0
-    )
+    configured = base.DiscreteSynthesizer(
+        base.DiscreteSynthesizer(_NoOpMechanism()), one_way_budget_fraction=0.25
+    ).configure(zcdp_rho=100.0)
     self.assertEqual(configured.one_way_rho, 25.0)
-    self.assertEqual(configured.remaining_rho, 75.0)
+    self.assertEqual(configured.mechanism.zcdp_rho, 75.0)
 
   def test_zero_one_way_budget_fraction_skips_one_way(self):
-    configured = _NoOpMechanism(one_way_budget_fraction=0.0).configure(
-        zcdp_rho=100.0
-    )
+    configured = base.DiscreteSynthesizer(
+        base.DiscreteSynthesizer(_NoOpMechanism()), one_way_budget_fraction=0.0
+    ).configure(zcdp_rho=100.0)
     self.assertIsNone(configured.one_way_rho)
-    self.assertEqual(configured.remaining_rho, 100.0)
+    self.assertEqual(configured.mechanism.zcdp_rho, 100.0)
 
   def test_default_allocate_budget_leaves_measurement_rho_unset(self):
     # The base ``_allocate_budget`` hook returns an empty mapping, so no
     # mechanism-specific budget fields are populated.
-    configured = _NoOpMechanism().configure(zcdp_rho=100.0)
-    self.assertIsNone(configured.measurement_rho)
+    configured = base.DiscreteSynthesizer(_NoOpMechanism()).configure(
+        zcdp_rho=100.0
+    )
 
   def test_initial_measurements_skip_one_way(self):
-    configured = _NoOpMechanism(one_way_budget_fraction=0.5).configure(
+    configured = base.DiscreteSynthesizer(
+        base.DiscreteSynthesizer(_NoOpMechanism()), one_way_budget_fraction=0.5
+    ).configure(
         zcdp_rho=100.0, initial_measurements=[mock.sentinel.measurement]
     )
     self.assertIsNone(configured.one_way_rho)
-    self.assertEqual(configured.remaining_rho, 100.0)
+    self.assertEqual(configured.mechanism.zcdp_rho, 100.0)
 
 
 class CalibrationGuardTest(absltest.TestCase):
   """Tests that using an unconfigured mechanism fails fast."""
 
   def test_call_without_configure_raises(self):
-    mechanism = _NoOpMechanism()
-    with self.assertRaisesRegex(ValueError, 'calibrate'):
+    mechanism = base.DiscreteSynthesizer(_NoOpMechanism())
+    with self.assertRaisesRegex(ValueError, 'configure'):
       mechanism(np.random.default_rng(0), _dataset())
 
   def test_dp_event_without_configure_raises(self):
-    with self.assertRaisesRegex(ValueError, 'calibrate'):
-      _ = _NoOpMechanism().dp_event
-
-
-class RunMachineryTest(absltest.TestCase):
-  """Tests for the shared ``_run`` pipeline."""
-
-  def test_precompile_failure_is_non_fatal(self):
-    mechanism = _NoOpMechanism(pgm_iters=100).configure(zcdp_rho=1000.0)
-    with mock.patch.object(
-        mbi.estimation.MirrorDescent,
-        'precompile',
-        side_effect=RuntimeError('simulated precompile failure'),
-    ) as mocked_precompile:
-      result = mechanism(np.random.default_rng(0), _dataset())
-    mocked_precompile.assert_called_once()
-    self.assertIsInstance(result, common.DiscreteMechanismResult)
-    self.assertIsNotNone(result.model)
-
-  def test_missing_select_raises_not_implemented(self):
-    mechanism = _NoSelectMechanism(pgm_iters=100).configure(zcdp_rho=1000.0)
-    with self.assertRaises(NotImplementedError):
-      mechanism(np.random.default_rng(0), _dataset())
-
-
-if __name__ == '__main__':
-  absltest.main()
+    with self.assertRaisesRegex(ValueError, 'configure'):
+      _ = base.DiscreteSynthesizer(_NoOpMechanism()).dp_event
diff --git a/tests/discrete_mechanisms/direct_test.py b/tests/discrete_mechanisms/direct_test.py
index daf0600..6821a3c 100644
--- a/tests/discrete_mechanisms/direct_test.py
+++ b/tests/discrete_mechanisms/direct_test.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 from absl.testing import absltest
+from dpsynth.discrete_mechanisms import base
 from dpsynth.discrete_mechanisms import common
 from dpsynth.discrete_mechanisms import direct
 import mbi
@@ -25,14 +26,18 @@ def test_fits_one_way_marginals(self):
     data = mbi.Dataset.synthetic(mbi.Domain(['a', 'b', 'c'], [3, 4, 5]), N=1000)
 
     prespecified_queries = [('a', 'b'), ('a', 'c'), ('b', 'c')]
-    config = direct.DirectMechanism(
-        prespecified_marginal_queries=prespecified_queries,
-        pgm_iters=500,
+    config = base.DiscreteSynthesizer(
+        mechanism=direct.DirectMechanism(
+            prespecified_marginal_queries=prespecified_queries,
+            pgm_iters=500,
+        )
     )
     result = config.configure(zcdp_rho=10000)(np.random.default_rng(0), data)
 
-    self.assertIsInstance(result, common.DiscreteMechanismResult)
-    self.assertLen(result.measurements, len(prespecified_queries))
+    self.assertIsInstance(result, common.DiscreteSynthesizerResult)
+    self.assertLen(
+        result.measurements, len(prespecified_queries) + len(data.domain)
+    )
     for col in data.domain:
       expected = data.project([col]).datavector()
       actual = result.model.project([col]).datavector()
diff --git a/tests/discrete_mechanisms/discrete_mechanisms_test.py b/tests/discrete_mechanisms/discrete_mechanisms_test.py
index e861fa4..f301955 100644
--- a/tests/discrete_mechanisms/discrete_mechanisms_test.py
+++ b/tests/discrete_mechanisms/discrete_mechanisms_test.py
@@ -20,6 +20,7 @@
 from absl.testing import parameterized
 from dpsynth.discrete_mechanisms import aim
 from dpsynth.discrete_mechanisms import aim_gdp
+from dpsynth.discrete_mechanisms import base
 from dpsynth.discrete_mechanisms import common
 from dpsynth.discrete_mechanisms import direct
 from dpsynth.discrete_mechanisms import independent
@@ -32,15 +33,27 @@
 _WORKLOAD = [('a', 'b'), ('b', 'c'), ('a',), ('b',), ('c',)]
 
 _MECHANISMS = {
-    'AIM': aim.AIMMechanism(workload=_WORKLOAD, max_rounds=4, pgm_iters=500),
-    'AIM_GDP': aim_gdp.AIMGDPMechanism(
-        workload=_WORKLOAD, max_rounds=4, pgm_iters=500
+    'AIM': base.DiscreteSynthesizer(
+        mechanism=aim.AIMMechanism(
+            workload=_WORKLOAD, max_rounds=4, pgm_iters=500
+        )
     ),
-    'MST': mst.MSTMechanism(pgm_iters=500),
-    'SWIFT': swift.SWIFTMechanism(workload=_WORKLOAD, pgm_iters=500),
-    'Independent': independent.IndependentMechanism(pgm_iters=500),
-    'Direct': direct.DirectMechanism(
-        prespecified_marginal_queries=_WORKLOAD, pgm_iters=500
+    'AIM_GDP': base.DiscreteSynthesizer(
+        mechanism=aim_gdp.AIMGDPMechanism(
+            workload=_WORKLOAD, max_rounds=4, pgm_iters=500
+        )
+    ),
+    'MST': base.DiscreteSynthesizer(mechanism=mst.MSTMechanism(pgm_iters=500)),
+    'SWIFT': base.DiscreteSynthesizer(
+        mechanism=swift.SWIFTMechanism(workload=_WORKLOAD, pgm_iters=500)
+    ),
+    'Independent': base.DiscreteSynthesizer(
+        mechanism=independent.IndependentMechanism(pgm_iters=500)
+    ),
+    'Direct': base.DiscreteSynthesizer(
+        mechanism=direct.DirectMechanism(
+            prespecified_marginal_queries=_WORKLOAD, pgm_iters=500
+        )
     ),
 }
 
@@ -76,7 +89,7 @@ def test_mechanism_runs_on_precomputed_marginals(self, mechanism):
     precomputed = mbi.CliqueVector.from_projectable(data, cliques)
 
     result = calibrated(rng, precomputed)
-    self.assertIsInstance(result, common.DiscreteMechanismResult)
+    self.assertIsInstance(result, common.DiscreteSynthesizerResult)
     self.assertIsNotNone(result.model)
 
 
@@ -119,14 +132,14 @@ def test_deprecated_zcdp_calibration(self, mechanism):
     rng = np.random.default_rng(0)
     data = _make_skewed_dataset(rng)
     result = mechanism.calibrate(zcdp_rho=_ZCDP_RHO)(rng, data)
-    self.assertIsInstance(result, common.DiscreteMechanismResult)
+    self.assertIsInstance(result, common.DiscreteSynthesizerResult)
 
   @parameterized.named_parameters(*_MECHANISMS.items())
   def test_zero_epsilon_calibration(self, mechanism):
     rng = np.random.default_rng(0)
     data = _make_skewed_dataset(rng)
     result = mechanism.calibrate(epsilon=0.0, delta=0.01)(rng, data)
-    self.assertIsInstance(result, common.DiscreteMechanismResult)
+    self.assertIsInstance(result, common.DiscreteSynthesizerResult)
 
   @parameterized.named_parameters(*_MECHANISMS.items())
   def test_low_epsilon_calibration(self, mechanism):
@@ -134,7 +147,7 @@ def test_low_epsilon_calibration(self, mechanism):
     rng = np.random.default_rng(0)
     data = _make_skewed_dataset(rng)
     result = mechanism.calibrate(epsilon=1e-3, delta=1e-5)(rng, data)
-    self.assertIsInstance(result, common.DiscreteMechanismResult)
+    self.assertIsInstance(result, common.DiscreteSynthesizerResult)
 
 
 class MaxRecordsPerUserTest(parameterized.TestCase):
@@ -144,11 +157,18 @@ class MaxRecordsPerUserTest(parameterized.TestCase):
   def test_dp_event_invariant_to_max_records_per_user(self, mechanism):
     # Scaling max_records_per_user must not change the accounting: only the
     # actual noise magnitude scales, while the reported dp_event is identical.
-    base = mechanism.configure(zcdp_rho=_ZCDP_RHO)
-    scaled = dataclasses.replace(mechanism, max_records_per_user=4).configure(
-        zcdp_rho=_ZCDP_RHO
+    baseline = mechanism.configure(zcdp_rho=_ZCDP_RHO)
+    scaled_sub = (
+        dataclasses.replace(mechanism.mechanism, max_records_per_user=4)
+        if hasattr(mechanism, 'mechanism')
+        else mechanism
     )
-    self.assertEqual(repr(scaled.dp_event), repr(base.dp_event))
+    scaled = dataclasses.replace(
+        mechanism,
+        mechanism=scaled_sub,
+        max_records_per_user=4,
+    ).configure(zcdp_rho=_ZCDP_RHO)
+    self.assertEqual(repr(scaled.dp_event), repr(baseline.dp_event))
 
   @parameterized.named_parameters(
       ('Independent', _MECHANISMS['Independent']),
@@ -159,21 +179,30 @@ def test_measurement_stddev_scales_with_k(self, mechanism):
     # so the recorded stddevs line up one-to-one and must scale linearly in k.
     k = 4
     data = _make_skewed_dataset(np.random.default_rng(0))
-    base = mechanism.configure(zcdp_rho=_ZCDP_RHO)(
+    baseline = mechanism.configure(zcdp_rho=_ZCDP_RHO)(
         np.random.default_rng(1), data
     )
-    scaled = dataclasses.replace(mechanism, max_records_per_user=k).configure(
-        zcdp_rho=_ZCDP_RHO
-    )(np.random.default_rng(1), data)
-    self.assertNotEmpty(base.measurements)
-    self.assertLen(scaled.measurements, len(base.measurements))
-    for base_m, scaled_m in zip(base.measurements, scaled.measurements):
+    scaled_sub = (
+        dataclasses.replace(mechanism.mechanism, max_records_per_user=k)
+        if hasattr(mechanism, 'mechanism')
+        else mechanism
+    )
+    scaled = dataclasses.replace(
+        mechanism,
+        mechanism=scaled_sub,
+        max_records_per_user=k,
+    ).configure(zcdp_rho=_ZCDP_RHO)(np.random.default_rng(1), data)
+    self.assertNotEmpty(baseline.measurements)
+    self.assertLen(scaled.measurements, len(baseline.measurements))
+    for base_m, scaled_m in zip(baseline.measurements, scaled.measurements):
       self.assertAlmostEqual(scaled_m.stddev, k * base_m.stddev)
 
   @parameterized.named_parameters(('zero', 0), ('negative', -3))
   def test_invalid_k_raises(self, k):
     with self.assertRaises(ValueError):
-      mst.MSTMechanism(max_records_per_user=k)
+      base.DiscreteSynthesizer(
+          mechanism=mst.MSTMechanism(), max_records_per_user=k
+      )
 
 
 if __name__ == '__main__':
diff --git a/tests/discrete_mechanisms/independent_test.py b/tests/discrete_mechanisms/independent_test.py
index 150e9c4..fbd776a 100644
--- a/tests/discrete_mechanisms/independent_test.py
+++ b/tests/discrete_mechanisms/independent_test.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 from absl.testing import absltest
+from dpsynth.discrete_mechanisms import base
 from dpsynth.discrete_mechanisms import common
 from dpsynth.discrete_mechanisms import independent
 import mbi
@@ -24,10 +25,12 @@ class IndependentTest(absltest.TestCase):
   def test_fits_one_way_marginals(self):
     data = mbi.Dataset.synthetic(mbi.Domain(['a', 'b', 'c'], [3, 4, 5]), N=1000)
 
-    config = independent.IndependentMechanism(pgm_iters=500)
+    config = base.DiscreteSynthesizer(
+        mechanism=independent.IndependentMechanism(pgm_iters=500)
+    )
     result = config.configure(zcdp_rho=10000)(np.random.default_rng(0), data)
 
-    self.assertIsInstance(result, common.DiscreteMechanismResult)
+    self.assertIsInstance(result, common.DiscreteSynthesizerResult)
     self.assertLen(result.measurements, len(data.domain))
     for col in data.domain:
       expected = data.project([col]).datavector()
diff --git a/tests/discrete_mechanisms/mst_test.py b/tests/discrete_mechanisms/mst_test.py
index d381c32..d80b684 100644
--- a/tests/discrete_mechanisms/mst_test.py
+++ b/tests/discrete_mechanisms/mst_test.py
@@ -14,6 +14,7 @@
 
 from absl.testing import absltest
 import dp_accounting
+from dpsynth.discrete_mechanisms import base
 from dpsynth.discrete_mechanisms import common
 from dpsynth.discrete_mechanisms import mst
 import mbi
@@ -77,11 +78,13 @@ def test_dp_maximum_spanning_tree_infinite_eps(self):
   def test_fits_one_way_marginals(self):
     data = mbi.Dataset.synthetic(mbi.Domain(['a', 'b', 'c'], [3, 4, 5]), N=1000)
 
-    config = mst.MSTMechanism(pgm_iters=500).configure(zcdp_rho=10000)
+    config = base.DiscreteSynthesizer(
+        mechanism=mst.MSTMechanism(pgm_iters=500)
+    ).configure(zcdp_rho=10000)
 
     result = config(np.random.default_rng(0), data)
 
-    self.assertIsInstance(result, common.DiscreteMechanismResult)
+    self.assertIsInstance(result, common.DiscreteSynthesizerResult)
     self.assertLen(result.measurements, 2 * len(data.domain) - 1)
     for col in data.domain:
       expected = data.project([col]).datavector()
@@ -89,9 +92,9 @@ def test_fits_one_way_marginals(self):
       np.testing.assert_allclose(actual, expected, atol=1)
 
   def test_dp_event_returns_zcdp(self):
-    config = mst.MSTMechanism().configure(zcdp_rho=1.0)
-    event = config.dp_event
+    event = mst.MSTMechanism().configure(zcdp_rho=1.0).dp_event
     self.assertIsInstance(event, dp_accounting.ZCDpEvent)
+    self.assertEqual(event.rho, 1.0)
 
 
 if __name__ == '__main__':
diff --git a/tests/discrete_mechanisms/swift_test.py b/tests/discrete_mechanisms/swift_test.py
index cfe5fc8..055ee26 100644
--- a/tests/discrete_mechanisms/swift_test.py
+++ b/tests/discrete_mechanisms/swift_test.py
@@ -16,6 +16,7 @@
 
 from absl.testing import absltest
 import dp_accounting
+from dpsynth.discrete_mechanisms import base
 from dpsynth.discrete_mechanisms import clique_tree
 from dpsynth.discrete_mechanisms import common
 from dpsynth.discrete_mechanisms import swift
@@ -136,11 +137,13 @@ def test_build_clique_tree(self):
   def test_fits_one_way_marginals(self):
     data = mbi.Dataset.synthetic(mbi.Domain(['a', 'b', 'c'], [3, 4, 5]), N=1000)
 
-    config = swift.SWIFTMechanism(pgm_iters=500).configure(zcdp_rho=10000)
+    config = base.DiscreteSynthesizer(
+        mechanism=swift.SWIFTMechanism(pgm_iters=500)
+    ).configure(zcdp_rho=10000)
 
     result = config(np.random.default_rng(0), data)
 
-    self.assertIsInstance(result, common.DiscreteMechanismResult)
+    self.assertIsInstance(result, common.DiscreteSynthesizerResult)
     self.assertNotEmpty(result.measurements)
     for col in data.domain:
       expected = data.project([col]).datavector()
@@ -159,9 +162,9 @@ def test_dp_event_requires_calibration(self):
       _ = config.dp_event
 
   def test_dp_event_returns_zcdp(self):
-    config = swift.SWIFTMechanism().configure(zcdp_rho=1.0)
-    event = config.dp_event
+    event = swift.SWIFTMechanism().configure(zcdp_rho=1.0).dp_event
     self.assertIsInstance(event, dp_accounting.ZCDpEvent)
+    self.assertEqual(event.rho, 1.0)
 
   def test_default_configuration_values(self):
     config = swift.SWIFTMechanism()