Update __init__ for documentation order

ruflab · Mar 5, 2024 · 55f7ca8 · 55f7ca8
1 parent 1c3bce2
commit 55f7ca8
Show file tree

Hide file tree

Showing 2 changed files with 53 additions and 39 deletions.
diff --git a/shimmer/__init__.py b/shimmer/__init__.py
@@ -39,6 +39,7 @@
     LatentsDomainGroupsDT,
     LatentsDomainGroupsT,
     LatentsDomainGroupT,
+    ModelModeT,
     RawDomainGroupDT,
     RawDomainGroupsDT,
     RawDomainGroupsT,
@@ -48,44 +49,42 @@
 
 __all__ = [
     "__version__",
-    "DomainModule",
+    "LatentsDomainGroupDT",
+    "LatentsDomainGroupsDT",
+    "LatentsDomainGroupsT",
+    "LatentsDomainGroupT",
+    "RawDomainGroupDT",
+    "RawDomainGroupsDT",
+    "RawDomainGroupsT",
+    "RawDomainGroupT",
+    "ModelModeT",
+    "SchedulerArgs",
+    "GWPredictions",
+    "GlobalWorkspaceBase",
+    "GlobalWorkspace",
+    "VariationalGlobalWorkspace",
+    "pretrained_global_workspace",
     "LossOutput",
-    "GWInterfaceBase",
-    "GWModule",
+    "DomainModule",
     "GWDecoder",
     "GWEncoder",
-    "GWInterface",
-    "GWModuleBase",
-    "VariationalGWEncoder",
-    "VariationalGWInterface",
-    "VariationalGWModule",
     "VariationalGWEncoder",
+    "GWInterfaceBase",
+    "GWInterface",
     "VariationalGWInterface",
+    "GWModuleBase",
+    "GWModule",
     "VariationalGWModule",
-    "ContrastiveLoss",
     "ContrastiveLossType",
     "VarContrastiveLossType",
-    "ContrastiveLossWithUncertainty",
     "contrastive_loss",
+    "ContrastiveLoss",
     "contrastive_loss_with_uncertainty",
+    "ContrastiveLossWithUncertainty",
     "LossCoefs",
     "VariationalLossCoefs",
-    "GWLosses",
     "GWLossesBase",
+    "GWLosses",
     "VariationalGWLosses",
-    "GlobalWorkspace",
-    "GlobalWorkspaceBase",
-    "VariationalGlobalWorkspace",
-    "SchedulerArgs",
-    "GWPredictions",
-    "pretrained_global_workspace",
     "RepeatedDataset",
-    "LatentsDomainGroupDT",
-    "LatentsDomainGroupsDT",
-    "LatentsDomainGroupsT",
-    "LatentsDomainGroupT",
-    "RawDomainGroupDT",
-    "RawDomainGroupsDT",
-    "RawDomainGroupsT",
-    "RawDomainGroupT",
 ]
diff --git a/shimmer/modules/__init__.py b/shimmer/modules/__init__.py
@@ -1,3 +1,4 @@
+from shimmer.dataset import RepeatedDataset
 from shimmer.modules.contrastive_loss import (
     ContrastiveLoss,
     ContrastiveLossType,
@@ -33,35 +34,49 @@
     VariationalGWLosses,
     VariationalLossCoefs,
 )
+from shimmer.modules.vae import (
+    VAE,
+    VAEDecoder,
+    VAEEncoder,
+    gaussian_nll,
+    kl_divergence_loss,
+    reparameterize,
+)
 
 __all__ = [
-    "DomainModule",
+    "SchedulerArgs",
+    "GWPredictions",
+    "GlobalWorkspaceBase",
+    "GlobalWorkspace",
+    "VariationalGlobalWorkspace",
+    "pretrained_global_workspace",
     "LossOutput",
-    "GWInterfaceBase",
-    "GWModule",
+    "DomainModule",
     "GWDecoder",
     "GWEncoder",
+    "VariationalGWEncoder",
+    "GWInterfaceBase",
     "GWInterface",
+    "VariationalGWInterface",
     "GWModuleBase",
     "GWModule",
-    "VariationalGWEncoder",
-    "VariationalGWInterface",
     "VariationalGWModule",
-    "ContrastiveLoss",
     "ContrastiveLossType",
     "VarContrastiveLossType",
-    "ContrastiveLossWithUncertainty",
     "contrastive_loss",
+    "ContrastiveLoss",
     "contrastive_loss_with_uncertainty",
+    "ContrastiveLossWithUncertainty",
     "LossCoefs",
     "VariationalLossCoefs",
-    "GWLosses",
     "GWLossesBase",
+    "GWLosses",
     "VariationalGWLosses",
-    "GlobalWorkspace",
-    "GlobalWorkspaceBase",
-    "VariationalGlobalWorkspace",
-    "SchedulerArgs",
-    "GWPredictions",
-    "pretrained_global_workspace",
+    "RepeatedDataset",
+    "reparameterize",
+    "kl_divergence_loss",
+    "gaussian_nll",
+    "VAEEncoder",
+    "VAEDecoder",
+    "VAE",
 ]