# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/130_models.MultiInputNet.ipynb (unless otherwise specified).

__all__ = ['MultiInputNet']

# Cell
from ..imports import *
from .layers import *

# Cell
class MultiInputNet(Module):
    def __init__(self, dls, models, device=None, flatten=True, custom_head=None, **kwargs):
        r"""
        Args:
            dls: DataLoaders of type MixedDataLoaders.
            models: list of models (one model per dataloader in dls).
            flatten: if True, the output of each model body will be flattend before concatenating and passing to a joint head.
            device: cpu or cuda. If None, default_device() will be chosen.
            custom_head: type of thead that will be applied.
            kwargs: custom_head kwargs
        """

        head = ifnone(custom_head, mlp_head)
        device = ifnone(device, default_device())
        self.models = nn.ModuleList()
        for m in L(models):
            m.head = Identity()
            self.models.append(m)
        self.flatten = Reshape(-1) if flatten else None
        self.concat = Concat(dim=1)
        with torch.no_grad():
            self.head = Noop
            out = self.forward(first(dls.train)[0])
            self.head_nf = out.shape[-1] if flatten else out.shape[1]
        self.head = head(self.head_nf, dls.c, **kwargs)
        self.to(device=device)

    def forward(self, xs):
        for i, (x,m) in enumerate(zip(xs, self.models)):
            _out = m(*x) if isinstance(x, L) else m(x)
            if self.flatten is not None and _out.ndim == 3: _out = self.flatten(_out)
            out = _out if i==0 else self.concat([out, _out])
        return self.head(out)