diff --git a/dingo/gw/SVD.py b/dingo/gw/SVD.py
index e77b7761d..80d3a53bc 100644
--- a/dingo/gw/SVD.py
+++ b/dingo/gw/SVD.py
@@ -256,19 +256,22 @@ def __init__(self, svd_basis: SVDBasis, inverse: bool = False):
         self.svd_basis = svd_basis
         self.inverse = inverse
 
-    def __call__(self, waveform: dict):
+    def __call__(self, input_sample: dict):
         """
         Parameters
         ----------
-        waveform : dict
-            Values should be arrays containing waveforms to be transformed.
+        input_sample : dict
+            input_sample["waveform"] should be a dict with arrays containing waveforms
+            to be transformed.
 
         Returns
         -------
         dict of the same form as the input, but with transformed waveforms.
         """
+        sample = input_sample.copy()
         if not self.inverse:
             func = self.svd_basis.compress
         else:
             func = self.svd_basis.decompress
-        return {k: func(v) for k, v in waveform.items()}
+        sample["waveform"] = {k: func(v) for k, v in sample["waveform"].items()}
+        return sample
diff --git a/dingo/gw/dataset/generate_dataset.py b/dingo/gw/dataset/generate_dataset.py
index a52ca9865..7b43a1b1d 100644
--- a/dingo/gw/dataset/generate_dataset.py
+++ b/dingo/gw/dataset/generate_dataset.py
@@ -18,6 +18,7 @@
 from dingo.gw.prior import build_prior_with_defaults
 from dingo.gw.SVD import ApplySVD, SVDBasis
 from dingo.gw.transforms import WhitenFixedASD
+from dingo.gw.transforms.waveform_transforms import HeterodynePhase
 from dingo.gw.waveform_generator import (
     NewInterfaceWaveformGenerator,
     WaveformGenerator,
@@ -183,6 +184,15 @@ def generate_dataset(settings: Dict, num_processes: int) -> WaveformDataset:
                 )
             )
 
+        if "phase_heterodyning" in settings["compression"]:
+            compression_transforms.append(
+                HeterodynePhase(
+                    domain,
+                    inverse=False,
+                    **settings["compression"]["phase_heterodyning"],
+                )
+            )
+
         if "svd" in settings["compression"]:
             svd_settings = settings["compression"]["svd"]
 
diff --git a/dingo/gw/dataset/waveform_dataset.py b/dingo/gw/dataset/waveform_dataset.py
index af750f5be..e3604c3f8 100644
--- a/dingo/gw/dataset/waveform_dataset.py
+++ b/dingo/gw/dataset/waveform_dataset.py
@@ -8,6 +8,7 @@
 from dingo.gw.SVD import SVDBasis, ApplySVD
 from dingo.gw.domains import build_domain
 from dingo.gw.transforms import WhitenFixedASD
+from dingo.gw.transforms.waveform_transforms import HeterodynePhase
 
 
 class WaveformDataset(DingoDataset, torch.utils.data.Dataset):
@@ -196,6 +197,15 @@ def initialize_decompression(self, svd_size_update: Optional[int] = None):
             svd_basis = SVDBasis(dictionary=self.svd)
             decompression_transform_list.append(ApplySVD(svd_basis, inverse=True))
 
