Add broadcasts and cycled broadcast to GW forward (#104)

* Add batch_broadcasts utils

* GlobalWorkspace forwards now return broadcasts

* Add forward for gw_mod

* fix: circular import

* Do broadcast cycles only if inverse is non empty

docs/make.py

@@ -14,7 +14,6 @@
     "shimmer.modules.contrastive_loss",
     "shimmer.dataset",
     "shimmer.modules.vae",
-    "shimmer.modules.utils",
     "shimmer.utils",
     "shimmer.cli.ckpt_migration",
 ]

shimmer/__init__.py

@@ -15,8 +15,11 @@
     GlobalWorkspace2Domains,
     GlobalWorkspaceBase,
     GlobalWorkspaceBayesian,
-    GWPredictions,
     SchedulerArgs,
+    batch_broadcasts,
+    batch_cycles,
+    batch_demi_cycles,
+    batch_translations,
     pretrained_global_workspace,
 )
 from shimmer.modules.gw_module import (
@@ -26,6 +29,11 @@
     GWModule,
     GWModuleBase,
     GWModuleBayesian,
+    GWModulePrediction,
+    broadcast,
+    broadcast_cycles,
+    cycle,
+    translation,
 )
 from shimmer.modules.losses import (
     BroadcastLossCoefs,
@@ -39,13 +47,6 @@
     SelectionBase,
     SingleDomainSelection,
 )
-from shimmer.modules.utils import (
-    batch_cycles,
-    batch_demi_cycles,
-    batch_translations,
-    cycle,
-    translation,
-)
 from shimmer.types import (
     LatentsDomainGroupDT,
     LatentsDomainGroupsDT,
@@ -72,7 +73,6 @@
     "RawDomainGroupT",
     "ModelModeT",
     "SchedulerArgs",
-    "GWPredictions",
     "GlobalWorkspaceBase",
     "GlobalWorkspace2Domains",
     "GlobalWorkspaceBayesian",
@@ -85,6 +85,7 @@
     "GWModuleBase",
     "GWModule",
     "GWModuleBayesian",
+    "GWModulePrediction",
     "ContrastiveLossType",
     "contrastive_loss",
     "ContrastiveLoss",
@@ -97,6 +98,9 @@
     "batch_cycles",
     "batch_demi_cycles",
     "batch_translations",
+    "batch_broadcasts",
+    "broadcast",
+    "broadcast_cycles",
     "cycle",
     "translation",
     "MIGRATION_DIR",

shimmer/modules/__init__.py

@@ -10,8 +10,11 @@
     GlobalWorkspace2Domains,
     GlobalWorkspaceBase,
     GlobalWorkspaceBayesian,
-    GWPredictions,
     SchedulerArgs,
+    batch_broadcasts,
+    batch_cycles,
+    batch_demi_cycles,
+    batch_translations,
     pretrained_global_workspace,
 )
 from shimmer.modules.gw_module import (
@@ -21,6 +24,11 @@
     GWModule,
     GWModuleBase,
     GWModuleBayesian,
+    GWModulePrediction,
+    broadcast,
+    broadcast_cycles,
+    cycle,
+    translation,
 )
 from shimmer.modules.losses import (
     BroadcastLossCoefs,
@@ -34,13 +42,6 @@
     SelectionBase,
     SingleDomainSelection,
 )
-from shimmer.modules.utils import (
-    batch_cycles,
-    batch_demi_cycles,
-    batch_translations,
-    cycle,
-    translation,
-)
 from shimmer.modules.vae import (
     VAE,
     VAEDecoder,
@@ -52,7 +53,6 @@
 
 __all__ = [
     "SchedulerArgs",
-    "GWPredictions",
     "GlobalWorkspaceBase",
     "GlobalWorkspace2Domains",
     "GlobalWorkspaceBayesian",
@@ -65,6 +65,7 @@
     "GWModuleBase",
     "GWModule",
     "GWModuleBayesian",
+    "GWModulePrediction",
     "ContrastiveLossType",
     "ContrastiveLossBayesianType",
     "contrastive_loss",
@@ -84,6 +85,9 @@
     "batch_cycles",
     "batch_demi_cycles",
     "batch_translations",
+    "batch_broadcasts",
+    "broadcast",
+    "broadcast_cycles",
     "cycle",
     "translation",
     "RandomSelection",