shimmer.modules.global_workspace

  1from collections.abc import Iterable, Mapping
  2from pathlib import Path
  3from typing import Any, Generic, TypedDict, TypeVar, cast
  4
  5import torch
  6from lightning.pytorch import LightningModule
  7from lightning.pytorch.utilities.types import OptimizerLRSchedulerConfig
  8from torch.nn import Module, ModuleDict
  9from torch.optim.lr_scheduler import OneCycleLR
 10
 11from shimmer.modules.contrastive_loss import ContrastiveLoss, ContrastiveLossType
 12from shimmer.modules.domain import DomainModule
 13from shimmer.modules.gw_module import (
 14    GWModule,
 15    GWModuleBase,
 16    GWModuleBayesian,
 17)
 18from shimmer.modules.losses import (
 19    BroadcastLossCoefs,
 20    GWLosses,
 21    GWLosses2Domains,
 22    GWLossesBase,
 23    GWLossesBayesian,
 24    LossCoefs,
 25)
 26from shimmer.modules.selection import (
 27    FixedSharedSelection,
 28    RandomSelection,
 29    SelectionBase,
 30    SingleDomainSelection,
 31)
 32from shimmer.modules.utils import batch_cycles, batch_demi_cycles, batch_translations
 33from shimmer.types import (
 34    LatentsDomainGroupsDT,
 35    LatentsDomainGroupsT,
 36    ModelModeT,
 37    RawDomainGroupsDT,
 38    RawDomainGroupsT,
 39    RawDomainGroupT,
 40)
 41from shimmer.utils import groups_batch_size
 42
 43
 44class SchedulerArgs(TypedDict, total=False):
 45    """TypedDict of arguments passed to the OneCycle scheduler"""
 46
 47    max_lr: float
 48    """Maximum learning rate"""
 49
 50    total_steps: int
 51    """Total number of steps"""
 52
 53
 54class GWPredictionsBase(TypedDict):
 55    """TypedDict of the output given when calling `GlobalWorkspaceBase.predict`"""
 56
 57    states: dict[str, torch.Tensor]
 58    """
 59    GW state representation from domain groups with only one domain.
 60    The key represent the domain's name.
 61    """
 62
 63
 64_T_gw_mod = TypeVar("_T_gw_mod", bound=GWModuleBase)
 65_T_selection_mod = TypeVar("_T_selection_mod", bound=SelectionBase)
 66_T_loss_mod = TypeVar("_T_loss_mod", bound=GWLossesBase)
 67
 68
 69class GlobalWorkspaceBase(
 70    Generic[_T_gw_mod, _T_selection_mod, _T_loss_mod], LightningModule
 71):
 72    """
 73    Global Workspace Lightning Module.
 74
 75    This is the base class to build the Global Workspace.
 76    """
 77
 78    def __init__(
 79        self,
 80        gw_mod: _T_gw_mod,
 81        selection_mod: _T_selection_mod,
 82        loss_mod: _T_loss_mod,
 83        optim_lr: float = 1e-3,
 84        optim_weight_decay: float = 0.0,
 85        scheduler_args: SchedulerArgs | None = None,
 86    ) -> None:
 87        """
 88        Initializes a GW
 89
 90        Args:
 91            gw_mod (`GWModuleBase`): the GWModule
 92            selection_mod (`SelectionBase`): selection module
 93            loss_mod (`GWLossesBase`): module to compute the GW losses.
 94            optim_lr (`float`): learning rate
 95            optim_weight_decay (`float`): weight decay
 96            scheduler_args (`SchedulerArgs`): `SchedulerArgs` instance to define
 97                scheduler parameters.
 98        """
 99        super().__init__()
