from torch import nn
import numpy as np
import base.basenetwork as BaseN
from networks.cholesky import CholeskyBlock

class FCNet(BaseN.BaseNetwork):
    name ="FCNet"
    def __init__(self,input_shape,output_shape,owner_name=""):
        super(FCNet,self).__init__(input_shape,output_shape,owner_name)
        
        x = input_shape
        self.model = nn.Sequential(BaseN.Flatten(),
                                   nn.Linear(np.prod(x), 1024),nn.Softplus(),
                                   nn.Linear(1024,512),nn.Tanh(),
                                   nn.Linear(512,256),
                                   BaseN.EigenLayer(256,self.output_shape[0]))
        self.compile()

class FCSpectralNet(BaseN.BaseNetwork):
    name ="FCSpectralNet"
    def __init__(self,input_shape,output_shape,owner_name=""):
        super(FCSpectralNet,self).__init__(input_shape,output_shape,owner_name)
        
        x = input_shape
        self.model = nn.Sequential(BaseN.Flatten(),
                                   nn.Linear(np.prod(x), 1024),BaseN.AdaptiveTanh(),
                                   nn.Linear(1024,1024),BaseN.AdaptiveTanh(),
                                   nn.Linear(1024,512),BaseN.AdaptiveTanh(),
                                   BaseN.EigenLayer(512,self.output_shape[0]))
        self.compile()
class FCSpectralMNet(BaseN.BaseNetwork):
    name ="FCSpectralMNet"
    def __init__(self,input_shape,output_shape,owner_name=""):
        super(FCSpectralMNet,self).__init__(input_shape,output_shape,owner_name)
        
        x = input_shape
        self.model = nn.Sequential(BaseN.Flatten(),
                                   nn.Linear(np.prod(x), 1024),nn.ReLU(),
                                   nn.Linear(1024,1024),nn.ReLU(),
                                   nn.Linear(1024,512),nn.ReLU(),
                                   nn.Linear(512,self.output_shape[0]-1),nn.Tanh(),
                                   BaseN.EigenLayer())
        self.compile()
class FCNetQ(BaseN.BaseNetwork):
    name ="FCNetQ"
    def __init__(self,input_shape,output_shape,owner_name=""):
        super(FCNetQ,self).__init__(input_shape,output_shape,owner_name)
        
        x = int(np.prod(input_shape))
        self.model = nn.Sequential(BaseN.Flatten(),
                                   nn.Linear(x,x),nn.Tanh(),
                                   nn.Linear(x,self.output_shape[0]))
        self.compile()


class ConvNet(BaseN.BaseNetwork):
    name="ConvNet"
    def __init__(self,input_shape,output_shape,owner_name=""):
        super(ConvNet,self).__init__(input_shape,output_shape,owner_name)
        
        self.conv = [nn.Sequential(BaseN.conv3_2(input_shape[0], 8),nn.ReLU(),
                                                 BaseN.conv3_2(8, 16),nn.ReLU(),
                                                 BaseN.conv3_2(8, 8))]
        x = BaseN.output_shape(self.conv[0],input_shape)
        self.model = nn.Sequential(self.conv[0],
                                   BaseN.Flatten(),
                                   nn.Linear(np.prod(x), 512),BaseN.AdaptiveTanh(),
                                   nn.Linear(512,256),
                                   BaseN.EigenLayer(256,self.output_shape[0],bias=False))
        self.compile()
class ConvNetBias(BaseN.BaseNetwork):
    name="ConvNetBias"
    def __init__(self,input_shape,output_shape,owner_name=""):
        super(ConvNetBias,self).__init__(input_shape,output_shape,owner_name)
        
        self.conv = [nn.Sequential(BaseN.conv3_2(input_shape[0], 8),nn.Softplus(),
                                   BaseN.conv3_2(8, 12),BaseN.AdaptiveTanh(),
                                   BaseN.conv3_2(12, 16),
                                   BaseN.conv3_2(16, 20))]
        x = BaseN.output_shape(self.conv[0],input_shape)
        self.model = nn.Sequential(self.conv[0],
                                   BaseN.Flatten(),
                                   nn.Linear(np.prod(x), 512),BaseN.AdaptiveTanh(),
                                   nn.Linear(512,256),
                                   BaseN.EigenLayer(256,self.output_shape[0],bias=False))
        self.compile()

class FCConvNet(BaseN.BaseNetwork):
    name="FCConvNet"
    def __init__(self,input_shape,output_shape,owner_name=""):
        super(FCConvNet,self).__init__(input_shape,output_shape,owner_name)
        
        self.conv = [nn.Sequential(BaseN.conv3_2(input_shape[0], 4),nn.Softplus(),
                                   BaseN.conv3_2(4, 8),BaseN.AdaptiveTanh())]
        x = BaseN.output_shape(self.conv[0],input_shape)
        self.model = nn.Sequential(self.conv[0],
                                   BaseN.Flatten(),
                                   nn.Linear(np.prod(x), 512),
                                   nn.Linear(512,512),BaseN.AdaptiveTanh(),
                                   nn.Linear(512,256),
                                   BaseN.EigenLayer(256,self.output_shape[0],bias=False))
        self.compile()
        
class FCConvNetBias(BaseN.BaseNetwork):
    name="FCConvNetBias"
    def __init__(self,input_shape,output_shape,owner_name=""):
        super(FCConvNetBias,self).__init__(input_shape,output_shape,owner_name)
        
        self.conv = [nn.Sequential(BaseN.conv3_2(input_shape[0], 4),nn.ReLU(),
                                   BaseN.conv3_2(4, 4),BaseN.AdaptiveTanh())]
        x = BaseN.output_shape(self.conv[0],input_shape)
        self.model = nn.Sequential(self.conv[0],
                                   BaseN.Flatten(),
                                   nn.Linear(np.prod(x), 512),
                                   nn.Linear(512,1024),BaseN.AdaptiveTanh(),
                                   nn.Linear(1024,256),
                                   BaseN.EigenLayer(256,self.output_shape[0],bias=False))
        self.compile()
        
