# Path: cortex/built_ins/networks/modules.py
# class View(nn.Module):
#     def __init__(self, *shape):
#         super(View, self).__init__()
#         self.shape = shape
# 
#     def forward(self, input):
#         return input.view(*self.shape)
# 
# Path: cortex/built_ins/networks/base_network.py
# class BaseNet(nn.Module):
#     '''Basic convenience network for cortex.
# 
#     Attributes:
#         models: A sequence of
# 
#     '''
# 
#     def __init__(self, nonlinearity='ReLU', output_nonlinearity=None):
#         super(BaseNet, self).__init__()
# 
#         self.models = nn.Sequential()
# 
#         self.output_nonlinearity = output_nonlinearity
#         self.layer_nonlinearity = get_nonlinearity(nonlinearity)
# 
#     def forward(self,
#                 x: torch.Tensor,
#                 nonlinearity: str = None,
#                 **nonlinearity_args: dict) -> torch.Tensor:
#         self.states = []
#         if nonlinearity is None:
#             nonlinearity = self.output_nonlinearity
#         elif not nonlinearity:
#             nonlinearity = None
# 
#         for model in self.models:
#             x = model(x)
#             self.states.append(x)
#         x = apply_nonlinearity(x, nonlinearity, **nonlinearity_args)
#         return x
# 
#     def get_h(self, dim_h, n_levels=None):
#         if isinstance(dim_h, (list, tuple)):
#             pass
#         elif n_levels:
#             dim_h = [dim_h for _ in range(n_levels)]
#         else:
#             dim_h = [dim_h]
# 
#         return dim_h
# 
#     def add_linear_layers(self,
#                           dim_in,
#                           dim_h,
#                           dim_ex=None,
#                           Linear=None,
#                           **layer_args):
#         Linear = Linear or nn.Linear
# 
#         if dim_h is None or len(dim_h) == 0:
#             return dim_in
# 
#         for dim_out in dim_h:
#             name = 'linear_({}/{})'.format(dim_in, dim_out)
#             self.models.add_module(name, Linear(dim_in, dim_out))
#             finish_layer_1d(
#                 self.models,
#                 name,
#                 dim_out,
#                 nonlinearity=self.layer_nonlinearity,
#                 **layer_args)
#             dim_in = dim_out
#             if dim_ex is not None:
#                 dim_in += dim_ex
# 
#         return dim_out
# 
#     def add_output_layer(self, dim_in, dim_out, Linear=None):
# 
#         Linear = Linear or nn.Linear
#         if dim_out is not None:
#             name = 'linear_({}/{})_{}'.format(dim_in, dim_out, 'out')
#             self.models.add_module(name, Linear(dim_in, dim_out))
# 
# Path: cortex/built_ins/networks/utils.py
# def finish_layer_2d(models, name, dim_x, dim_y, dim_out,
#                     dropout=False, layer_norm=False, batch_norm=False,
#                     nonlinearity=None):
#     if layer_norm and batch_norm:
#         logger.warning('Ignoring layer_norm because batch_norm is True')
# 
#     if dropout:
#         models.add_module(name + '_do', nn.Dropout2d(p=dropout))
# 
#     if layer_norm:
#         models.add_module(name + '_ln', nn.LayerNorm((dim_out, dim_x, dim_y)))
#     elif batch_norm:
#         models.add_module(name + '_bn', nn.BatchNorm2d(dim_out))
# 
#     if nonlinearity:
#         nonlinearity = get_nonlinearity(nonlinearity)
#         models.add_module(
#             '{}_{}'.format(name, nonlinearity.__class__.__name__),
#             nonlinearity)


import logging
import torch.nn as nn
from .modules import View
from .base_network import BaseNet
from .utils import finish_layer_2d
            nonlinearity=nonlinearity, output_nonlinearity=output_nonlinearity)

        dim_h_ = dim_h
        logger.debug('Input shape: {}'.format(shape))
        dim_x_, dim_y_, dim_out_ = shape

        dim_x = dim_x_
        dim_y = dim_y_
        dim_h = dim_h_

        saved_spatial_dimensions = [(dim_x, dim_y)]

        for n in range(n_steps):
            dim_x, dim_y = self.next_size(dim_x, dim_y, f_size, stride, pad)
            saved_spatial_dimensions.append((dim_x, dim_y))
            if n < n_steps - 1:
                dim_h *= 2

        dim_out = dim_x * dim_y * dim_h

        if initial_layer is not None:
            dim_h_ = [initial_layer, dim_out]
        else:
            dim_h_ = [dim_out]

        self.add_linear_layers(dim_in, dim_h_, **layer_args)

        name = 'reshape'
        self.models.add_module(name, View(-1, dim_h, dim_x, dim_y))