+        if "phase_heterodyning" in self.settings["compression"]:
+            decompression_transform_list.append(
+                HeterodynePhase(
+                    self.domain,
+                    inverse=True,
+                    **self.settings["compression"]["phase_heterodyning"],
+                )
+            )
+
         if "whitening" in self.settings["compression"]:
             decompression_transform_list.append(
                 WhitenFixedASD(
@@ -322,13 +332,12 @@ def __getitems__(
                 for pol, waveforms in self.polarizations.items()
             }
 
-        # Decompression transforms are assumed to apply only to the waveform,
-        # and do not involve parameters.
-        if self.decompression_transform is not None:
-            polarizations = self.decompression_transform(polarizations)
-
-        # Main transforms can depend also on parameters.
+        # Decompression transforms operate on sample dicts since some (e.g.,
+        # HeterodynePhase) need access to parameters.
         data = {"parameters": parameters, "waveform": polarizations}
+        if self.decompression_transform is not None:
+            data = self.decompression_transform(data)
+            polarizations = data["waveform"]
         if self.transform is not None:
             data = self.transform(data)
 
diff --git a/dingo/gw/training/train_builders.py b/dingo/gw/training/train_builders.py
index ca62856bb..6cb27c7df 100755
--- a/dingo/gw/training/train_builders.py
+++ b/dingo/gw/training/train_builders.py
@@ -19,6 +19,7 @@
     RepackageStrainsAndASDS,
     UnpackDict,
     GNPECoalescenceTimes,
+    GNPEChirp,
     SampleExtrinsicParameters,
     GetDetectorTimes,
     CropMaskStrainRandom,
@@ -125,6 +126,11 @@ def set_train_transforms(wfd, data_settings, asd_dataset_path, omit_transforms=N
         )
         extra_context_parameters += transforms[-1].context_parameters
 
+    if "gnpe_chirp" in data_settings:
+        d = data_settings["gnpe_chirp"]
+        transforms.append(GNPEChirp(d["kernel"], domain, d.get("order", 0)))
+        extra_context_parameters += transforms[-1].context_parameters
+
     # Add the GNPE context to context_parameters the first time the transforms are
     # constructed. We do not want to overwrite the ordering of the parameters in
     # subsequent runs.
@@ -255,6 +261,7 @@ def build_svd_for_embedding_network(
             SelectStandardizeRepackageParameters,
             UnpackDict,
             CropMaskStrainRandom,
+            GNPEChirp,
         ],
     )
 
diff --git a/dingo/gw/transforms/gnpe_transforms.py b/dingo/gw/transforms/gnpe_transforms.py
index eb4db2ba4..a79e304c1 100644
--- a/dingo/gw/transforms/gnpe_transforms.py
+++ b/dingo/gw/transforms/gnpe_transforms.py
@@ -3,6 +3,8 @@
 from bilby.core.prior import PriorDict
 from abc import ABC, abstractmethod
 
+from dingo.gw.transforms.waveform_transforms import HeterodynePhase
+
 
 class GNPEBase(ABC):
     """
@@ -227,3 +229,111 @@ def __call__(self, input_sample):
         extrinsic_parameters.update(new_parameters)
         sample["extrinsic_parameters"] = extrinsic_parameters
         return sample
+
+
+class GNPEChirp(GNPEBase):
+    """
+    Relative binning / heterodyning GNPE transform, which factors out the overall chirp
+    from the waveform. This is done based on the proxy parameters chirp_mass_proxy and
+    optionally mass_ratio_proxy. These are defined as blurred version of the parameters
+    chirp_mass and mass_ratio.
+
+    At leading order, the data are transformed by dividing by a fiducial waveform of the
+    form
+
+        exp( - 1j * (3/128) * (pi G chirp_mass_proxy f / c**3)**(-5/3) ) ;
+
+    see 2001.11412, eq. (7.2). This is the leading order chirp due to the emission of
+    quadrupole radiation.
+
+    At next to leading order, this transform also optionally implements 1PN corrections
+    involving the mass ratio. We do not include any amplitude in the fiducial waveform,
+    since at inference time this transform will be applied to noisy data. Multiplying
+    the frequency-domain noise by a complex number of unit norm is allowed because it
+    only changes the phase, not the overall amplitude, which would change the noise PSD.
+    """
+
+    def __init__(self, kernel, domain, order: int = 0, inference: bool = False):
+        """
+        Parameters
+        ----------
+        kernel : dict or str
+            Defines a Bilby prior. If a dict, keys should include chirp_mass,
+            and (possibly) mass_ratio.
+        domain : Domain
+            Only works for a FrequencyDomain at present.
+        order : int
+            Twice the post-Newtonian order for the expansion. Valid orders are 0 and 2.
+        inference : bool = False
+            Whether to use inference or training mode.
+        """
+        # We copy the kernel because the PriorDict constructor modifies the argument.
+        kernel = kernel.copy()
+
+        if order == 0:
+            if "chirp_mass" not in kernel:
+                raise KeyError("Kernel must include chirp_mass key.")
+            if "mass_ratio" in kernel:
+                print(
+                    "Warning: mass_ratio kernel provided, but will be ignored for "
+                    "order 0 GNPE."
+                )
+                kernel.pop("mass_ratio")
+        elif order == 2:
+            if "chirp_mass" not in kernel or "mass_ratio" not in kernel:
+                raise KeyError("Kernel must include chirp_mass and mass_ratio keys.")
+            raise NotImplementedError(
+                "Mass ratio conditioning for 2nd order not set up."
+            )
+        else:
+            raise ValueError(f"Order {order} invalid. Acceptable values are 0 and 2.")
+
+        operators = {"chirp_mass": "+", "mass_ratio": "+"}
+        super().__init__(kernel, operators)
+
+        self.inference = inference
+        self.phase_heterodyning_transform = HeterodynePhase(
+            domain, order, inverse=False
+        )
+
+    def __call__(self, input_sample):
+        sample = input_sample.copy()
+
+        extrinsic_parameters = sample["extrinsic_parameters"].copy()
+
+        # If proxies already exist, use them. Otherwise, sample them. Proxies may
+        # already exist if provided by an unconditional initialization network when
+        # attempting to recover the density from GNPE samples, or when using fixed
+        # initialization parameters.
+        if set(self.proxy_list).issubset(extrinsic_parameters.keys()):
+            proxies = {p: extrinsic_parameters[p] for p in self.proxy_list}
+        else:
+            # The relevant parameters could be in either the intrinsic or extrinsic
+            # parameters list. At inference time, we put all GNPE parameters into the
+            # extrinsic parameters list.
+            parameters = {**sample.get("parameters", {}), **extrinsic_parameters}
+            proxies = self.sample_proxies(parameters)
+            extrinsic_parameters.update(proxies)
+            delta_parameters = {
+                "delta_" + p[: -len("_proxy")]: parameters[p[: -len("_proxy")]]
+                - proxies[p]
+                for p in self.proxy_list
+            }
+            extrinsic_parameters.update(delta_parameters)
+            sample["extrinsic_parameters"] = extrinsic_parameters
+
+        # The only situation where we would expect to not have a waveform to transform
+        # would be when calculating parameter standardizations, since we just want to
+        # draw samples of the parameters at that point, and not prepare any data.
+        if "waveform" in sample:
+            sample["waveform"] = self.phase_heterodyning_transform(
+                {
+                    "waveform": sample["waveform"],
+                    "parameters": {
+                        "chirp_mass": proxies["chirp_mass_proxy"],
+                        "mass_ratio": proxies.get("mass_ratio_proxy"),
+                    },
+                }
+            )["waveform"]
+
+        return sample
diff --git a/dingo/gw/transforms/noise_transforms.py b/dingo/gw/transforms/noise_transforms.py
index d864980eb..7f5ad573e 100644
--- a/dingo/gw/transforms/noise_transforms.py
+++ b/dingo/gw/transforms/noise_transforms.py
@@ -104,25 +104,27 @@ def __init__(
 
         self.inverse = inverse
 
-    def __call__(self, sample):
+    def __call__(self, input_sample):
         """
         Parameters
         ----------
-        sample : dict
-            Dictionary of numpy arrays, e.g., with keys corresponding to polarizations.
-            Method whitens each array with the same ASD.
+        input_sample : dict
+            input_sample["waveform"] is a dictionary of numpy arrays, e.g., with keys
+            corresponding to polarizations. Method whitens each array with the same ASD.
 
         Returns
         -------
-        dict of the same form as sample, but with whitened / un-whitened data.
+        dict of the same form as input_sample, but with whitened / un-whitened waveforms.
         """
-        result = {}
-        for k, v in sample.items():
+        sample = input_sample.copy()
+        waveform = {}
+        for k, v in sample["waveform"].items():
             if self.inverse:
-                result[k] = v * self.asd_array
+                waveform[k] = v * self.asd_array
             else:
-                result[k] = v / self.asd_array
-        return result
+                waveform[k] = v / self.asd_array
+        sample["waveform"] = waveform
+        return sample
 
 
 class WhitenAndScaleStrain(object):
diff --git a/dingo/gw/transforms/waveform_transforms.py b/dingo/gw/transforms/waveform_transforms.py
index f6d022435..0df07f71b 100644
--- a/dingo/gw/transforms/waveform_transforms.py
+++ b/dingo/gw/transforms/waveform_transforms.py
@@ -1,5 +1,8 @@
 from typing import Optional
+
+import lal
 import numpy as np
+import torch
 
 from dingo.gw.domains import MultibandedFrequencyDomain, UniformFrequencyDomain
 
@@ -555,3 +558,186 @@ def create_mask_based_on_frequency_update(
             frequency_masks[d] = frequency_masks[d] & mask_max
 
     return frequency_masks
+
+
+def factor_fiducial_waveform(
+    data, domain, chirp_mass, mass_ratio=None, order=0, inverse=False
+):
+    """
+    Relative binning / heterodyning. Divides the data by a fiducial waveform defined by
+    the chirp mass and (optionally) mass ratio. Allows for batching.
+
+    At leading order, this factors out the overall chirp by dividing the data by a
+    fiducial waveform of the form
+        exp( - 1j * (3/128) * (pi G chirp_mass f / c**3)**(-5/3) ) ;
+    see 2001.11412, eq. (7.2). This is the leading order chirp due to the emission of
+    quadrupole radiation.
+
+    At next-to-leading order, this implements 1PN corrections involving the mass ratio.
+
+    We do not include any amplitude in the fiducial waveform, since at inference time
+    this transformation will be applied to noisy data. Multiplying the frequency-domain
+    noise by a complex number of unit norm is allowed because it only changes the
+    phase, not the overall amplitude, which would change the noise PSD.
+
+    Parameters
+    ----------
+    data : Union[dict, torch.Tensor]
+        If a dict, the keys would correspond to different detectors or
+        polarizations. For a Tensor, these would be within different components.
+        This method uses the same fiducial waveform for each detector.
+    domain : Domain
+        Only works for a FrequencyDomain or MultibandedFrequencyDomain at present.
+    chirp_mass : Union[np.array, torch.Tensor, float]
+        Chirp mass parameter(s).
+    mass_ratio : Union[np.array, torch.Tensor, float]
+        Mass ratio parameter(s).
+    order : int
+        Twice the post-Newtonian order for the expansion. Valid orders are 0 and 2.
+    inverse : bool
+        If True, apply the inverse transformation (restore chirp phase).
+
+    Returns
+    -------
+    data: Union[dict, torch.Tensor]
+        Transformed data, of the same form as data.
+    """
+    if order not in (0, 2):
+        raise ValueError(f"Order {order} invalid. Acceptable values are 0 and 2.")
+    if order == 2:
+        if mass_ratio is None:
+            raise ValueError("Mass ratio required for 2nd order heterodyning.")
+        elif (
+            isinstance(chirp_mass, (np.ndarray, torch.Tensor))
+            and chirp_mass.shape != mass_ratio.shape
+        ):
+            raise ValueError(
+                f"Shape of chirp_mass ({chirp_mass.shape}) and mass_ratio "
+                f"({mass_ratio.shape}) don't match"
+            )
+    if not np.isfinite(chirp_mass).all():
+        raise ValueError("Got nan or inf elements in chirp_mass.")
+
+    if type(data) == dict:
+        f = domain.get_sample_frequencies_astype(list(data.values())[0])
+    else:
+        f = domain.get_sample_frequencies_astype(data)
+
+    # Expand across possible batch dimension.
+    if type(chirp_mass) == np.float64 or type(chirp_mass) == float:
+        # Use np.outer.squeeze instead of chirp_mass * f, as the latter has a
+        # different precision, and we don't want the behaviour of this function to
+        # depend on whether chirp_mass is an array or float.
+        mc_f = np.outer(chirp_mass, f).squeeze()
+        if f[0] == 0.0:
+            mc_f[0] = 1.0
+    elif type(chirp_mass) == np.ndarray:
+        mc_f = np.outer(chirp_mass, f)
+        if mass_ratio is not None:
+            mass_ratio = mass_ratio[:, None]
+    elif type(chirp_mass) == torch.Tensor:
+        if chirp_mass.dim() == 0:
+            mc_f = chirp_mass * f  # scalar tensor, no need for outer
+        else:
+            mc_f = torch.outer(chirp_mass, f)
+            if mass_ratio is not None:
+                mass_ratio = mass_ratio[:, None]
+    else:
+        raise TypeError(
+            f"Invalid type {type(chirp_mass)}. "
+            f"Only implemented for floats, arrays and tensors"
+        )
+
+    # Avoid taking a negative power of 0 in the first index. This will get
+    # chopped off or multiplied by 0 later anyway.
+    if (f[..., 0] == 0.0).all():
+        mc_f[..., 0] = 1.0
+
+    # Leading (0PN) phase
+    pi_mc_f_SI = np.pi * (lal.GMSUN_SI / lal.C_SI ** 3) * mc_f
+    fiducial_phase = (3 / 128) * (pi_mc_f_SI) ** (-5 / 3)
+
+    # 1PN correction
+    if order >= 2:
+        assert mass_ratio is not None
+        symmetric_mass_ratio = mass_ratio / (1 + mass_ratio) ** 2
+        pi_m_f_SI = pi_mc_f_SI / symmetric_mass_ratio ** (3 / 5)
+        correction = 1 + (55 * symmetric_mass_ratio / 9 + 3715 / 756) * (
+            pi_m_f_SI
+        ) ** (2 / 3)
+
+        fiducial_phase *= correction
+
+    if inverse:
+        fiducial_phase *= -1
+
+    if type(data) == dict:
+        result = {}
+        for k, v in data.items():
+            result[k] = domain.add_phase(v, -fiducial_phase)
+    else:
+        result = domain.add_phase(data, -fiducial_phase)
+
+    return result
+
+
+class HeterodynePhase(object):
+    """
+    Transform operator for applying phase heterodyning. See docstring of
+    factor_fiducial_waveform for details.
+    """
+
+    def __init__(
+        self,
+        domain,
+        order: int = 0,
+        inverse: bool = False,
+        fixed_parameters: dict = None,
+    ):
+        """
+        Parameters
+        ----------
+        domain: Domain
+            Domain of the waveform data.
+        order : int
+            Twice the post-Newtonian order for the expansion. Either 0 or 2.
+        inverse : bool
+            Whether to apply the forward or inverse transform. Default: False.
+        fixed_parameters: dict = None
+            If set, extract chirp_mass and mass_ratio parameters from this dict instead
+            from the input sample.
+        """
+        self.domain = domain
+        self.order = order
+        self.inverse = inverse
+        self.fixed_parameters = fixed_parameters
+
+    def __call__(self, input_sample):
+        """
+        Parameters
+        ----------
+        input_sample : dict
+            sample["waveform"]: Dict with values being arrays containing waveforms,
+            or torch Tensor with the waveform.
+            sample["parameters"] contains parameters of the binary system. Required for
+            chirp mass (and mass ratio if self.order == 2).
+
+        Returns
+        -------
+        dict of the same form as the input, but with transformed (heterodyned) waveforms.
+        """
+        sample = input_sample.copy()
+        waveform = sample["waveform"]
+        if self.fixed_parameters:
+            parameters = self.fixed_parameters
+        else:
+            parameters = sample["parameters"]
+        sample["waveform"] = factor_fiducial_waveform(
+            waveform,
+            self.domain,
+            parameters["chirp_mass"],
+            mass_ratio=parameters.get("mass_ratio"),
+            order=self.order,
+            inverse=self.inverse,
+        )
+        return sample
diff --git a/dingo/gw/waveform_generator/waveform_generator.py b/dingo/gw/waveform_generator/waveform_generator.py
index 43f6a6488..b36385980 100644
--- a/dingo/gw/waveform_generator/waveform_generator.py
+++ b/dingo/gw/waveform_generator/waveform_generator.py
@@ -250,7 +250,8 @@ def generate_hplus_hcross(
             wf_dict = self._domain_transform(wf_dict)
 
         if self.transform is not None:
-            return self.transform(wf_dict)
+            data_sample = {"waveform": wf_dict, "parameters": parameters}
+            return self.transform(data_sample)["waveform"]
         else:
             return wf_dict
 
diff --git a/tests/gw/transforms/test_heterodyne.py b/tests/gw/transforms/test_heterodyne.py
new file mode 100644
index 000000000..0567ba58e
--- /dev/null
+++ b/tests/gw/transforms/test_heterodyne.py
@@ -0,0 +1,326 @@
+import pytest
+import numpy as np
+import pandas as pd
+import torch
+
+from dingo.gw.domains import UniformFrequencyDomain, build_domain
+from dingo.gw.transforms.waveform_transforms import (
+    factor_fiducial_waveform,
+    HeterodynePhase,
+)
+from dingo.gw.transforms.gnpe_transforms import GNPEChirp
+from dingo.gw.gwutils import get_mismatch
+
+
+@pytest.fixture
+def domain():
+    return UniformFrequencyDomain(f_min=20.0, f_max=1024.0, delta_f=0.125)
+
+
+@pytest.fixture
+def complex_waveform_dict(domain):
+    """Dict of complex numpy waveforms keyed by detector name."""
+    n = len(domain)
+    rng = np.random.default_rng(42)
+    return {
+        "H1": rng.standard_normal(n) + 1j * rng.standard_normal(n),
+        "L1": rng.standard_normal(n) + 1j * rng.standard_normal(n),
+    }
+
+
+@pytest.fixture
+def complex_waveform_tensor(domain):
+    """Complex torch tensor waveform, batched."""
+    torch.manual_seed(42)
+    batch = 8
+    n = len(domain)
+    return torch.randn(batch, n) + 1j * torch.randn(batch, n)
+
+
+# --- Tests for factor_fiducial_waveform ---
+
+
+def test_factor_fiducial_waveform_roundtrip_dict(domain, complex_waveform_dict):
+    """Heterodyne then inverse heterodyne recovers original data (dict input)."""
+    mc = 1.2
+    hetero = factor_fiducial_waveform(complex_waveform_dict, domain, mc)
+    recovered = factor_fiducial_waveform(hetero, domain, mc, inverse=True)
+    for k in complex_waveform_dict:
+        np.testing.assert_allclose(recovered[k], complex_waveform_dict[k], atol=1e-10)
+
+
+def test_factor_fiducial_waveform_roundtrip_tensor(domain, complex_waveform_tensor):
+    """Heterodyne then inverse heterodyne recovers original data (tensor input)."""
+    mc = torch.tensor(1.2)
+    hetero = factor_fiducial_waveform(complex_waveform_tensor, domain, mc)
+    recovered = factor_fiducial_waveform(hetero, domain, mc, inverse=True)
+    torch.testing.assert_close(recovered, complex_waveform_tensor, atol=1e-5, rtol=1e-5)
+
+
+def test_factor_fiducial_waveform_roundtrip_batched_tensor(domain):
+    """Roundtrip with batched chirp mass tensor."""
+    torch.manual_seed(0)
+    batch = 4
+    n = len(domain)
+    data = torch.randn(batch, n) + 1j * torch.randn(batch, n)
+    mc = torch.tensor([1.0, 1.2, 1.4, 1.6])
+    hetero = factor_fiducial_waveform(data, domain, mc)
+    recovered = factor_fiducial_waveform(hetero, domain, mc, inverse=True)
+    torch.testing.assert_close(recovered, data, atol=1e-4, rtol=1e-4)
+
+
+def test_factor_fiducial_waveform_unit_norm(domain, complex_waveform_dict):
+    """Verify the phase factor has unit norm (doesn't change amplitude)."""
+    mc = 1.2
+    hetero = factor_fiducial_waveform(complex_waveform_dict, domain, mc)
+    for k in complex_waveform_dict:
+        np.testing.assert_allclose(
+            np.abs(hetero[k]), np.abs(complex_waveform_dict[k]), atol=1e-12
+        )
+
+
+def test_factor_fiducial_waveform_order_2(domain, complex_waveform_dict):
+    """Verify order=2 (1PN correction) runs and differs from order=0 result."""
+    mc = 1.2
+    q = 0.8
+    result_0 = factor_fiducial_waveform(complex_waveform_dict, domain, mc, order=0)
+    result_2 = factor_fiducial_waveform(
+        complex_waveform_dict, domain, mc, mass_ratio=q, order=2
+    )
+    # They should differ
+    for k in complex_waveform_dict:
+        assert not np.allclose(result_0[k], result_2[k]), (
+            "order=0 and order=2 results should differ"
+        )
+    # But amplitudes should still be preserved
+    for k in complex_waveform_dict:
+        np.testing.assert_allclose(
+            np.abs(result_2[k]), np.abs(complex_waveform_dict[k]), atol=1e-12
+        )
+
+
+def test_factor_fiducial_waveform_invalid_order(domain, complex_waveform_dict):
+    """Verify invalid order raises ValueError."""
+    with pytest.raises(ValueError, match="Order 1 invalid"):
+        factor_fiducial_waveform(complex_waveform_dict, domain, 1.2, order=1)
+
+
+# --- Tests for HeterodynePhase ---
+
+
+def test_heterodyne_phase_transform(domain, complex_waveform_dict):
+    """Test HeterodynePhase.__call__ with a sample dict."""
+    transform = HeterodynePhase(domain, order=0)
+    sample = {
+        "waveform": complex_waveform_dict,
+        "parameters": {"chirp_mass": 1.2},
+    }
+    result = transform(sample)
+    # Waveform should be modified
+    for k in complex_waveform_dict:
+        assert not np.allclose(result["waveform"][k], complex_waveform_dict[k])
+    # Amplitudes preserved
+    for k in complex_waveform_dict:
+        np.testing.assert_allclose(
+            np.abs(result["waveform"][k]),
+            np.abs(complex_waveform_dict[k]),
+            atol=1e-12,
+        )
+
+
+def test_heterodyne_phase_fixed_parameters(domain, complex_waveform_dict):
+    """Test HeterodynePhase with fixed_parameters overriding sample parameters."""
+    mc_fixed = 1.4
+    transform = HeterodynePhase(
+        domain, order=0, fixed_parameters={"chirp_mass": mc_fixed}
+    )
+    sample = {
+        "waveform": complex_waveform_dict,
+        "parameters": {"chirp_mass": 999.0},  # should be ignored
+    }
+    result = transform(sample)
+
+    # Compare with direct call using fixed chirp mass
+    expected = factor_fiducial_waveform(complex_waveform_dict, domain, mc_fixed)
+    for k in complex_waveform_dict:
+        np.testing.assert_allclose(result["waveform"][k], expected[k], atol=1e-12)
+
+
+# --- Tests for GNPEChirp ---
+
+
+def test_gnpe_chirp_training(domain, complex_waveform_dict):
+    """Test GNPEChirp in training mode: proxies sampled, delta stored, waveform modified."""
+    kernel = {
+        "chirp_mass": "bilby.core.prior.Uniform(minimum=-0.01, maximum=0.01)",
+    }
+    transform = GNPEChirp(kernel, domain, order=0, inference=False)
+
+    mc_true = 1.2
+    sample = {
+        "waveform": complex_waveform_dict,
+        "parameters": {"chirp_mass": mc_true},
+        "extrinsic_parameters": {"chirp_mass": mc_true},
+    }
+    result = transform(sample)
+    ep = result["extrinsic_parameters"]
+
+    # Proxy should exist and be close to true value
+    assert "chirp_mass_proxy" in ep
+    assert abs(ep["chirp_mass_proxy"] - mc_true) <= 0.01
+
+    # Delta should be chirp_mass - proxy
+    assert "delta_chirp_mass" in ep
+    np.testing.assert_allclose(
+        ep["delta_chirp_mass"], mc_true - ep["chirp_mass_proxy"]
+    )
+
+    # Waveform should be modified
+    for k in complex_waveform_dict:
+        assert not np.allclose(result["waveform"][k], complex_waveform_dict[k])
+
+
+def test_gnpe_chirp_inference(domain):
+    """Test GNPEChirp in inference mode with pre-existing proxies."""
+    kernel = {
+        "chirp_mass": "bilby.core.prior.Uniform(minimum=-0.01, maximum=0.01)",
+    }
+    transform = GNPEChirp(kernel, domain, order=0, inference=True)
+
+    batch = 4
+    n = len(domain)
+    torch.manual_seed(0)
+    waveform = torch.randn(batch, n) + 1j * torch.randn(batch, n)
+
+    mc_true = torch.full((batch,), 1.2)
+    mc_proxy = torch.full((batch,), 1.19)
+
+    sample = {
+        "waveform": waveform,
+        "extrinsic_parameters": {
+            "chirp_mass": mc_true,
+            "chirp_mass_proxy": mc_proxy,
+        },
+    }
+    result = transform(sample)
+    ep = result["extrinsic_parameters"]
+
+    # Should reuse existing proxy
+    torch.testing.assert_close(ep["chirp_mass_proxy"], mc_proxy)
+
+
+def test_gnpe_chirp_context_parameters(domain):
+    """Verify context_parameters list is correct (should be proxy names from GNPEBase)."""
+    kernel = {
+        "chirp_mass": "bilby.core.prior.Uniform(minimum=-0.01, maximum=0.01)",
+    }
+    transform = GNPEChirp(kernel, domain, order=0)
+    assert transform.context_parameters == ["chirp_mass_proxy"]
+
+
+def test_gnpe_chirp_no_waveform(domain):
+    """Verify GNPEChirp works when no 'waveform' key (standardization case)."""
+    kernel = {
+        "chirp_mass": "bilby.core.prior.Uniform(minimum=-0.01, maximum=0.01)",
+    }
+    transform = GNPEChirp(kernel, domain, order=0)
+
+    sample = {
+        "parameters": {"chirp_mass": 1.2},
+        "extrinsic_parameters": {"chirp_mass": 1.2},
+    }
+    result = transform(sample)
+    ep = result["extrinsic_parameters"]
+
+    assert "chirp_mass_proxy" in ep
+    assert "delta_chirp_mass" in ep
+    assert "waveform" not in result
+
+
+# --- Integration test with real BNS waveforms (from Max's branch) ---
+
+
+def _number_of_zero_crossings(x):
+    """Count zero crossings along the last axis."""
+    return np.sum(np.diff(np.sign(x), axis=-1) != 0, axis=-1)
+
+
+def _max_zero_crossings(polarizations, min_idx=0):
+    return np.max(
+        [
+            np.max(_number_of_zero_crossings(v[..., min_idx:].real))
+            for v in polarizations.values()
+        ]
+    )
+
+
+def _max_mismatch(a, b, domain):
+    if a.keys() != b.keys():
+        raise ValueError()
+    return np.max([np.max(get_mismatch(a[k], b[k], domain)) for k in a.keys()])
+
+
+@pytest.mark.slow
+def test_heterodyning_bns_waveforms():
+    """
+    Integration test: generate BNS waveforms, heterodyne them, and verify that
+    (1) heterodyning is invertible (zero mismatch on roundtrip),
+    (2) heterodyned waveforms have significantly fewer zero crossings.
+
+    Adapted from Max Dax's test on the bns_add_dingo_pipe_max branch.
+    """
+    from dingo.gw.prior import build_prior_with_defaults
+    from dingo.gw.waveform_generator import (
+        WaveformGenerator,
+        generate_waveforms_parallel,
+    )
+
+    domain_settings = {
+        "type": "UniformFrequencyDomain",
+        "f_min": 20.0,
+        "f_max": 1024.0,
+        "delta_f": 0.00390625,
+    }
+    prior_settings = {
+        "mass_1": "bilby.core.prior.Constraint(minimum=1.0, maximum=2.5)",
+        "mass_2": "bilby.core.prior.Constraint(minimum=1.0, maximum=2.5)",
+        "chirp_mass": "bilby.gw.prior.UniformInComponentsChirpMass(minimum=1.0, maximum=2.0)",
+        "mass_ratio": "bilby.gw.prior.UniformInComponentsMassRatio(minimum=0.125, maximum=1.0)",
+        "phase": "default",
+        "a_1": "bilby.core.prior.Uniform(minimum=0.0, maximum=0.99)",
+        "a_2": "bilby.core.prior.Uniform(minimum=0.0, maximum=0.99)",
+        "tilt_1": "default",
+        "tilt_2": "default",
+        "phi_12": "default",
+        "phi_jl": "default",
+        "theta_jn": "default",
+        "lambda_1": "bilby.core.prior.Uniform(minimum=0.0, maximum=5000)",
+        "luminosity_distance": 100.0,
+        "geocent_time": 0.0,
+    }
+    wfg_settings = {"approximant": "IMRPhenomPv2_NRTidal", "f_ref": 10.0}
+
+    ufd = build_domain(domain_settings)
+    prior = build_prior_with_defaults(prior_settings)
+    wfg = WaveformGenerator(domain=ufd, **wfg_settings)
+
+    parameters = prior.sample(10)
+    polarizations = generate_waveforms_parallel(wfg, pd.DataFrame(parameters))
+
+    # Heterodyne and inverse heterodyne
+    polarizations_het = HeterodynePhase(ufd)(
+        {"waveform": polarizations, "parameters": parameters}
+    )["waveform"]
+    polarizations_roundtrip = HeterodynePhase(ufd, inverse=True)(
+        {"waveform": polarizations_het, "parameters": parameters}
+    )["waveform"]
+
+    # Roundtrip should recover original waveforms to machine precision.
+    assert _max_mismatch(polarizations, polarizations_roundtrip, ufd) < 1e-14
+    # Heterodyned waveforms should differ from originals.
+    assert _max_mismatch(polarizations, polarizations_het, ufd) > 1e-2
+
+    # Heterodyned waveforms should have far fewer zero crossings (the whole point).
+    n_roots = _max_zero_crossings(polarizations, ufd.min_idx)
+    n_roots_het = _max_zero_crossings(polarizations_het, ufd.min_idx)
+    assert n_roots > 10 * n_roots_het