# Path: cortex/built_ins/networks/utils.py
# 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 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 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)
# 
# 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/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)


from collections import OrderedDict
from cortex.built_ins.networks.utils import get_nonlinearity, \
    apply_nonlinearity, finish_layer_1d, finish_layer_2d
from torch import nn
from cortex.built_ins.networks.modules import View
import torch
                                bias=True))]))

    finish_layer_1d(model, name, dim_out, dropout, layer_norm, batch_norm,
                    nonlinearity['relu'])

    assert isinstance(model[0], torch.nn.modules.linear.Linear)
    assert isinstance(model[1], torch.nn.modules.batchnorm.BatchNorm1d)
    assert isinstance(model[2], torch.nn.modules.activation.ReLU)
    assert model[0].in_features == 64 and model[0].out_features == 4096


def test_finish_layer_2d():

    """

    Asserts: True if right GAN layers are being added.

    """

    # Test settings for a GAN
    dim_x = 4
    dim_y = 4
    dim_out = 256
    dropout = False
    layer_norm = False
    batch_norm = True
    nonlinearity = nn.ReLU()
    name = 'reshape'
    batch_norm_1d_layer = nn.BatchNorm1d(4096)
    nonlinear_relu_layer = nn.ReLU()

# Path: cortex/built_ins/networks/utils.py
# 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 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 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)
# 
# 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/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: tests/built_ins/networks/test_network_utils.py
from collections import OrderedDict
from cortex.built_ins.networks.utils import get_nonlinearity, \
    apply_nonlinearity, finish_layer_1d, finish_layer_2d
from torch import nn
from cortex.built_ins.networks.modules import View
import torch

                                bias=True))]))

    finish_layer_1d(model, name, dim_out, dropout, layer_norm, batch_norm,
                    nonlinearity['relu'])

    assert isinstance(model[0], torch.nn.modules.linear.Linear)
    assert isinstance(model[1], torch.nn.modules.batchnorm.BatchNorm1d)
    assert isinstance(model[2], torch.nn.modules.activation.ReLU)
    assert model[0].in_features == 64 and model[0].out_features == 4096


def test_finish_layer_2d():

    """

    Asserts: True if right GAN layers are being added.

    """

    # Test settings for a GAN
    dim_x = 4
    dim_y = 4
    dim_out = 256
    dropout = False
    layer_norm = False
    batch_norm = True
    nonlinearity = nn.ReLU()
    name = 'reshape'
    batch_norm_1d_layer = nn.BatchNorm1d(4096)
    nonlinear_relu_layer = nn.ReLU()

    view = View()