class ConvNet2(BaseN.BaseNetwork):
    name="ConvNet2"
    def __init__(self,input_shape,output_shape,owner_name=""):
        super(ConvNet2,self).__init__(input_shape,output_shape,owner_name)
        
        self.conv = [nn.Sequential(BaseN.conv3_2(input_shape[0], 3),nn.Softplus(),
                                                 BaseN.conv3_2(3, 6),BaseN.conv3_2(6, 12))]
        x = BaseN.output_shape(self.conv[0],input_shape)
        self.model = nn.Sequential(self.conv[0],
                                   BaseN.Flatten(),
                                   nn.Linear(np.prod(x), 512),
                                   nn.Linear(512,256),nn.Tanh(),
                                   nn.Linear(256,512),
                                   nn.Linear(512,1024),nn.Tanh(),
                                   nn.Linear(1024,512),
                                   nn.Linear(512,256),nn.Tanh(),
                                   nn.Linear(256,256),
                                   BaseN.EigenLayer(256,self.output_shape[0]))
        self.compile()
class ConvNetBig(BaseN.BaseNetwork):
    name="ConvNetBig"
    def __init__(self,input_shape,output_shape,owner_name=""):
        super(ConvNetBig,self).__init__(input_shape,output_shape,owner_name)
        
        self.conv = [nn.Sequential(BaseN.conv3_2(input_shape[0], 8),nn.Softplus(),
                                                 BaseN.conv3_2(8, 16),nn.Softplus(),
                                                 BaseN.conv3_2(16, 32))]
        x = BaseN.output_shape(self.conv[0],input_shape)
        self.model = nn.Sequential(self.conv[0],
                                   BaseN.Flatten(),
                                   nn.Linear(np.prod(x), 512),
                                   nn.Linear(512,256),nn.Tanh(),
                                   nn.Linear(256,512),
                                   BaseN.EigenLayer(512,self.output_shape[0]))
        self.compile()

class ConvNetBigBias(BaseN.BaseNetwork):
    name="ConvNetBigBias"
    def __init__(self,input_shape,output_shape,owner_name=""):
        super(ConvNetBigBias,self).__init__(input_shape,output_shape,owner_name)
        
        self.conv = [nn.Sequential(BaseN.conv3_2(input_shape[0], 4),nn.Softplus(),
                                                 BaseN.conv3_2(4, 4),BaseN.AdaptiveTanh())]
        x = BaseN.output_shape(self.conv[0],input_shape)
        self.model = nn.Sequential(self.conv[0],
                                   BaseN.Flatten(),
                                   nn.Linear(np.prod(x), 512),
                                   nn.Linear(512,256),nn.Tanh(),
                                   nn.Linear(256,512),
                                   BaseN.EigenLayer(512,self.output_shape[0],bias=False))
        self.compile()
class ConvNetBigAtari(BaseN.BaseNetwork):
    name="ConvNetBigAtari"
    def __init__(self,input_shape,output_shape,owner_name=""):
        super(ConvNetBigAtari,self).__init__(input_shape,output_shape,owner_name)
        
        self.conv = [nn.Sequential(BaseN.conv3_2(input_shape[0], 8),nn.Softplus(),
                                                 BaseN.conv3_2(8, 16),
                                                 BaseN.conv3_2(16, 32))]
        x = BaseN.output_shape(self.conv[0],input_shape)
        self.model = nn.Sequential(self.conv[0],
                                   BaseN.Flatten(),
                                   nn.Linear(np.prod(x), 512),
                                   nn.Linear(512,512),nn.Tanh(),
                                   nn.Linear(512,1024),
                                   BaseN.EigenLayer(1024,self.output_shape[0]))
        self.compile()

class ConvNetBigS(BaseN.BaseNetwork):
    name="ConvNetBigS"
    def __init__(self,input_shape,output_shape,owner_name=""):
        super(ConvNetBigS,self).__init__(input_shape,output_shape,owner_name)
        
        self.conv = [nn.Sequential(BaseN.conv3_2(input_shape[0], 8),nn.Softplus(),
                                                 BaseN.conv3_2(8, 16),
                                                 BaseN.conv3_2(16, 32))]
        x = BaseN.output_shape(self.conv[0],input_shape)
        self.model = nn.Sequential(self.conv[0],
                                   BaseN.Flatten(),
                                   nn.Linear(np.prod(x), 512),
                                   nn.Linear(512,256),nn.Tanh(),
                                   nn.Linear(256,512),
                                   nn.Linear(512,self.output_shape[0]))
        self.compile()
        
class ConvNetMNIST(BaseN.BaseNetwork):
    name = "ConvNetMNIST"
    def __init__(self,input_shape,output_shape,**kwargs):
        super(ConvNetMNIST,self).__init__(**kwargs)
        self.n = output_shape
        self.conv = [BaseN.ResNetBlock(1,32),
                         BaseN.conv3_2(32,64)]
        x = BaseN.output_shape(self.conv[0],input_shape)
        self.model = nn.Sequential(self.conv[0], nn.Softplus(),
                                    BaseN.Flatten(),
                                    nn.Linear(np.prod(x), 512),
                                    nn.Linear(512,256),nn.Tanh(),
                                    BaseN.EigenLayer(256,self.output_shape[0]))
        self.compile()
        