100        self.save_hyperparameters(
101            ignore=[
102                "gw_mod",
103                "selection_mod",
104                "domain_mods",
105                "loss_mod",
106                "domain_descriptions",
107                "contrastive_loss",
108                "cont_loss_bayesian",
109                "gw_encoders",
110                "gw_decoders",
111            ]
112        )
113
114        self.gw_mod = gw_mod
115        """ a `GWModuleBase` implementation."""
116
117        self.selection_mod = selection_mod
118        """A `SelectionBase` implementation."""
119
120        self.loss_mod = loss_mod
121        """The module that computes losses of the GW"""
122
123        self.optim_lr = optim_lr
124        self.optim_weight_decay = optim_weight_decay
125        self.scheduler_args = SchedulerArgs(max_lr=optim_lr, total_steps=1)
126        if scheduler_args is not None:
127            self.scheduler_args.update(scheduler_args)
128
129    @property
130    def domain_mods(self) -> Mapping[str, DomainModule]:
131        return self.gw_mod.domain_mods
132
133    @property
134    def workspace_dim(self) -> int:
135        """Dimension of the GW."""
136        return self.gw_mod.workspace_dim
137
138    def encode_and_fuse(
139        self, x: LatentsDomainGroupsT, selection_module: SelectionBase
140    ) -> dict[frozenset[str], torch.Tensor]:
141        """
142        Encode a group of latent representations into the GW representation.
143
144        Args:
145            x (`LatentsDomainGroupsT`): the input domain representations.
146            selection_scores (`Mapping[str, torch.Tensor]`):
147
148        Returns:
149            `dict[frozenset[str], torch.Tensor]`: the GW representations.
150        """
151        return {
152            domains: self.gw_mod.encode_and_fuse(latents, selection_module)
153            for domains, latents in x.items()
154        }
155
156    def encode(self, x: LatentsDomainGroupsT) -> LatentsDomainGroupsDT:
157        """
158        Encode a group of latent representations into the pre-fusion GW representation.
159
160        Args:
161            x (`LatentsDomainGroupsT`): the input domain representations.
162
163        Returns:
164            `LatensDomainGroupsDT`: the GW representations.
165        """
166        return {domains: self.gw_mod.encode(latents) for domains, latents in x.items()}
167
168    def fuse(
169        self,
170        x: LatentsDomainGroupsT,
171        selection_scores: Mapping[frozenset[str], Mapping[str, torch.Tensor]],
172    ) -> dict[frozenset[str], torch.Tensor]:
173        """
174        Fuses a group of latent representations into the GW representation.
175
176        Args:
177            x (`LatentsDomainGroupsT`): the pre-fusion latent representations
178            selection_scores (`Mapping[frozenset[str], Mapping[str, torch.Tensor]]`):
179                selection scores for each group
180
181        Returns:
182            `dict[frozenset[str], torch.Tensor]`: GW representation of each group
183        """
184        return {
185            domains: self.gw_mod.fuse(latents, selection_scores[domains])
186            for domains, latents in x.items()
187        }
188
189    def decode(
190        self,
191        z: Mapping[frozenset[str], torch.Tensor],
192        domains: Iterable[str] | None = None,
193    ) -> LatentsDomainGroupsDT:
194        """
195        Decode the group GW representation into given `domains`.
196
197        Args:
198            z (`torch.Tensor`): the GW representation.
199            domains (`Iterable[str]`): iterable of domains to decode.
200
201        Returns:
202            `dict[str, torch.Tensor]`: the decoded unimodal representations.
203        """
204        return {
205            domain_names: self.gw_mod.decode(gw_rep, domains)
206            for domain_names, gw_rep in z.items()
207        }
208
209    def forward(  # type: ignore
210        self,
211        latent_domains: LatentsDomainGroupsT,
212    ) -> GWPredictionsBase:
213        """
214        Computes demi-cycles, cycles, and translations.
215
216        Args:
217            latent_domains (`LatentsT`): Groups of domains for the computation.
218
219        Returns:
220            `GWPredictionsBase`: the predictions on the batch.
221        """
222
223        return GWPredictionsBase(states=self.batch_gw_states(latent_domains))
224
225    def batch_gw_states(
226        self, latent_domains: LatentsDomainGroupsT
227    ) -> dict[str, torch.Tensor]:
228        """
229        Comptues GW states of a batch of groups of domains.
230
231        Args:
232            latent_domains (`LatentsT`): the batch of groups of domains
233
234        Returns:
235            `dict[str, torch.Tensor]`: states for each domain.
236        """
237        predictions: dict[str, torch.Tensor] = {}
238        for domains, latents in latent_domains.items():
239            if len(domains) > 1:
240                continue
241            domain_name = list(domains)[0]
242            z = self.gw_mod.encode_and_fuse(
243                latents, selection_module=self.selection_mod
244            )
245            predictions[domain_name] = z
246        return predictions
247
248    def encode_domain(self, domain: Any, name: str) -> torch.Tensor:
249        """
250        Encodes a domain from the domain data into the unimodal representation.
251
252        This is a convenient proxy for the `DomainModule.encode` method and is
253        equivalent to:
254        ```python
255        self.domain_mods[name].encode(domain)
256        ```
257
258        Args:
259            domain (`Any`): the domain data
260            name (`str`): domain name to encode
261
262        Returns:
263            `torch.Tensor`: the domain's unimodal representation.
264        """
265        return self.domain_mods[name].encode(domain)
266
267    def encode_domains(self, batch: RawDomainGroupsT) -> LatentsDomainGroupsDT:
268        """
269        Encode all domains in the batch.
270
271        Args:
272            batch (`RawDomainGroupsT`): the batch of
273                domain groups with raw unimodal data to encode into groups of latent
274                representations.
275
276        Returns:
277            `LatentsDomainGroupsDT`: the domains' unimodal representations.
278        """
279        return {
280            domains: {
281                name: self.domain_mods[name].encode(domain)
282                for name, domain in data.items()
283            }
284            for domains, data in batch.items()
285        }
286
287    def decode_domain(self, domain: torch.Tensor, name: str) -> Any:
288        """
289        Decodes a domain from the unimodal representation into the domain data.
290
291        This is a convenient proxy for the `DomainModule.encode` method and is
292        equivalent to:
293        ```python
294        self.domain_mods[name].decode(domain)
295        ```
296
297        Args:
298            domain (`torch.Tensor`): the domain data
299            name (`str`): domain name to encode
300
301        Returns:
302            `Any`: the domain's raw data.
303        """
304        return self.domain_mods[name].decode(domain)
305
306    def decode_domains(self, latents_domain: LatentsDomainGroupsT) -> RawDomainGroupsDT:
307        """
308        Decodes all domains in the batch.
309
310        Args:
311            batch (`LatentsDomainGroupsT`): the batch of
312                domain groups with unimodal latent representation to decode into
313                groups of raw data.
314
315        Returns:
316            `LatentsDomainGroupsDT`: the domains' raw data.
317        """
318        return {
319            domains: {
320                name: self.domain_mods[name].decode(domain)
321                for name, domain in latents.items()
322            }
323            for domains, latents in latents_domain.items()
324        }
325
326    def generic_step(self, batch: RawDomainGroupsT, mode: ModelModeT) -> torch.Tensor:
327        """
328        The generic step used in `training_step`, `validation_step` and
329        `test_step`.
330
331        Args:
332            batch (`RawDomainGroupsT`): the batch of groups of raw unimodal data.
333            mode (`ModelModeT`):
334
335        Returns:
336            `torch.Tensor`: the loss to train on.
337        """
338        domain_latents = self.encode_domains(batch)
339        batch_size = groups_batch_size(domain_latents)
340
341        loss_output = self.loss_mod.step(domain_latents, mode)
342
343        for name, metric in loss_output.all.items():
344            self.log(
345                f"{mode}/{name}",
346                metric,
347                batch_size=batch_size,
348                add_dataloader_idx=False,
349            )
350
351        return loss_output.loss
352
353    def validation_step(  # type: ignore
354        self, data: RawDomainGroupT, batch_idx: int, dataloader_idx: int = 0
355    ) -> torch.Tensor:
356        """Validation step used by lightning"""
357
358        batch = {frozenset(data.keys()): data}
359        for domain in data:
360            batch[frozenset([domain])] = {domain: data[domain]}
361        if dataloader_idx == 0:
362            return self.generic_step(batch, mode="val")
363        return self.generic_step(batch, mode="val/ood")
364
365    def test_step(  # type: ignore
366        self, data: Mapping[str, Any], batch_idx: int, dataloader_idx: int = 0
367    ) -> torch.Tensor:
368        """Test step used by lightning"""
369
370        batch = {frozenset(data.keys()): data}
371        for domain in data:
372            batch[frozenset([domain])] = {domain: data[domain]}
373        if dataloader_idx == 0:
374            return self.generic_step(batch, mode="test")
375        return self.generic_step(batch, mode="test/ood")
376
377    def training_step(  # type: ignore
378        self, batch: Mapping[frozenset[str], Mapping[str, Any]], batch_idx: int
379    ) -> torch.Tensor:
380        """Training step used by lightning"""
381
382        return self.generic_step(batch, mode="train")
383
384    def predict_step(  # type: ignore
385        self, data: Mapping[str, Any], batch_idx: int
386    ) -> GWPredictionsBase:
387        """Predict step used by lightning"""
388
389        batch = {frozenset(data.keys()): data}
390        for domain in data:
391            batch[frozenset([domain])] = {domain: data[domain]}
392
393        domain_latents = self.encode_domains(batch)
394        return self.forward(domain_latents)
395
396    def configure_optimizers(self) -> OptimizerLRSchedulerConfig:
397        """
398        Configure models optimizers.
399
400        Here we use `AdamW` for the optimizer and `OneCycleLR` for the learning-rate
401        scheduler.
402        """
403
404        optimizer = torch.optim.AdamW(
405            self.parameters(),
406            lr=self.optim_lr,
407            weight_decay=self.optim_weight_decay,
408        )
409
410        lr_scheduler = OneCycleLR(optimizer, **self.scheduler_args)
411
412        return {
413            "optimizer": optimizer,
414            "lr_scheduler": {
415                "scheduler": lr_scheduler,
416                "interval": "step",
417            },
418        }
419
420
421def freeze_domain_modules(
422    domain_mods: Mapping[str, DomainModule],
423) -> dict[str, DomainModule]:
424    """
425    Freezes weights and set to eval mode the domain modules.
426
427    .. note::
428        The output is casted as `dict[str, DomainModule]` type for better
429        auto-completion, but is actually a torch `ModuleDict`.
430
431    Args:
432        domain_mods (`Mapping[str, DomainModule]`): mapping of domain modules to freeze
433
434    Returns:
435        `ModuleDict`: frozen modules.
436    """
437
438    for mod in domain_mods.values():
439        mod.freeze()
440    # Cast for better auto-completion at the expense of ModuleDict
441    return cast(dict[str, DomainModule], ModuleDict(domain_mods))
442
443
444class GWPredictions(GWPredictionsBase):
445    """TypedDict of the output given when calling `GlobalWorkspaceBase.predict`"""
446
447    demi_cycles: dict[str, torch.Tensor]
448    """
449    Demi-cycle predictions of the model for each domain. Only computed on domain
450    groups with only one domain.
451    """
452
453    cycles: dict[tuple[str, str], torch.Tensor]
454    """
455    Cycle predictions of the model from one domain through another one.
456    Only computed on domain groups with more than one domain.
457    The keys are tuple with start domain and intermediary domain.
458    """
459
460    translations: dict[tuple[str, str], torch.Tensor]
461    """
462    Translation predictions of the model from one domain through another one.
463
464    Only computed on domain groups with more than one domain.
465    The keys are tuples with start domain and target domain.
466    """
467
468
469class GlobalWorkspace2Domains(
470    GlobalWorkspaceBase[GWModule, SingleDomainSelection, GWLosses2Domains]
471):
472    """
473    A simple 2-domains max flavor of GlobalWorkspaceBase.
474
475    This is used to simplify a Global Workspace instanciation and only overrides the
476    `__init__` method.
477    """
478
479    def __init__(
480        self,
481        domain_mods: Mapping[str, DomainModule],
482        gw_encoders: Mapping[str, Module],
483        gw_decoders: Mapping[str, Module],
484        workspace_dim: int,
485        loss_coefs: LossCoefs,
486        optim_lr: float = 1e-3,
487        optim_weight_decay: float = 0.0,
488        scheduler_args: SchedulerArgs | None = None,
489        learn_logit_scale: bool = False,
490        contrastive_loss: ContrastiveLossType | None = None,
491    ) -> None:
492        """
493        Initializes a Global Workspace
494
495        Args:
496            domain_mods (`Mapping[str, DomainModule]`): mapping of the domains
497                connected to the GW. Keys are domain names, values are the
498                `DomainModule`.
499            gw_encoders (`Mapping[str, torch.nn.Module]`): mapping for each domain
500                name to a `torch.nn.Module` class which role is to encode a
501                unimodal latent representations into a GW representation (pre fusion).
502            gw_decoders (`Mapping[str, torch.nn.Module]`): mapping for each domain
503                name to a `torch.nn.Module` class which role is to decode a
504                GW representation into a unimodal latent representations.
505            workspace_dim (`int`): dimension of the GW.
506            loss_coefs (`LossCoefs`): loss coefficients
507            optim_lr (`float`): learning rate
508            optim_weight_decay (`float`): weight decay
509            scheduler_args (`SchedulerArgs | None`): optimization scheduler's arguments
510            learn_logit_scale (`bool`): whether to learn the contrastive learning
511                contrastive loss when using the default contrastive loss.
512            contrastive_loss (`ContrastiveLossType | None`): a contrastive loss
513                function used for alignment. `learn_logit_scale` will not affect custom
514                contrastive losses.
515        """
516        domain_mods = freeze_domain_modules(domain_mods)
517
518        gw_mod = GWModule(domain_mods, workspace_dim, gw_encoders, gw_decoders)
519        if contrastive_loss is None:
520            contrastive_loss = ContrastiveLoss(
521                torch.tensor([1 / 0.07]).log(), "mean", learn_logit_scale
522            )
523        selection_mod = SingleDomainSelection()
524        loss_mod = GWLosses2Domains(
525            gw_mod, selection_mod, domain_mods, loss_coefs, contrastive_loss
526        )
527
528        super().__init__(
529            gw_mod,
530            selection_mod,
531            loss_mod,
532            optim_lr,
533            optim_weight_decay,
534            scheduler_args,
535        )
536
537    def forward(  # type: ignore
538        self,
539        latent_domains: LatentsDomainGroupsT,
540    ) -> GWPredictions:
541        """
542        Computes demi-cycles, cycles, and translations.
543
544        Args:
545            latent_domains (`LatentsT`): Groups of domains for the computation.
546
547        Returns:
548            `GWPredictions`: the predictions on the batch.
549        """
550        return GWPredictions(
551            demi_cycles=batch_demi_cycles(
552                self.gw_mod, self.selection_mod, latent_domains
553            ),
554            cycles=batch_cycles(
555                self.gw_mod, self.selection_mod, latent_domains, self.domain_mods.keys()
556            ),
557            translations=batch_translations(
558                self.gw_mod, self.selection_mod, latent_domains
559            ),
560            **super().forward(latent_domains),
561        )
562
563
564class GlobalWorkspace(GlobalWorkspaceBase[GWModule, RandomSelection, GWLosses]):
565    """The 2-domain fusion (with broadcast loss) flavor of GlobalWorkspaceBase.
566
567    This is used to simplify a Global Workspace instanciation and only overrides the
568    `__init__` method.
569    """
570
571    def __init__(
572        self,
573        domain_mods: Mapping[str, DomainModule],
574        gw_encoders: Mapping[str, Module],
575        gw_decoders: Mapping[str, Module],
576        workspace_dim: int,
577        loss_coefs: BroadcastLossCoefs,
578        selection_temperature: float = 0.2,
579        optim_lr: float = 1e-3,
580        optim_weight_decay: float = 0.0,
581        scheduler_args: SchedulerArgs | None = None,
582        learn_logit_scale: bool = False,
583        contrastive_loss: ContrastiveLossType | None = None,
584    ) -> None:
585        """
586        Initializes a Global Workspace
587
588        Args:
589            domain_mods (`Mapping[str, DomainModule]`): mapping of the domains
590                connected to the GW. Keys are domain names, values are the
591                `DomainModule`.
592            gw_encoders (`Mapping[str, torch.nn.Module]`): mapping for each domain
593                name to a `torch.nn.Module` class which role is to encode a
594                unimodal latent representations into a GW representation (pre fusion).
595            gw_decoders (`Mapping[str, torch.nn.Module]`): mapping for each domain
596                name to a `torch.nn.Module` class which role is to decode a
597                GW representation into a unimodal latent representations.
598            workspace_dim (`int`): dimension of the GW.
599            loss_coefs (`BroadcastLossCoefs`): loss coefs for the losses.
600            selection_temperature (`float`): temperature value for the RandomSelection
601                module.
602            optim_lr (`float`): learning rate
603            optim_weight_decay (`float`): weight decay
604            scheduler_args (`SchedulerArgs | None`): optimization scheduler's arguments
605            learn_logit_scale (`bool`): whether to learn the contrastive learning
606                contrastive loss when using the default contrastive loss.
607            contrastive_loss (`ContrastiveLossType | None`): a contrastive loss
608                function used for alignment. `learn_logit_scale` will not affect custom
609                contrastive losses.
610        """
611        domain_mods = freeze_domain_modules(domain_mods)
612        gw_mod = GWModule(domain_mods, workspace_dim, gw_encoders, gw_decoders)
613
614        if contrastive_loss is None:
615            contrastive_loss = ContrastiveLoss(
616                torch.tensor([1 / 0.07]).log(), "mean", learn_logit_scale
617            )
618
619        selection_mod = RandomSelection(selection_temperature)
620        loss_mod = GWLosses(
621            gw_mod, selection_mod, domain_mods, loss_coefs, contrastive_loss
622        )
623
624        super().__init__(
625            gw_mod,
626            selection_mod,
627            loss_mod,
628            optim_lr,
629            optim_weight_decay,
630            scheduler_args,
631        )
632
633    def forward(  # type: ignore
634        self,
635        latent_domains: LatentsDomainGroupsT,
636    ) -> GWPredictions:
637        """
638        Computes demi-cycles, cycles, and translations.
639
640        Args:
641            latent_domains (`LatentsT`): Groups of domains for the computation.
642
643        Returns:
644            `GWPredictions`: the predictions on the batch.
645        """
646        return GWPredictions(
647            demi_cycles=batch_demi_cycles(
648                self.gw_mod, self.selection_mod, latent_domains
649            ),
650            cycles=batch_cycles(
651                self.gw_mod, self.selection_mod, latent_domains, self.domain_mods.keys()
652            ),
653            translations=batch_translations(
654                self.gw_mod, self.selection_mod, latent_domains
655            ),
656            # TODO: add other combinations
657            **super().forward(latent_domains),
658        )
659
660
661class GlobalWorkspaceBayesian(
662    GlobalWorkspaceBase[GWModuleBayesian, FixedSharedSelection, GWLossesBayesian]
663):
664    """
665    A simple 2-domains max GlobalWorkspaceBase with a Bayesian base uncertainty
666    prediction.
667
668    This is used to simplify a Global Workspace instanciation and only overrides the
669    `__init__` method.
670    """
671
672    def __init__(
673        self,
674        domain_mods: Mapping[str, DomainModule],
675        gw_encoders: Mapping[str, Module],
676        gw_decoders: Mapping[str, Module],
677        workspace_dim: int,
678        loss_coefs: BroadcastLossCoefs,
679        sensitivity_selection: float = 1,
680        sensitivity_precision: float = 1,
681        optim_lr: float = 1e-3,
682        optim_weight_decay: float = 0.0,
683        scheduler_args: SchedulerArgs | None = None,
684        learn_logit_scale: bool = False,
685        use_normalized_constrastive: bool = True,
686        contrastive_loss: ContrastiveLossType | None = None,
687        precision_softmax_temp: float = 0.01,
688    ) -> None:
689        """
690        Initializes a Global Workspace
691
692        Args:
693            domain_mods (`Mapping[str, DomainModule]`): mapping of the domains
694                connected to the GW. Keys are domain names, values are the
695                `DomainModule`.
696            gw_encoders (`Mapping[str, torch.nn.Module]`): mapping for each domain
697                name to a `torch.nn.Module` class which role is to encode a
698                unimodal latent representations into a GW representation (pre fusion).
699            gw_decoders (`Mapping[str, torch.nn.Module]`): mapping for each domain
700                name to a `torch.nn.Module` class which role is to decode a
701                GW representation into a unimodal latent representations.
702            workspace_dim (`int`): dimension of the GW.
703            loss_coefs (`LossCoefs`): loss coefficients
704            sensitivity_selection (`float`): sensivity coef $c'_1$
705            sensitivity_precision (`float`): sensitivity coef $c'_2$
706            optim_lr (`float`): learning rate
707            optim_weight_decay (`float`): weight decay
708            scheduler_args (`SchedulerArgs | None`): optimization scheduler's arguments
709            learn_logit_scale (`bool`): whether to learn the contrastive learning
710                contrastive loss when using the default contrastive loss.
711            use_normalized_constrastive (`bool`): whether to use the normalized cont
712                loss by the precision coefs
713            contrastive_loss (`ContrastiveLossType | None`): a contrastive loss
714                function used for alignment. `learn_logit_scale` will not affect custom
715                contrastive losses.
716            precision_softmax_temp (`float`): temperature to use in softmax of
717                precision
718        """
719        domain_mods = freeze_domain_modules(domain_mods)
720
721        gw_mod = GWModuleBayesian(
722            domain_mods,
723            workspace_dim,
724            gw_encoders,
725            gw_decoders,
726            sensitivity_selection,
727            sensitivity_precision,
728            precision_softmax_temp,
729        )
730
731        selection_mod = FixedSharedSelection()
732
733        contrastive_loss = ContrastiveLoss(
734            torch.tensor([1]).log(), "mean", learn_logit_scale
735        )
736
737        loss_mod = GWLossesBayesian(
738            gw_mod,
739            selection_mod,
740            domain_mods,
741            loss_coefs,
742            contrastive_loss,
743            use_normalized_constrastive,
744        )
745
746        super().__init__(
747            gw_mod,
748            selection_mod,
749            loss_mod,
750            optim_lr,
751            optim_weight_decay,
752            scheduler_args,
753        )
754
755    def forward(  # type: ignore
756        self,
757        latent_domains: LatentsDomainGroupsT,
758    ) -> GWPredictions:
759        """
760        Computes demi-cycles, cycles, and translations.
761
762        Args:
763            latent_domains (`LatentsT`): Groups of domains for the computation.
764
765        Returns:
766            `GWPredictions`: the predictions on the batch.
767        """
768        return GWPredictions(
769            demi_cycles=batch_demi_cycles(
770                self.gw_mod, self.selection_mod, latent_domains
771            ),
772            cycles=batch_cycles(
773                self.gw_mod, self.selection_mod, latent_domains, self.domain_mods.keys()
774            ),
775            translations=batch_translations(
776                self.gw_mod, self.selection_mod, latent_domains
777            ),
778            **super().forward(latent_domains),
779        )
780
781
782def pretrained_global_workspace(
783    checkpoint_path: str | Path,
784    domain_mods: Mapping[str, DomainModule],
785    gw_encoders: Mapping[str, Module],
786    gw_decoders: Mapping[str, Module],
787    workspace_dim: int,
788    loss_coefs: LossCoefs,
789    contrastive_fn: ContrastiveLossType,
790    **kwargs,
791) -> GlobalWorkspace2Domains:
792    """
793    Load a `GlobalWorkspace` flavor of `GlobalWorkspaceBase` from a checkpoint.
794
795    Args:
796        checkpoint_path (`str | Path`): path to checkpoint
797        domain_mods (`Mapping[str, DomainModule]`): mapping of the domains
798            connected to the GW. Keys are domain names, values are the
799            `DomainModule`.
800        gw_encoders (`Mapping[str, torch.nn.Module]`): mapping for each domain
801            name to a `torch.nn.Module` class which role is to encode a
802            unimodal latent representations into a GW representation (pre fusion).
803        gw_decoders (`Mapping[str, torch.nn.Module]`): mapping for each domain
804            name to a `torch.nn.Module` class which role is to decode a
805            GW representation into a unimodal latent representations.
806        workspace_dim (`int`): dimension of the GW.
807        loss_coefs (`LossCoefs`): loss coefficients
808        contrastive_loss (`ContrastiveLossType`): a contrastive loss
809            function used for alignment. `learn_logit_scale` will not affect custom
810            contrastive losses.
811        **kwargs: additional arguments to pass to
812            `GlobalWorkspace.load_from_checkpoint`.
813
814    Returns:
815        `GlobalWorkspace`: the pretrained `GlobalWorkspace`.
816
817    Raises:
818        `TypeError`: if loaded type is not `GlobalWorkspace`.
819    """
820    domain_mods = freeze_domain_modules(domain_mods)
821    gw_mod = GWModule(domain_mods, workspace_dim, gw_encoders, gw_decoders)
822    selection_mod = SingleDomainSelection()
823    loss_mod = GWLosses2Domains(
824        gw_mod, selection_mod, domain_mods, loss_coefs, contrastive_fn
825    )
826
827    gw = GlobalWorkspace2Domains.load_from_checkpoint(
828        checkpoint_path,
829        gw_mod=gw_mod,
830        selection_mid=selection_mod,
831        loss_coefs=loss_coefs,
832        loss_mod=loss_mod,
833        **kwargs,
834    )
835    if not isinstance(gw, GlobalWorkspace2Domains):
836        raise TypeError("model should be of type GlobalWorkspace")
837    return gw