Rename GlobalWorkspaceFusion to GlobalWorkspace (#70)

docs/shimmer_basics.md

@@ -13,7 +13,7 @@ to make a GW in shimmer:
 ![architecture](assets/shimmer_architecture.png)
 
 Let's detail:
-- [`DomainModule`](https://bdvllrs.github.io/shimmer/shimmer/modules/domain.html#DomainModule)s
+- [`DomainModule`](https://bdvllrs.github.io/shimmer/latest/shimmer/modules/domain.html#DomainModule)s
     are the individual domain modules which encode domain data into a latent vector;
 - the `GWModule` has access to the domain modules, and defines how to encode, decode and merge representations of the domains into a unique GW representation.
 - finally `GlobalWorkspaceBase` takes all building blocks to make a [Pytorch Lightning](https://lightning.ai/docs/pytorch/stable/) module
@@ -87,7 +87,7 @@ class DomainDataModule(LightningDataModule):
 Now that our data module is defined, let's create `DomainModule`s.
 
 ## `DomainModule`
-For more details about DomainModules, see the [DomainModule API docs](https://bdvllrs.github.io/shimmer/shimmer/modules/domain.html#DomainModule).
+For more details about DomainModules, see the [DomainModule API docs](https://bdvllrs.github.io/shimmer/latest/shimmer/modules/domain.html#DomainModule).
 The `DomainModule` class extends from a LightningModule and requires you to define some
 methods:
 
@@ -422,7 +422,7 @@ class GenericDomain(DomainModule):
         return LossOutput(loss=F.mse_loss(pred, target))
 ```
 
-To learn more about LossOutput, see [API docs](https://bdvllrs.github.io/shimmer/shimmer/modules/domain.html#LossOutput).
+To learn more about LossOutput, see [API docs](https://bdvllrs.github.io/shimmer/latest/shimmer/modules/domain.html#LossOutput).
 
 ## Let's make a GW!
 
@@ -436,7 +436,7 @@ from lightning.pytorch import Trainer
 from lightning.pytorch.callbacks import ModelCheckpoint
 from torch import nn
 
-from shimmer import GlobalWorkspace, GWDecoder, GWEncoder, LossCoefs
+from shimmer import GlobalWorkspace2Domains, GWDecoder, GWEncoder, LossCoefs
 from shimmer.modules.global_workspace import SchedulerArgs
 
 
@@ -508,7 +508,7 @@ def train_gw():
 
     n_epochs = 4
 
-    global_workspace = GlobalWorkspace(
+    global_workspace = GlobalWorkspace2Domains(
         domain_mods,
         gw_encoders,
         gw_decoders,
@@ -631,7 +631,7 @@ We define loss coefficients for the different losses. Note that `LossCoefs` is a
 
 Finally we make the GlobalWorkspace and train it.
 ```python
-    global_workspace = GlobalWorkspace(
+    global_workspace = GlobalWorkspace2Domains(
         domain_mods,
         gw_encoders,
         gw_decoders,

examples/main_example/train_gw.py

@@ -4,7 +4,7 @@
 from lightning.pytorch.callbacks import ModelCheckpoint
 from torch import nn
 
-from shimmer import GlobalWorkspace, GWDecoder, GWEncoder, LossCoefs
+from shimmer import GlobalWorkspace2Domains, GWDecoder, GWEncoder, LossCoefs
 from shimmer.modules.global_workspace import SchedulerArgs
 
 
@@ -76,7 +76,7 @@ def train_gw():
 
     n_epochs = 4
 
-    global_workspace = GlobalWorkspace(
+    global_workspace = GlobalWorkspace2Domains(
         domain_mods,
         gw_encoders,
         gw_decoders,

shimmer/__init__.py

@@ -6,7 +6,7 @@
 )
 from shimmer.modules.domain import DomainModule, LossOutput
 from shimmer.modules.global_workspace import (
-    GlobalWorkspace,
+    GlobalWorkspace2Domains,
     GlobalWorkspaceBase,
     GlobalWorkspaceBayesian,
     GWPredictions,
@@ -23,7 +23,7 @@
 )
 from shimmer.modules.losses import (
     BroadcastLossCoefs,
-    GWLosses,
+    GWLosses2Domains,
     GWLossesBase,
     GWLossesBayesian,
     LossCoefs,
@@ -68,7 +68,7 @@
     "SchedulerArgs",
     "GWPredictions",
     "GlobalWorkspaceBase",
-    "GlobalWorkspace",
+    "GlobalWorkspace2Domains",
     "GlobalWorkspaceBayesian",
     "pretrained_global_workspace",
     "LossOutput",
@@ -85,7 +85,7 @@
     "LossCoefs",
     "BroadcastLossCoefs",
     "GWLossesBase",
-    "GWLosses",
+    "GWLosses2Domains",
     "GWLossesBayesian",
     "RepeatedDataset",
     "batch_cycles",

shimmer/modules/__init__.py

@@ -7,7 +7,7 @@
 )
 from shimmer.modules.domain import DomainModule, LossOutput
 from shimmer.modules.global_workspace import (
-    GlobalWorkspace,
+    GlobalWorkspace2Domains,
     GlobalWorkspaceBase,
     GlobalWorkspaceBayesian,
     GWPredictions,
@@ -24,7 +24,7 @@
 )
 from shimmer.modules.losses import (
     BroadcastLossCoefs,
-    GWLosses,
+    GWLosses2Domains,
     GWLossesBase,
     GWLossesBayesian,
     LossCoefs,
@@ -54,7 +54,7 @@
     "SchedulerArgs",
     "GWPredictions",
     "GlobalWorkspaceBase",
-    "GlobalWorkspace",
+    "GlobalWorkspace2Domains",
     "GlobalWorkspaceBayesian",
     "pretrained_global_workspace",
     "LossOutput",
@@ -72,7 +72,7 @@
     "LossCoefs",
     "BroadcastLossCoefs",
     "GWLossesBase",
-    "GWLosses",
+    "GWLosses2Domains",
     "GWLossesBayesian",
     "RepeatedDataset",
     "reparameterize",

shimmer/modules/global_workspace.py

@@ -18,9 +18,9 @@
 from shimmer.modules.losses import (
     BroadcastLossCoefs,
     GWLosses,
+    GWLosses2Domains,
     GWLossesBase,
     GWLossesBayesian,
-    GWLossesFusion,
     LossCoefs,
 )
 from shimmer.modules.selection import (
@@ -466,7 +466,9 @@ class GWPredictions(GWPredictionsBase):
     """
 
 
-class GlobalWorkspace(GlobalWorkspaceBase[GWModule, SingleDomainSelection, GWLosses]):
+class GlobalWorkspace2Domains(
+    GlobalWorkspaceBase[GWModule, SingleDomainSelection, GWLosses2Domains]
+):
     """
     A simple 2-domains max flavor of GlobalWorkspaceBase.
 
@@ -519,7 +521,7 @@ def __init__(
                 torch.tensor([1 / 0.07]).log(), "mean", learn_logit_scale
             )
         selection_mod = SingleDomainSelection()
-        loss_mod = GWLosses(
+        loss_mod = GWLosses2Domains(
             gw_mod, selection_mod, domain_mods, loss_coefs, contrastive_loss
         )
 
@@ -559,9 +561,7 @@ def forward(  # type: ignore
         )
 
 
-class GlobalWorkspaceFusion(
-    GlobalWorkspaceBase[GWModule, RandomSelection, GWLossesFusion]
-):
+class GlobalWorkspace(GlobalWorkspaceBase[GWModule, RandomSelection, GWLosses]):
     """The 2-domain fusion (with broadcast loss) flavor of GlobalWorkspaceBase.
 
     This is used to simplify a Global Workspace instanciation and only overrides the
@@ -617,7 +617,7 @@ def __init__(
             )
 
         selection_mod = RandomSelection(selection_temperature)
-        loss_mod = GWLossesFusion(
+        loss_mod = GWLosses(
             gw_mod, selection_mod, domain_mods, loss_coefs, contrastive_loss
         )
 
@@ -630,6 +630,33 @@ def __init__(
             scheduler_args,
         )
 
+    def forward(  # type: ignore
+        self,
+        latent_domains: LatentsDomainGroupsT,
+    ) -> GWPredictions:
+        """
+        Computes demi-cycles, cycles, and translations.
+
+        Args:
+            latent_domains (`LatentsT`): Groups of domains for the computation.
+
+        Returns:
+            `GWPredictions`: the predictions on the batch.
+        """
+        return GWPredictions(
+            demi_cycles=batch_demi_cycles(
+                self.gw_mod, self.selection_mod, latent_domains
+            ),
+            cycles=batch_cycles(
+                self.gw_mod, self.selection_mod, latent_domains, self.domain_mods.keys()
+            ),
+            translations=batch_translations(
+                self.gw_mod, self.selection_mod, latent_domains
+            ),
+            # TODO: add other combinations
+            **super().forward(latent_domains),
+        )
+
 
 class GlobalWorkspaceBayesian(
     GlobalWorkspaceBase[GWModuleBayesian, FixedSharedSelection, GWLossesBayesian]
@@ -753,7 +780,7 @@ def pretrained_global_workspace(
     loss_coefs: LossCoefs,
     contrastive_fn: ContrastiveLossType,
     **kwargs,
-) -> GlobalWorkspace:
+) -> GlobalWorkspace2Domains:
     """
     Load a `GlobalWorkspace` flavor of `GlobalWorkspaceBase` from a checkpoint.
 
@@ -785,16 +812,18 @@ def pretrained_global_workspace(
     domain_mods = freeze_domain_modules(domain_mods)
     gw_mod = GWModule(domain_mods, workspace_dim, gw_encoders, gw_decoders)
     selection_mod = SingleDomainSelection()
-    loss_mod = GWLosses(gw_mod, selection_mod, domain_mods, loss_coefs, contrastive_fn)
+    loss_mod = GWLosses2Domains(
+        gw_mod, selection_mod, domain_mods, loss_coefs, contrastive_fn
+    )
 
-    gw = GlobalWorkspace.load_from_checkpoint(
+    gw = GlobalWorkspace2Domains.load_from_checkpoint(
         checkpoint_path,
         gw_mod=gw_mod,
         selection_mid=selection_mod,
         loss_coefs=loss_coefs,
         loss_mod=loss_mod,
         **kwargs,
     )
-    if not isinstance(gw, GlobalWorkspace):
+    if not isinstance(gw, GlobalWorkspace2Domains):
         raise TypeError("model should be of type GlobalWorkspace")
     return gw

shimmer/modules/losses.py

@@ -348,7 +348,7 @@ class LossCoefs(TypedDict, total=False):
     """Contrastive loss coefficient."""
 
 
-class GWLosses(GWLossesBase):
+class GWLosses2Domains(GWLossesBase):
     """
     Implementation of `GWLossesBase` used for `GWModule`.
     """
@@ -659,7 +659,7 @@ class BroadcastLossCoefs(TypedDict, total=False):
     """translation loss coefficient. Translation, like cycles, can be many-to-one."""
 
 
-class GWLossesFusion(GWLossesBase):
+class GWLosses(GWLossesBase):
     """
     Implementation of `GWLossesBase` for fusion-based models.
     """

tests/test_broadcast.py

@@ -2,7 +2,7 @@
 from torch import nn
 
 from shimmer.modules.domain import DomainModule, LossOutput
-from shimmer.modules.global_workspace import GlobalWorkspaceFusion
+from shimmer.modules.global_workspace import GlobalWorkspace
 from shimmer.modules.losses import BroadcastLossCoefs
 
 
@@ -39,7 +39,7 @@ def test_broadcast_loss():
         "contrastives": 0.1,
     }
 
-    gw_fusion = GlobalWorkspaceFusion(
+    gw_fusion = GlobalWorkspace(
         domain_mods,
         gw_encoders,
         gw_decoders,

tests/test_training.py

@@ -1,7 +1,7 @@
 import torch.utils.data
 from utils import DummyData, DummyDataset, DummyDomainModule
 
-from shimmer import GlobalWorkspace, GWDecoder, GWEncoder
+from shimmer import GlobalWorkspace2Domains, GWDecoder, GWEncoder
 from shimmer.modules.selection import SingleDomainSelection
 
 
@@ -61,7 +61,7 @@ def test_training():
         ),
     }
 
-    gw = GlobalWorkspace(
+    gw = GlobalWorkspace2Domains(
         domains,
         gw_encoders,
         gw_decoders,