class ConvNetSimple(BaseN.BaseNetwork):
    name="ConvNetSimple"
    def __init__(self,input_shape,output_shape,owner_name=""):
        super(ConvNetSimple,self).__init__(input_shape,output_shape,owner_name)
        
        self.conv = [nn.Sequential(BaseN.conv3_2(input_shape[0], 4),nn.Softplus())]
        x = BaseN.output_shape(self.conv[0],input_shape)
        self.model = nn.Sequential(self.conv[0],
                                   BaseN.Flatten(),
                                   nn.Linear(np.prod(x), 512),
                                   nn.Linear(512,256),nn.Tanh(),
                                   nn.Linear(256,self.output_shape[0]))
        self.compile()

class FCNetSimple(BaseN.BaseNetwork):
    name ="FCNetSimple"
    def __init__(self,input_shape,output_shape,owner_name=""):
        super(FCNetSimple,self).__init__(input_shape,output_shape,owner_name)
        
        x = input_shape
        self.model = nn.Sequential(BaseN.Flatten(),
                                   nn.Linear(np.prod(x), 1024),nn.Softplus(),
                                   nn.Linear(1024,512),
                                   nn.Linear(512,256),nn.Tanh(),
                                   nn.Linear(256,self.output_shape[0]))
        self.compile()

