# Path: cortex/built_ins/networks/convnets.py
# class SimpleNet(nn.Module):
# 
#     def __init__(self):
#         super(SimpleNet, self).__init__()
#         self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
#         self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
#         self.conv2_drop = nn.Dropout2d()
#         self.fc1 = nn.Linear(320, 50)
#         self.fc2 = nn.Linear(50, 10)
# 
#     def forward(self, x):
#         x = F.relu(F.max_pool2d(self.conv1(x), 2))
#         x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
#         x = x.view(-1, 320)
#         x = F.relu(self.fc1(x))
#         x = F.dropout(x, training=self.training)
#         x = self.fc2(x)
#         return F.log_softmax(x)
# 
# def infer_conv_size(w, k, s, p):
#     x = (w - k + 2 * p) // s + 1
#     return x
# 
# 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 cortex.built_ins.networks.convnets import SimpleNet, infer_conv_size
from torch import nn
from cortex.built_ins.networks.modules import View
import torch
    # (conv_(64 / 128)_2_ReLU): ReLU()
    assert isinstance(layers[7][1],
                      nn.ReLU) and layers[7][0] == 'conv_(64/128)_2_ReLU'

    # (conv_(128/256)_3):
    # Conv2d(128, 256, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1), bias=False)
    assert isinstance(layers[8][1],
                      nn.Conv2d) and layers[8][0] == 'conv_(128/256)_3'
    assert layers[8][1].kernel_size == (4, 4) and layers[8][1].stride == (
        2, 2) and layers[4][1].padding == (1, 1) and layers[8][1].bias is None

    # (conv_(128/256)_3_do): Dropout2d(p=0.2)
    assert isinstance(layers[9][1],
                      nn.Dropout2d) and layers[9][0] == 'conv_(128/256)_3_do'
    assert layers[9][1].p == 0.2

    # (conv_(128 / 256)_3_bn):
    # BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    assert isinstance(layers[10][1],
                      nn.BatchNorm2d) and layers[10][0] == 'conv_(128/256)_3_bn'
    assert layers[10][1].eps == 1e-05 and layers[10][1].momentum == 0.1
    assert layers[10][1].affine is True
    assert layers[10][1].track_running_stats is True

    # (conv_(128/256)_3_ReLU): ReLU()
    assert isinstance(layers[11][1],
                      nn.ReLU) and layers[11][0] == 'conv_(128/256)_3_ReLU'

    # (final_reshape_4x4x256to4096): View()
    assert isinstance(layers[12][1],

# Path: cortex/built_ins/networks/convnets.py
# class SimpleNet(nn.Module):
# 
#     def __init__(self):
#         super(SimpleNet, self).__init__()
#         self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
#         self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
#         self.conv2_drop = nn.Dropout2d()
#         self.fc1 = nn.Linear(320, 50)
#         self.fc2 = nn.Linear(50, 10)
# 
#     def forward(self, x):
#         x = F.relu(F.max_pool2d(self.conv1(x), 2))
#         x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
#         x = x.view(-1, 320)
#         x = F.relu(self.fc1(x))
#         x = F.dropout(x, training=self.training)
#         x = self.fc2(x)
#         return F.log_softmax(x)
# 
# def infer_conv_size(w, k, s, p):
#     x = (w - k + 2 * p) // s + 1
#     return x
# 
# 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_convnets.py
from cortex.built_ins.networks.convnets import SimpleNet, infer_conv_size
from torch import nn
from cortex.built_ins.networks.modules import View
import torch

    # (conv_(64 / 128)_2_ReLU): ReLU()
    assert isinstance(layers[7][1],
                      nn.ReLU) and layers[7][0] == 'conv_(64/128)_2_ReLU'

    # (conv_(128/256)_3):
    # Conv2d(128, 256, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1), bias=False)
    assert isinstance(layers[8][1],
                      nn.Conv2d) and layers[8][0] == 'conv_(128/256)_3'
    assert layers[8][1].kernel_size == (4, 4) and layers[8][1].stride == (
        2, 2) and layers[4][1].padding == (1, 1) and layers[8][1].bias is None

    # (conv_(128/256)_3_do): Dropout2d(p=0.2)
    assert isinstance(layers[9][1],
                      nn.Dropout2d) and layers[9][0] == 'conv_(128/256)_3_do'
    assert layers[9][1].p == 0.2

    # (conv_(128 / 256)_3_bn):
    # BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    assert isinstance(layers[10][1],
                      nn.BatchNorm2d) and layers[10][0] == 'conv_(128/256)_3_bn'
    assert layers[10][1].eps == 1e-05 and layers[10][1].momentum == 0.1
    assert layers[10][1].affine is True
    assert layers[10][1].track_running_stats is True

    # (conv_(128/256)_3_ReLU): ReLU()
    assert isinstance(layers[11][1],
                      nn.ReLU) and layers[11][0] == 'conv_(128/256)_3_ReLU'

    # (final_reshape_4x4x256to4096): View()
    assert isinstance(layers[12][1],

                      View) and layers[12][0] == 'final_reshape_4x4x256to4096'