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

__all__ = ['ProjectionWrapper']

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

# Cell
class ProjectionWrapper(Module):
    def __init__(self, arch, c_in, c_out, seq_len, new_c_in=None, new_seq_len=None, device=None, **kwargs):

        device = ifnone(device, default_device())
        new_c_in = ifnone(new_c_in, c_in)
        new_seq_len = ifnone(new_seq_len, seq_len)
        self.new_shape = c_in != new_c_in or seq_len != new_seq_len
        if self.new_shape:
            layers = []
            if c_in != new_c_in:
                lin = nn.Linear(c_in, new_c_in)
                nn.init.constant_(lin.weight, 0)
                layers += [Transpose(1,2), lin, Transpose(1,2)]
            if seq_len != new_seq_len:
                lin2 = nn.Linear(seq_len, new_seq_len)
                nn.init.constant_(lin2.weight, 0)
                layers += [lin2]
            self.new_shape_fn = nn.Sequential(*layers)
        self.model = build_ts_model(arch, c_in=new_c_in, c_out=c_out, seq_len=new_seq_len, **kwargs)
        self.to(device)
    def forward(self, x):
        if self.new_shape: x = self.new_shape_fn(x)
        return self.model(x)