# Path: cortex/built_ins/networks/utils.py
# def apply_nonlinearity(x, nonlinearity, **nonlinearity_args):
#     if nonlinearity:
#         if isinstance(nonlinearity, str):
#             nonlinearity = get_nonlinearity(nonlinearity)
#         if callable(nonlinearity):
#             if isinstance(nonlinearity, nn.PReLU):
#                 nonlinearity.to(x.device)
#             try:
#                 x = nonlinearity(x, **nonlinearity_args)
#             except RuntimeError:
#                 nonlinearity.to('cpu')
#                 x = nonlinearity(x, **nonlinearity_args)
#         else:
#             raise ValueError(nonlinearity, type(nonlinearity))
#     return x
# 
# def get_nonlinearity(nonlinearity=None):
#     if not nonlinearity:
#         pass
#     elif callable(nonlinearity):
#         if nonlinearity == nn.LeakyReLU:
#             nonlinearity = nonlinearity(0.02, inplace=True)
#     elif hasattr(nn, nonlinearity):
#         nonlinearity = getattr(nn, nonlinearity)
#         if nonlinearity == 'LeakyReLU':
#             nonlinearity = nonlinearity(0.02, inplace=True)
#         else:
#             nonlinearity = nonlinearity()
#     elif hasattr(nn.functional, nonlinearity):
#         nonlinearity = getattr(nn.functional, nonlinearity)
#     else:
#         raise ValueError(nonlinearity)
#     return nonlinearity
# 
# def finish_layer_1d(models, name, 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.Dropout(p=dropout))
# 
#     if layer_norm:
#         models.add_module(name + '_ln', nn.LayerNorm(dim_out))
#     elif batch_norm:
#         models.add_module(name + '_bn', nn.BatchNorm1d(dim_out))
# 
#     if nonlinearity:
#         nonlinearity = get_nonlinearity(nonlinearity)
#         models.add_module(
#             '{}_{}'.format(name, nonlinearity.__class__.__name__),
#             nonlinearity)


import torch.nn as nn
import torch
from .utils import apply_nonlinearity, get_nonlinearity, finish_layer_1d
        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))

# Path: cortex/built_ins/networks/utils.py
# def apply_nonlinearity(x, nonlinearity, **nonlinearity_args):
#     if nonlinearity:
#         if isinstance(nonlinearity, str):
#             nonlinearity = get_nonlinearity(nonlinearity)
#         if callable(nonlinearity):
#             if isinstance(nonlinearity, nn.PReLU):
#                 nonlinearity.to(x.device)
#             try:
#                 x = nonlinearity(x, **nonlinearity_args)
#             except RuntimeError:
#                 nonlinearity.to('cpu')
#                 x = nonlinearity(x, **nonlinearity_args)
#         else:
#             raise ValueError(nonlinearity, type(nonlinearity))
#     return x
# 
# def get_nonlinearity(nonlinearity=None):
#     if not nonlinearity:
#         pass
#     elif callable(nonlinearity):
#         if nonlinearity == nn.LeakyReLU:
#             nonlinearity = nonlinearity(0.02, inplace=True)
#     elif hasattr(nn, nonlinearity):
#         nonlinearity = getattr(nn, nonlinearity)
#         if nonlinearity == 'LeakyReLU':
#             nonlinearity = nonlinearity(0.02, inplace=True)
#         else:
#             nonlinearity = nonlinearity()
#     elif hasattr(nn.functional, nonlinearity):
#         nonlinearity = getattr(nn.functional, nonlinearity)
#     else:
#         raise ValueError(nonlinearity)
#     return nonlinearity
# 
# def finish_layer_1d(models, name, 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.Dropout(p=dropout))
# 
#     if layer_norm:
#         models.add_module(name + '_ln', nn.LayerNorm(dim_out))
#     elif batch_norm:
#         models.add_module(name + '_bn', nn.BatchNorm1d(dim_out))
# 
#     if nonlinearity:
#         nonlinearity = get_nonlinearity(nonlinearity)
#         models.add_module(
#             '{}_{}'.format(name, nonlinearity.__class__.__name__),
#             nonlinearity)


# Path: cortex/built_ins/networks/base_network.py
import torch.nn as nn
import torch
from .utils import apply_nonlinearity, get_nonlinearity, finish_layer_1d

        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(