[((2877, 2922), 'base.basenetwork.output_shape', 'BaseN.output_shape', (['self.conv[0]', 'input_shape'], {}), '(self.conv[0], input_shape)\n', (2895, 2922), True, 'import base.basenetwork as BaseN\n'), ((3776, 3821), 'base.basenetwork.output_shape', 'BaseN.output_shape', (['self.conv[0]', 'input_shape'], {}), '(self.conv[0], input_shape)\n', (3794, 3821), True, 'import base.basenetwork as BaseN\n'), ((4553, 4598), 'base.basenetwork.output_shape', 'BaseN.output_shape', (['self.conv[0]', 'input_shape'], {}), '(self.conv[0], input_shape)\n', (4571, 4598), True, 'import base.basenetwork as BaseN\n'), ((5401, 5446), 'base.basenetwork.output_shape', 'BaseN.output_shape', (['self.conv[0]', 'input_shape'], {}), '(self.conv[0], input_shape)\n', (5419, 5446), True, 'import base.basenetwork as BaseN\n'), ((6254, 6299), 'base.basenetwork.output_shape', 'BaseN.output_shape', (['self.conv[0]', 'input_shape'], {}), '(self.conv[0], input_shape)\n', (6272, 6299), True, 'import base.basenetwork as BaseN\n'), ((7388, 7433), 'base.basenetwork.output_shape', 'BaseN.output_shape', (['self.conv[0]', 'input_shape'], {}), '(self.conv[0], input_shape)\n', (7406, 7433), True, 'import base.basenetwork as BaseN\n'), ((8227, 8272), 'base.basenetwork.output_shape', 'BaseN.output_shape', (['self.conv[0]', 'input_shape'], {}), '(self.conv[0], input_shape)\n', (8245, 8272), True, 'import base.basenetwork as BaseN\n'), ((9131, 9176), 'base.basenetwork.output_shape', 'BaseN.output_shape', (['self.conv[0]', 'input_shape'], {}), '(self.conv[0], input_shape)\n', (9149, 9176), True, 'import base.basenetwork as BaseN\n'), ((10015, 10060), 'base.basenetwork.output_shape', 'BaseN.output_shape', (['self.conv[0]', 'input_shape'], {}), '(self.conv[0], input_shape)\n', (10033, 10060), True, 'import base.basenetwork as BaseN\n'), ((10758, 10803), 'base.basenetwork.output_shape', 'BaseN.output_shape', (['self.conv[0]', 'input_shape'], {}), '(self.conv[0], input_shape)\n', (10776, 10803), True, 'import base.basenetwork as BaseN\n'), ((11475, 11520), 'base.basenetwork.output_shape', 'BaseN.output_shape', (['self.conv[0]', 'input_shape'], {}), '(self.conv[0], input_shape)\n', (11493, 11520), True, 'import base.basenetwork as BaseN\n'), ((371, 386), 'base.basenetwork.Flatten', 'BaseN.Flatten', ([], {}), '()\n', (384, 386), True, 'import base.basenetwork as BaseN\n'), ((451, 464), 'torch.nn.Softplus', 'nn.Softplus', ([], {}), '()\n', (462, 464), False, 'from torch import nn\n'), ((501, 521), 'torch.nn.Linear', 'nn.Linear', (['(1024)', '(512)'], {}), '(1024, 512)\n', (510, 521), False, 'from torch import nn\n'), ((521, 530), 'torch.nn.Tanh', 'nn.Tanh', ([], {}), '()\n', (528, 530), False, 'from torch import nn\n'), ((567, 586), 'torch.nn.Linear', 'nn.Linear', (['(512)', '(256)'], {}), '(512, 256)\n', (576, 586), False, 'from torch import nn\n'), ((622, 665), 'base.basenetwork.EigenLayer', 'BaseN.EigenLayer', (['(256)', 'self.output_shape[0]'], {}), '(256, self.output_shape[0])\n', (638, 665), True, 'import base.basenetwork as BaseN\n'), ((967, 982), 'base.basenetwork.Flatten', 'BaseN.Flatten', ([], {}), '()\n', (980, 982), True, 'import base.basenetwork as BaseN\n'), ((1047, 1067), 'base.basenetwork.AdaptiveTanh', 'BaseN.AdaptiveTanh', ([], {}), '()\n', (1065, 1067), True, 'import base.basenetwork as BaseN\n'), ((1104, 1125), 'torch.nn.Linear', 'nn.Linear', (['(1024)', '(1024)'], {}), '(1024, 1024)\n', (1113, 1125), False, 'from torch import nn\n'), ((1125, 1145), 'base.basenetwork.AdaptiveTanh', 'BaseN.AdaptiveTanh', ([], {}), '()\n', (1143, 1145), True, 'import base.basenetwork as BaseN\n'), ((1182, 1202), 'torch.nn.Linear', 'nn.Linear', (['(1024)', '(512)'], {}), '(1024, 512)\n', (1191, 1202), False, 'from torch import nn\n'), ((1202, 1222), 'base.basenetwork.AdaptiveTanh', 'BaseN.AdaptiveTanh', ([], {}), '()\n', (1220, 1222), True, 'import base.basenetwork as BaseN\n'), ((1259, 1302), 'base.basenetwork.EigenLayer', 'BaseN.EigenLayer', (['(512)', 'self.output_shape[0]'], {}), '(512, self.output_shape[0])\n', (1275, 1302), True, 'import base.basenetwork as BaseN\n'), ((1606, 1621), 'base.basenetwork.Flatten', 'BaseN.Flatten', ([], {}), '()\n', (1619, 1621), True, 'import base.basenetwork as BaseN\n'), ((1686, 1695), 'torch.nn.ReLU', 'nn.ReLU', ([], {}), '()\n', (1693, 1695), False, 'from torch import nn\n'), ((1732, 1753), 'torch.nn.Linear', 'nn.Linear', (['(1024)', '(1024)'], {}), '(1024, 1024)\n', (1741, 1753), False, 'from torch import nn\n'), ((1753, 1762), 'torch.nn.ReLU', 'nn.ReLU', ([], {}), '()\n', (1760, 1762), False, 'from torch import nn\n'), ((1799, 1819), 'torch.nn.Linear', 'nn.Linear', (['(1024)', '(512)'], {}), '(1024, 512)\n', (1808, 1819), False, 'from torch import nn\n'), ((1819, 1828), 'torch.nn.ReLU', 'nn.ReLU', ([], {}), '()\n', (1826, 1828), False, 'from torch import nn\n'), ((1865, 1905), 'torch.nn.Linear', 'nn.Linear', (['(512)', '(self.output_shape[0] - 1)'], {}), '(512, self.output_shape[0] - 1)\n', (1874, 1905), False, 'from torch import nn\n'), ((1903, 1912), 'torch.nn.Tanh', 'nn.Tanh', ([], {}), '()\n', (1910, 1912), False, 'from torch import nn\n'), ((1949, 1967), 'base.basenetwork.EigenLayer', 'BaseN.EigenLayer', ([], {}), '()\n', (1965, 1967), True, 'import base.basenetwork as BaseN\n'), ((2205, 2225), 'numpy.prod', 'np.prod', (['input_shape'], {}), '(input_shape)\n', (2212, 2225), True, 'import numpy as np\n'), ((2262, 2277), 'base.basenetwork.Flatten', 'BaseN.Flatten', ([], {}), '()\n', (2275, 2277), True, 'import base.basenetwork as BaseN\n'), ((2314, 2329), 'torch.nn.Linear', 'nn.Linear', (['x', 'x'], {}), '(x, x)\n', (2323, 2329), False, 'from torch import nn\n'), ((2329, 2338), 'torch.nn.Tanh', 'nn.Tanh', ([], {}), '()\n', (2336, 2338), False, 'from torch import nn\n'), ((2375, 2409), 'torch.nn.Linear', 'nn.Linear', (['x', 'self.output_shape[0]'], {}), '(x, self.output_shape[0])\n', (2384, 2409), False, 'from torch import nn\n'), ((3006, 3021), 'base.basenetwork.Flatten', 'BaseN.Flatten', ([], {}), '()\n', (3019, 3021), True, 'import base.basenetwork as BaseN\n'), ((3085, 3105), 'base.basenetwork.AdaptiveTanh', 'BaseN.AdaptiveTanh', ([], {}), '()\n', (3103, 3105), True, 'import base.basenetwork as BaseN\n'), ((3142, 3161), 'torch.nn.Linear', 'nn.Linear', (['(512)', '(256)'], {}), '(512, 256)\n', (3151, 3161), False, 'from torch import nn\n'), ((3197, 3252), 'base.basenetwork.EigenLayer', 'BaseN.EigenLayer', (['(256)', 'self.output_shape[0]'], {'bias': '(False)'}), '(256, self.output_shape[0], bias=False)\n', (3213, 3252), True, 'import base.basenetwork as BaseN\n'), ((3905, 3920), 'base.basenetwork.Flatten', 'BaseN.Flatten', ([], {}), '()\n', (3918, 3920), True, 'import base.basenetwork as BaseN\n'), ((3984, 4004), 'base.basenetwork.AdaptiveTanh', 'BaseN.AdaptiveTanh', ([], {}), '()\n', (4002, 4004), True, 'import base.basenetwork as BaseN\n'), ((4041, 4060), 'torch.nn.Linear', 'nn.Linear', (['(512)', '(256)'], {}), '(512, 256)\n', (4050, 4060), False, 'from torch import nn\n'), ((4096, 4151), 'base.basenetwork.EigenLayer', 'BaseN.EigenLayer', (['(256)', 'self.output_shape[0]'], {'bias': '(False)'}), '(256, self.output_shape[0], bias=False)\n', (4112, 4151), True, 'import base.basenetwork as BaseN\n'), ((4682, 4697), 'base.basenetwork.Flatten', 'BaseN.Flatten', ([], {}), '()\n', (4695, 4697), True, 'import base.basenetwork as BaseN\n'), ((4797, 4816), 'torch.nn.Linear', 'nn.Linear', (['(512)', '(512)'], {}), '(512, 512)\n', (4806, 4816), False, 'from torch import nn\n'), ((4816, 4836), 'base.basenetwork.AdaptiveTanh', 'BaseN.AdaptiveTanh', ([], {}), '()\n', (4834, 4836), True, 'import base.basenetwork as BaseN\n'), ((4873, 4892), 'torch.nn.Linear', 'nn.Linear', (['(512)', '(256)'], {}), '(512, 256)\n', (4882, 4892), False, 'from torch import nn\n'), ((4928, 4983), 'base.basenetwork.EigenLayer', 'BaseN.EigenLayer', (['(256)', 'self.output_shape[0]'], {'bias': '(False)'}), '(256, self.output_shape[0], bias=False)\n', (4944, 4983), True, 'import base.basenetwork as BaseN\n'), ((5530, 5545), 'base.basenetwork.Flatten', 'BaseN.Flatten', ([], {}), '()\n', (5543, 5545), True, 'import base.basenetwork as BaseN\n'), ((5645, 5665), 'torch.nn.Linear', 'nn.Linear', (['(512)', '(1024)'], {}), '(512, 1024)\n', (5654, 5665), False, 'from torch import nn\n'), ((5665, 5685), 'base.basenetwork.AdaptiveTanh', 'BaseN.AdaptiveTanh', ([], {}), '()\n', (5683, 5685), True, 'import base.basenetwork as BaseN\n'), ((5722, 5742), 'torch.nn.Linear', 'nn.Linear', (['(1024)', '(256)'], {}), '(1024, 256)\n', (5731, 5742), False, 'from torch import nn\n'), ((5778, 5833), 'base.basenetwork.EigenLayer', 'BaseN.EigenLayer', (['(256)', 'self.output_shape[0]'], {'bias': '(False)'}), '(256, self.output_shape[0], bias=False)\n', (5794, 5833), True, 'import base.basenetwork as BaseN\n'), ((6383, 6398), 'base.basenetwork.Flatten', 'BaseN.Flatten', ([], {}), '()\n', (6396, 6398), True, 'import base.basenetwork as BaseN\n'), ((6498, 6517), 'torch.nn.Linear', 'nn.Linear', (['(512)', '(256)'], {}), '(512, 256)\n', (6507, 6517), False, 'from torch import nn\n'), ((6517, 6526), 'torch.nn.Tanh', 'nn.Tanh', ([], {}), '()\n', (6524, 6526), False, 'from torch import nn\n'), ((6563, 6582), 'torch.nn.Linear', 'nn.Linear', (['(256)', '(512)'], {}), '(256, 512)\n', (6572, 6582), False, 'from torch import nn\n'), ((6618, 6638), 'torch.nn.Linear', 'nn.Linear', (['(512)', '(1024)'], {}), '(512, 1024)\n', (6627, 6638), False, 'from torch import nn\n'), ((6638, 6647), 'torch.nn.Tanh', 'nn.Tanh', ([], {}), '()\n', (6645, 6647), False, 'from torch import nn\n'), ((6684, 6704), 'torch.nn.Linear', 'nn.Linear', (['(1024)', '(512)'], {}), '(1024, 512)\n', (6693, 6704), False, 'from torch import nn\n'), ((6740, 6759), 'torch.nn.Linear', 'nn.Linear', (['(512)', '(256)'], {}), '(512, 256)\n', (6749, 6759), False, 'from torch import nn\n'), ((6759, 6768), 'torch.nn.Tanh', 'nn.Tanh', ([], {}), '()\n', (6766, 6768), False, 'from torch import nn\n'), ((6805, 6824), 'torch.nn.Linear', 'nn.Linear', (['(256)', '(256)'], {}), '(256, 256)\n', (6814, 6824), False, 'from torch import nn\n'), ((6860, 6903), 'base.basenetwork.EigenLayer', 'BaseN.EigenLayer', (['(256)', 'self.output_shape[0]'], {}), '(256, self.output_shape[0])\n', (6876, 6903), True, 'import base.basenetwork as BaseN\n'), ((7517, 7532), 'base.basenetwork.Flatten', 'BaseN.Flatten', ([], {}), '()\n', (7530, 7532), True, 'import base.basenetwork as BaseN\n'), ((7632, 7651), 'torch.nn.Linear', 'nn.Linear', (['(512)', '(256)'], {}), '(512, 256)\n', (7641, 7651), False, 'from torch import nn\n'), ((7651, 7660), 'torch.nn.Tanh', 'nn.Tanh', ([], {}), '()\n', (7658, 7660), False, 'from torch import nn\n'), ((7697, 7716), 'torch.nn.Linear', 'nn.Linear', (['(256)', '(512)'], {}), '(256, 512)\n', (7706, 7716), False, 'from torch import nn\n'), ((7752, 7795), 'base.basenetwork.EigenLayer', 'BaseN.EigenLayer', (['(512)', 'self.output_shape[0]'], {}), '(512, self.output_shape[0])\n', (7768, 7795), True, 'import base.basenetwork as BaseN\n'), ((8356, 8371), 'base.basenetwork.Flatten', 'BaseN.Flatten', ([], {}), '()\n', (8369, 8371), True, 'import base.basenetwork as BaseN\n'), ((8471, 8490), 'torch.nn.Linear', 'nn.Linear', (['(512)', '(256)'], {}), '(512, 256)\n', (8480, 8490), False, 'from torch import nn\n'), ((8490, 8499), 'torch.nn.Tanh', 'nn.Tanh', ([], {}), '()\n', (8497, 8499), False, 'from torch import nn\n'), ((8536, 8555), 'torch.nn.Linear', 'nn.Linear', (['(256)', '(512)'], {}), '(256, 512)\n', (8545, 8555), False, 'from torch import nn\n'), ((8591, 8646), 'base.basenetwork.EigenLayer', 'BaseN.EigenLayer', (['(512)', 'self.output_shape[0]'], {'bias': '(False)'}), '(512, self.output_shape[0], bias=False)\n', (8607, 8646), True, 'import base.basenetwork as BaseN\n'), ((9260, 9275), 'base.basenetwork.Flatten', 'BaseN.Flatten', ([], {}), '()\n', (9273, 9275), True, 'import base.basenetwork as BaseN\n'), ((9375, 9394), 'torch.nn.Linear', 'nn.Linear', (['(512)', '(512)'], {}), '(512, 512)\n', (9384, 9394), False, 'from torch import nn\n'), ((9394, 9403), 'torch.nn.Tanh', 'nn.Tanh', ([], {}), '()\n', (9401, 9403), False, 'from torch import nn\n'), ((9440, 9460), 'torch.nn.Linear', 'nn.Linear', (['(512)', '(1024)'], {}), '(512, 1024)\n', (9449, 9460), False, 'from torch import nn\n'), ((9496, 9540), 'base.basenetwork.EigenLayer', 'BaseN.EigenLayer', (['(1024)', 'self.output_shape[0]'], {}), '(1024, self.output_shape[0])\n', (9512, 9540), True, 'import base.basenetwork as BaseN\n'), ((10144, 10159), 'base.basenetwork.Flatten', 'BaseN.Flatten', ([], {}), '()\n', (10157, 10159), True, 'import base.basenetwork as BaseN\n'), ((10259, 10278), 'torch.nn.Linear', 'nn.Linear', (['(512)', '(256)'], {}), '(512, 256)\n', (10268, 10278), False, 'from torch import nn\n'), ((10278, 10287), 'torch.nn.Tanh', 'nn.Tanh', ([], {}), '()\n', (10285, 10287), False, 'from torch import nn\n'), ((10324, 10343), 'torch.nn.Linear', 'nn.Linear', (['(256)', '(512)'], {}), '(256, 512)\n', (10333, 10343), False, 'from torch import nn\n'), ((10379, 10415), 'torch.nn.Linear', 'nn.Linear', (['(512)', 'self.output_shape[0]'], {}), '(512, self.output_shape[0])\n', (10388, 10415), False, 'from torch import nn\n'), ((10674, 10698), 'base.basenetwork.ResNetBlock', 'BaseN.ResNetBlock', (['(1)', '(32)'], {}), '(1, 32)\n', (10691, 10698), True, 'import base.basenetwork as BaseN\n'), ((10724, 10745), 'base.basenetwork.conv3_2', 'BaseN.conv3_2', (['(32)', '(64)'], {}), '(32, 64)\n', (10737, 10745), True, 'import base.basenetwork as BaseN\n'), ((10852, 10865), 'torch.nn.Softplus', 'nn.Softplus', ([], {}), '()\n', (10863, 10865), False, 'from torch import nn\n'), ((10903, 10918), 'base.basenetwork.Flatten', 'BaseN.Flatten', ([], {}), '()\n', (10916, 10918), True, 'import base.basenetwork as BaseN\n'), ((11020, 11039), 'torch.nn.Linear', 'nn.Linear', (['(512)', '(256)'], {}), '(512, 256)\n', (11029, 11039), False, 'from torch import nn\n'), ((11039, 11048), 'torch.nn.Tanh', 'nn.Tanh', ([], {}), '()\n', (11046, 11048), False, 'from torch import nn\n'), ((11086, 11129), 'base.basenetwork.EigenLayer', 'BaseN.EigenLayer', (['(256)', 'self.output_shape[0]'], {}), '(256, self.output_shape[0])\n', (11102, 11129), True, 'import base.basenetwork as BaseN\n'), ((11604, 11619), 'base.basenetwork.Flatten', 'BaseN.Flatten', ([], {}), '()\n', (11617, 11619), True, 'import base.basenetwork as BaseN\n'), ((11719, 11738), 'torch.nn.Linear', 'nn.Linear', (['(512)', '(256)'], {}), '(512, 256)\n', (11728, 11738), False, 'from torch import nn\n'), ((11738, 11747), 'torch.nn.Tanh', 'nn.Tanh', ([], {}), '()\n', (11745, 11747), False, 'from torch import nn\n'), ((11784, 11820), 'torch.nn.Linear', 'nn.Linear', (['(256)', 'self.output_shape[0]'], {}), '(256, self.output_shape[0])\n', (11793, 11820), False, 'from torch import nn\n'), ((12116, 12131), 'base.basenetwork.Flatten', 'BaseN.Flatten', ([], {}), '()\n', (12129, 12131), True, 'import base.basenetwork as BaseN\n'), ((12196, 12209), 'torch.nn.Softplus', 'nn.Softplus', ([], {}), '()\n', (12207, 12209), False, 'from torch import nn\n'), ((12246, 12266), 'torch.nn.Linear', 'nn.Linear', (['(1024)', '(512)'], {}), '(1024, 512)\n', (12255, 12266), False, 'from torch import nn\n'), ((12302, 12321), 'torch.nn.Linear', 'nn.Linear', (['(512)', '(256)'], {}), '(512, 256)\n', (12311, 12321), False, 'from torch import nn\n'), ((12321, 12330), 'torch.nn.Tanh', 'nn.Tanh', ([], {}), '()\n', (12328, 12330), False, 'from torch import nn\n'), ((12367, 12403), 'torch.nn.Linear', 'nn.Linear', (['(256)', 'self.output_shape[0]'], {}), '(256, self.output_shape[0])\n', (12376, 12403), False, 'from torch import nn\n'), ((433, 443), 'numpy.prod', 'np.prod', (['x'], {}), '(x)\n', (440, 443), True, 'import numpy as np\n'), ((1029, 1039), 'numpy.prod', 'np.prod', (['x'], {}), '(x)\n', (1036, 1039), True, 'import numpy as np\n'), ((1668, 1678), 'numpy.prod', 'np.prod', (['x'], {}), '(x)\n', (1675, 1678), True, 'import numpy as np\n'), ((2669, 2701), 'base.basenetwork.conv3_2', 'BaseN.conv3_2', (['input_shape[0]', '(8)'], {}), '(input_shape[0], 8)\n', (2682, 2701), True, 'import base.basenetwork as BaseN\n'), ((2702, 2711), 'torch.nn.ReLU', 'nn.ReLU', ([], {}), '()\n', (2709, 2711), False, 'from torch import nn\n'), ((2762, 2782), 'base.basenetwork.conv3_2', 'BaseN.conv3_2', (['(8)', '(16)'], {}), '(8, 16)\n', (2775, 2782), True, 'import base.basenetwork as BaseN\n'), ((2783, 2792), 'torch.nn.ReLU', 'nn.ReLU', ([], {}), '()\n', (2790, 2792), False, 'from torch import nn\n'), ((2843, 2862), 'base.basenetwork.conv3_2', 'BaseN.conv3_2', (['(8)', '(8)'], {}), '(8, 8)\n', (2856, 2862), True, 'import base.basenetwork as BaseN\n'), ((3068, 3078), 'numpy.prod', 'np.prod', (['x'], {}), '(x)\n', (3075, 3078), True, 'import numpy as np\n'), ((3521, 3553), 'base.basenetwork.conv3_2', 'BaseN.conv3_2', (['input_shape[0]', '(8)'], {}), '(input_shape[0], 8)\n', (3534, 3553), True, 'import base.basenetwork as BaseN\n'), ((3554, 3567), 'torch.nn.Softplus', 'nn.Softplus', ([], {}), '()\n', (3565, 3567), False, 'from torch import nn\n'), ((3604, 3624), 'base.basenetwork.conv3_2', 'BaseN.conv3_2', (['(8)', '(12)'], {}), '(8, 12)\n', (3617, 3624), True, 'import base.basenetwork as BaseN\n'), ((3625, 3645), 'base.basenetwork.AdaptiveTanh', 'BaseN.AdaptiveTanh', ([], {}), '()\n', (3643, 3645), True, 'import base.basenetwork as BaseN\n'), ((3682, 3703), 'base.basenetwork.conv3_2', 'BaseN.conv3_2', (['(12)', '(16)'], {}), '(12, 16)\n', (3695, 3703), True, 'import base.basenetwork as BaseN\n'), ((3740, 3761), 'base.basenetwork.conv3_2', 'BaseN.conv3_2', (['(16)', '(20)'], {}), '(16, 20)\n', (3753, 3761), True, 'import base.basenetwork as BaseN\n'), ((3967, 3977), 'numpy.prod', 'np.prod', (['x'], {}), '(x)\n', (3974, 3977), True, 'import numpy as np\n'), ((4415, 4447), 'base.basenetwork.conv3_2', 'BaseN.conv3_2', (['input_shape[0]', '(4)'], {}), '(input_shape[0], 4)\n', (4428, 4447), True, 'import base.basenetwork as BaseN\n'), ((4448, 4461), 'torch.nn.Softplus', 'nn.Softplus', ([], {}), '()\n', (4459, 4461), False, 'from torch import nn\n'), ((4498, 4517), 'base.basenetwork.conv3_2', 'BaseN.conv3_2', (['(4)', '(8)'], {}), '(4, 8)\n', (4511, 4517), True, 'import base.basenetwork as BaseN\n'), ((4518, 4538), 'base.basenetwork.AdaptiveTanh', 'BaseN.AdaptiveTanh', ([], {}), '()\n', (4536, 4538), True, 'import base.basenetwork as BaseN\n'), ((4744, 4754), 'numpy.prod', 'np.prod', (['x'], {}), '(x)\n', (4751, 4754), True, 'import numpy as np\n'), ((5267, 5299), 'base.basenetwork.conv3_2', 'BaseN.conv3_2', (['input_shape[0]', '(4)'], {}), '(input_shape[0], 4)\n', (5280, 5299), True, 'import base.basenetwork as BaseN\n'), ((5300, 5309), 'torch.nn.ReLU', 'nn.ReLU', ([], {}), '()\n', (5307, 5309), False, 'from torch import nn\n'), ((5346, 5365), 'base.basenetwork.conv3_2', 'BaseN.conv3_2', (['(4)', '(4)'], {}), '(4, 4)\n', (5359, 5365), True, 'import base.basenetwork as BaseN\n'), ((5366, 5386), 'base.basenetwork.AdaptiveTanh', 'BaseN.AdaptiveTanh', ([], {}), '()\n', (5384, 5386), True, 'import base.basenetwork as BaseN\n'), ((5592, 5602), 'numpy.prod', 'np.prod', (['x'], {}), '(x)\n', (5599, 5602), True, 'import numpy as np\n'), ((6102, 6134), 'base.basenetwork.conv3_2', 'BaseN.conv3_2', (['input_shape[0]', '(3)'], {}), '(input_shape[0], 3)\n', (6115, 6134), True, 'import base.basenetwork as BaseN\n'), ((6135, 6148), 'torch.nn.Softplus', 'nn.Softplus', ([], {}), '()\n', (6146, 6148), False, 'from torch import nn\n'), ((6199, 6218), 'base.basenetwork.conv3_2', 'BaseN.conv3_2', (['(3)', '(6)'], {}), '(3, 6)\n', (6212, 6218), True, 'import base.basenetwork as BaseN\n'), ((6219, 6239), 'base.basenetwork.conv3_2', 'BaseN.conv3_2', (['(6)', '(12)'], {}), '(6, 12)\n', (6232, 6239), True, 'import base.basenetwork as BaseN\n'), ((6445, 6455), 'numpy.prod', 'np.prod', (['x'], {}), '(x)\n', (6452, 6455), True, 'import numpy as np\n'), ((7170, 7202), 'base.basenetwork.conv3_2', 'BaseN.conv3_2', (['input_shape[0]', '(8)'], {}), '(input_shape[0], 8)\n', (7183, 7202), True, 'import base.basenetwork as BaseN\n'), ((7203, 7216), 'torch.nn.Softplus', 'nn.Softplus', ([], {}), '()\n', (7214, 7216), False, 'from torch import nn\n'), ((7267, 7287), 'base.basenetwork.conv3_2', 'BaseN.conv3_2', (['(8)', '(16)'], {}), '(8, 16)\n', (7280, 7287), True, 'import base.basenetwork as BaseN\n'), ((7288, 7301), 'torch.nn.Softplus', 'nn.Softplus', ([], {}), '()\n', (7299, 7301), False, 'from torch import nn\n'), ((7352, 7373), 'base.basenetwork.conv3_2', 'BaseN.conv3_2', (['(16)', '(32)'], {}), '(16, 32)\n', (7365, 7373), True, 'import base.basenetwork as BaseN\n'), ((7579, 7589), 'numpy.prod', 'np.prod', (['x'], {}), '(x)\n', (7586, 7589), True, 'import numpy as np\n'), ((8075, 8107), 'base.basenetwork.conv3_2', 'BaseN.conv3_2', (['input_shape[0]', '(4)'], {}), '(input_shape[0], 4)\n', (8088, 8107), True, 'import base.basenetwork as BaseN\n'), ((8108, 8121), 'torch.nn.Softplus', 'nn.Softplus', ([], {}), '()\n', (8119, 8121), False, 'from torch import nn\n'), ((8172, 8191), 'base.basenetwork.conv3_2', 'BaseN.conv3_2', (['(4)', '(4)'], {}), '(4, 4)\n', (8185, 8191), True, 'import base.basenetwork as BaseN\n'), ((8192, 8212), 'base.basenetwork.AdaptiveTanh', 'BaseN.AdaptiveTanh', ([], {}), '()\n', (8210, 8212), True, 'import base.basenetwork as BaseN\n'), ((8418, 8428), 'numpy.prod', 'np.prod', (['x'], {}), '(x)\n', (8425, 8428), True, 'import numpy as np\n'), ((8927, 8959), 'base.basenetwork.conv3_2', 'BaseN.conv3_2', (['input_shape[0]', '(8)'], {}), '(input_shape[0], 8)\n', (8940, 8959), True, 'import base.basenetwork as BaseN\n'), ((8960, 8973), 'torch.nn.Softplus', 'nn.Softplus', ([], {}), '()\n', (8971, 8973), False, 'from torch import nn\n'), ((9024, 9044), 'base.basenetwork.conv3_2', 'BaseN.conv3_2', (['(8)', '(16)'], {}), '(8, 16)\n', (9037, 9044), True, 'import base.basenetwork as BaseN\n'), ((9095, 9116), 'base.basenetwork.conv3_2', 'BaseN.conv3_2', (['(16)', '(32)'], {}), '(16, 32)\n', (9108, 9116), True, 'import base.basenetwork as BaseN\n'), ((9322, 9332), 'numpy.prod', 'np.prod', (['x'], {}), '(x)\n', (9329, 9332), True, 'import numpy as np\n'), ((9811, 9843), 'base.basenetwork.conv3_2', 'BaseN.conv3_2', (['input_shape[0]', '(8)'], {}), '(input_shape[0], 8)\n', (9824, 9843), True, 'import base.basenetwork as BaseN\n'), ((9844, 9857), 'torch.nn.Softplus', 'nn.Softplus', ([], {}), '()\n', (9855, 9857), False, 'from torch import nn\n'), ((9908, 9928), 'base.basenetwork.conv3_2', 'BaseN.conv3_2', (['(8)', '(16)'], {}), '(8, 16)\n', (9921, 9928), True, 'import base.basenetwork as BaseN\n'), ((9979, 10000), 'base.basenetwork.conv3_2', 'BaseN.conv3_2', (['(16)', '(32)'], {}), '(16, 32)\n', (9992, 10000), True, 'import base.basenetwork as BaseN\n'), ((10206, 10216), 'numpy.prod', 'np.prod', (['x'], {}), '(x)\n', (10213, 10216), True, 'import numpy as np\n'), ((10966, 10976), 'numpy.prod', 'np.prod', (['x'], {}), '(x)\n', (10973, 10976), True, 'import numpy as np\n'), ((11414, 11446), 'base.basenetwork.conv3_2', 'BaseN.conv3_2', (['input_shape[0]', '(4)'], {}), '(input_shape[0], 4)\n', (11427, 11446), True, 'import base.basenetwork as BaseN\n'), ((11447, 11460), 'torch.nn.Softplus', 'nn.Softplus', ([], {}), '()\n', (11458, 11460), False, 'from torch import nn\n'), ((11666, 11676), 'numpy.prod', 'np.prod', (['x'], {}), '(x)\n', (11673, 11676), True, 'import numpy as np\n'), ((12178, 12188), 'numpy.prod', 'np.prod', (['x'], {}), '(x)\n', (12185, 12188), True, 'import numpy as np\n')]