# Path: cortex/built_ins/networks/modules.py
# class View(nn.Module):
#     def __init__(self, *shape):
#         super(View, self).__init__()
#         self.shape = shape
# 
#     def forward(self, input):
#         return input.view(*self.shape)
# 
# Path: cortex/built_ins/networks/base_network.py
# class BaseNet(nn.Module):
#     '''Basic convenience network for cortex.
# 
#     Attributes:
#         models: A sequence of
# 
#     '''
# 
#     def __init__(self, nonlinearity='ReLU', output_nonlinearity=None):
#         super(BaseNet, self).__init__()
# 
#         self.models = nn.Sequential()
# 
#         self.output_nonlinearity = output_nonlinearity
#         self.layer_nonlinearity = get_nonlinearity(nonlinearity)
# 
#     def forward(self,
#                 x: torch.Tensor,
#                 nonlinearity: str = None,
#                 **nonlinearity_args: dict) -> torch.Tensor:
#         self.states = []
#         if nonlinearity is None:
#             nonlinearity = self.output_nonlinearity
#         elif not nonlinearity:
#             nonlinearity = None
# 
#         for model in self.models:
#             x = model(x)
#             self.states.append(x)
#         x = apply_nonlinearity(x, nonlinearity, **nonlinearity_args)
#         return x
# 
#     def get_h(self, dim_h, n_levels=None):
#         if isinstance(dim_h, (list, tuple)):
#             pass
#         elif n_levels:
#             dim_h = [dim_h for _ in range(n_levels)]
#         else:
#             dim_h = [dim_h]
# 
#         return dim_h
# 
#     def add_linear_layers(self,
#                           dim_in,
#                           dim_h,
#                           dim_ex=None,
#                           Linear=None,
#                           **layer_args):
#         Linear = Linear or nn.Linear
# 
#         if dim_h is None or len(dim_h) == 0:
#             return dim_in
# 
#         for dim_out in dim_h:
#             name = 'linear_({}/{})'.format(dim_in, dim_out)
#             self.models.add_module(name, Linear(dim_in, dim_out))
#             finish_layer_1d(
#                 self.models,
#                 name,
#                 dim_out,
#                 nonlinearity=self.layer_nonlinearity,
#                 **layer_args)
#             dim_in = dim_out
#             if dim_ex is not None:
#                 dim_in += dim_ex
# 
#         return dim_out
# 
#     def add_output_layer(self, dim_in, dim_out, Linear=None):
# 
#         Linear = Linear or nn.Linear
#         if dim_out is not None:
#             name = 'linear_({}/{})_{}'.format(dim_in, dim_out, 'out')
#             self.models.add_module(name, Linear(dim_in, dim_out))
# 
# Path: cortex/built_ins/networks/utils.py
# def finish_layer_2d(models, name, dim_x, dim_y, dim_out,
#                     dropout=False, layer_norm=False, batch_norm=False,
#                     nonlinearity=None):
#     if layer_norm and batch_norm:
#         logger.warning('Ignoring layer_norm because batch_norm is True')
# 
#     if dropout:
#         models.add_module(name + '_do', nn.Dropout2d(p=dropout))
# 
#     if layer_norm:
#         models.add_module(name + '_ln', nn.LayerNorm((dim_out, dim_x, dim_y)))
#     elif batch_norm:
#         models.add_module(name + '_bn', nn.BatchNorm2d(dim_out))
# 
#     if nonlinearity:
#         nonlinearity = get_nonlinearity(nonlinearity)
#         models.add_module(
#             '{}_{}'.format(name, nonlinearity.__class__.__name__),
#             nonlinearity)


# Path: cortex/built_ins/networks/conv_decoders.py
import logging
import torch.nn as nn
from .modules import View
from .base_network import BaseNet
from .utils import finish_layer_2d

            nonlinearity=nonlinearity, output_nonlinearity=output_nonlinearity)

        dim_h_ = dim_h
        logger.debug('Input shape: {}'.format(shape))
        dim_x_, dim_y_, dim_out_ = shape

        dim_x = dim_x_
        dim_y = dim_y_
        dim_h = dim_h_

        saved_spatial_dimensions = [(dim_x, dim_y)]

        for n in range(n_steps):
            dim_x, dim_y = self.next_size(dim_x, dim_y, f_size, stride, pad)
            saved_spatial_dimensions.append((dim_x, dim_y))
            if n < n_steps - 1:
                dim_h *= 2

        dim_out = dim_x * dim_y * dim_h

        if initial_layer is not None:
            dim_h_ = [initial_layer, dim_out]
        else:
            dim_h_ = [dim_out]

        self.add_linear_layers(dim_in, dim_h_, **layer_args)

        name = 'reshape'
        self.models.add_module(name, View(-1, dim_h, dim_x, dim_y))


        finish_layer_2d(self.models, name, dim_x, dim_y, dim_h,