Move all files into a separate directory, fix serialization

Author: vpratz

bayesflow/experimental/__init__.py

@@ -9,4 +9,4 @@
 
 from ..utils._docs import _add_imports_to_all
 
-_add_imports_to_all(include_modules=[])
+_add_imports_to_all(include_modules=["diffusion_model"])

bayesflow/experimental/diffusion_model/__init__.py

@@ -0,0 +1,6 @@
+from .diffusion_model import DiffusionModel
+from .noise_schedules import CosineNoiseSchedule, EDMNoiseSchedule, NoiseSchedule
+
+from ...utils._docs import _add_imports_to_all
+
+_add_imports_to_all(include_modules=[])

bayesflow/experimental/diffusion_model.py renamed to bayesflow/experimental/diffusion_model/diffusion_model.py

@@ -14,14 +14,15 @@
     weighted_mean,
     integrate,
     integrate_stochastic,
-    find_noise_schedule,
     logging,
     tensor_utils,
 )
+from .dispatch import find_noise_schedule
 from bayesflow.utils.serialization import serialize, deserialize, serializable
 
 
-@serializable("bayesflow.experimental")
+# disable module check, use potential module after moving from experimental
+@serializable("bayesflow.networks", disable_module_check=True)
 class DiffusionModel(InferenceNetwork):
     """Diffusion Model as described in this overview paper [1].
 

bayesflow/experimental/diffusion_model/dispatch.py

@@ -0,0 +1,59 @@
+from functools import singledispatch
+from .noise_schedules import NoiseSchedule
+
+
+@singledispatch
+def find_noise_schedule(arg, *args, **kwargs):
+    raise TypeError(f"Unknown noise schedule: {arg!r}")
+
+
+@find_noise_schedule.register
+def _(noise_schedule: NoiseSchedule):
+    return noise_schedule
+
+
+@find_noise_schedule.register
+def _(name: str, *args, **kwargs):
+    match name.lower():
+        case "cosine":
+            from .noise_schedules import CosineNoiseSchedule
+
+            return CosineNoiseSchedule()
+        case "edm":
+            from .noise_schedules import EDMNoiseSchedule
+
+            return EDMNoiseSchedule()
+        case other:
+            raise ValueError(f"Unsupported noise schedule name: '{other}'.")
+
+
+@find_noise_schedule.register
+def _(config: dict, *args, **kwargs):
+    name = config.get("type", "").lower()
+    params = {k: v for k, v in config.items() if k != "type"}
+    match name:
+        case "cosine":
+            from .noise_schedules import CosineNoiseSchedule
+
+            return CosineNoiseSchedule(**params)
+        case "edm":
+            from .noise_schedules import EDMNoiseSchedule
+
+            return EDMNoiseSchedule(**params)
+        case other:
+            raise ValueError(f"Unsupported noise schedule config: '{other}'.")
+
+
+@find_noise_schedule.register
+def _(cls: type, *args, **kwargs):
+    # Lazily import NoiseSchedule class and compare
+    from .noise_schedules import NoiseSchedule
+
+    if issubclass(cls, NoiseSchedule):
+        return cls(*args, **kwargs)
+    raise TypeError(f"Expected subclass of NoiseSchedule, got {cls}")
+
+
+@find_noise_schedule.register
+def _(schedule: type, *args, **kwargs):
+    return schedule

bayesflow/experimental/noise_schedules.py renamed to bayesflow/experimental/diffusion_model/noise_schedules.py

@@ -8,7 +8,8 @@
 from bayesflow.utils.serialization import deserialize, serializable
 
 
-@serializable("bayesflow.experimental")
+# disable module check, use potential module after moving from experimental
+@serializable("bayesflow.networks", disable_module_check=True)
 class NoiseSchedule(ABC):
     r"""Noise schedule for diffusion models. We follow the notation from [1].
 
@@ -227,7 +228,8 @@ def from_config(cls, config, custom_objects=None):
         return cls(**deserialize(config, custom_objects=custom_objects))
 
 
-@serializable("bayesflow.experimental")
+# disable module check, use potential module after moving from experimental
+@serializable("bayesflow.networks", disable_module_check=True)
 class EDMNoiseSchedule(NoiseSchedule):
     """EDM noise schedule for diffusion models. This schedule is based on the EDM paper [1].
     This should be used with the F-prediction type in the diffusion model.

bayesflow/utils/dispatch/__init__.py

@@ -6,4 +6,3 @@
 from .find_inference_network import find_inference_network
 from .find_summary_network import find_summary_network
 from .find_distribution import find_distribution
-from .find_noise_schedule import find_noise_schedule

bayesflow/utils/dispatch/find_noise_schedule.py

@@ -1,53 +0,0 @@
-from functools import singledispatch
-
-
-@singledispatch
-def find_noise_schedule(arg, *args, **kwargs):
-    raise TypeError(f"Unknown noise schedule: {arg!r}")
-
-
-@find_noise_schedule.register
-def _(name: str, *args, **kwargs):
-    match name.lower():
-        case "cosine":
-            from bayesflow.experimental.noise_schedules import CosineNoiseSchedule
-
-            return CosineNoiseSchedule()
-        case "edm":
-            from bayesflow.experimental.noise_schedules import EDMNoiseSchedule
-
-            return EDMNoiseSchedule()
-        case other:
-            raise ValueError(f"Unsupported noise schedule name: '{other}'.")
-
-
-@find_noise_schedule.register
-def _(config: dict, *args, **kwargs):
-    name = config.get("type", "").lower()
-    params = {k: v for k, v in config.items() if k != "type"}
-    match name:
-        case "cosine":
-            from bayesflow.experimental.noise_schedules import CosineNoiseSchedule
-
-            return CosineNoiseSchedule(**params)
-        case "edm":
-            from bayesflow.experimental.noise_schedules import EDMNoiseSchedule
-
-            return EDMNoiseSchedule(**params)
-        case other:
-            raise ValueError(f"Unsupported noise schedule config: '{other}'.")
-
-
-@find_noise_schedule.register
-def _(cls: type, *args, **kwargs):
-    # Lazily import NoiseSchedule class and compare
-    from bayesflow.experimental.noise_schedules import NoiseSchedule
-
-    if issubclass(cls, NoiseSchedule):
-        return cls(*args, **kwargs)
-    raise TypeError(f"Expected subclass of NoiseSchedule, got {cls}")
-
-
-@find_noise_schedule.register
-def _(schedule: type, *args, **kwargs):
-    return schedule