import dpsynth fails because mbi.Dataset is registered as a JAX dataclass twice

[hanzalaareeb]

A fresh minimal installation from main succeeds, but importing DPSynth fails.

Environment:

• WSL2 Linux x86_64
• CPython 3.12.3
• pip 26.2.1
• dpsynth 0.4.0 (ba212434752f67a128c962e3ca3e2be5b53d2253)
• mbi 2.0.0 (025b76f84f40256529732e98d93c1139fb2153d8)
• jax/jaxlib 0.11.0

Installation:

pip install "dpsynth @ git+https://github.com/google/dpsynth.git"

installation completes successfully, but:

python -c "import dpsynth; print(dpsynth.__version__)"

fails with:

ValueError: Duplicate custom dataclass PyTreeDef type registration for
<class 'mbi.dataset.Dataset'>.

MBI already registers Dataset:

@jax.tree_util.register_dataclass
@dataclasses.dataclass(frozen=True, eq=False)
class Dataset:

DPSynth registers it again in dpsynth/checkpoint.py:

jax.tree_util.register_dataclass(
    mbi.Dataset,
    data_fields=["data", "weights"],
    meta_fields=["domain"],
)

The DPSynth-side registration should likely be removed now that MBI owns the
registration. A clean-environment regression test for import dpsynth would
also prevent this from recurring.