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()

[((64, 12, 64, 56), 'base.basenetwork.output_shape', 'BaseN.output_shape', ({(64, 31, 64, 43): 'self.conv[0]', (64, 44, 64, 55): 'input_shape'}, {}), '(self.conv[0], input_shape)', True, 'import base.basenetwork as BaseN\n'), ((80, 12, 80, 56), 'base.basenetwork.output_shape', 'BaseN.output_shape', ({(80, 31, 80, 43): 'self.conv[0]', (80, 44, 80, 55): 'input_shape'}, {}), '(self.conv[0], input_shape)', True, 'import base.basenetwork as BaseN\n'), ((95, 12, 95, 56), 'base.basenetwork.output_shape', 'BaseN.output_shape', ({(95, 31, 95, 43): 'self.conv[0]', (95, 44, 95, 55): 'input_shape'}, {}), '(self.conv[0], input_shape)', True, 'import base.basenetwork as BaseN\n'), ((111, 12, 111, 56), 'base.basenetwork.output_shape', 'BaseN.output_shape', ({(111, 31, 111, 43): 'self.conv[0]', (111, 44, 111, 55): 'input_shape'}, {}), '(self.conv[0], input_shape)', True, 'import base.basenetwork as BaseN\n'), ((127, 12, 127, 56), 'base.basenetwork.output_shape', 'BaseN.output_shape', ({(127, 31, 127, 43): 'self.conv[0]', (127, 44, 127, 55): 'input_shape'}, {}), '(self.conv[0], input_shape)', True, 'import base.basenetwork as BaseN\n'), ((147, 12, 147, 56), 'base.basenetwork.output_shape', 'BaseN.output_shape', ({(147, 31, 147, 43): 'self.conv[0]', (147, 44, 147, 55): 'input_shape'}, {}), '(self.conv[0], input_shape)', True, 'import base.basenetwork as BaseN\n'), ((163, 12, 163, 56), 'base.basenetwork.output_shape', 'BaseN.output_shape', ({(163, 31, 163, 43): 'self.conv[0]', (163, 44, 163, 55): 'input_shape'}, {}), '(self.conv[0], input_shape)', True, 'import base.basenetwork as BaseN\n'), ((179, 12, 179, 56), 'base.basenetwork.output_shape', 'BaseN.output_shape', ({(179, 31, 179, 43): 'self.conv[0]', (179, 44, 179, 55): 'input_shape'}, {}), '(self.conv[0], input_shape)', True, 'import base.basenetwork as BaseN\n'), ((196, 12, 196, 56), 'base.basenetwork.output_shape', 'BaseN.output_shape', ({(196, 31, 196, 43): 'self.conv[0]', (196, 44, 196, 55): 'input_shape'}, {}), '(self.conv[0], input_shape)', True, 'import base.basenetwork as BaseN\n'), ((212, 12, 212, 56), 'base.basenetwork.output_shape', 'BaseN.output_shape', ({(212, 31, 212, 43): 'self.conv[0]', (212, 44, 212, 55): 'input_shape'}, {}), '(self.conv[0], input_shape)', True, 'import base.basenetwork as BaseN\n'), ((226, 12, 226, 56), 'base.basenetwork.output_shape', 'BaseN.output_shape', ({(226, 31, 226, 43): 'self.conv[0]', (226, 44, 226, 55): 'input_shape'}, {}), '(self.conv[0], input_shape)', True, 'import base.basenetwork as BaseN\n'), ((12, 35, 12, 50), 'base.basenetwork.Flatten', 'BaseN.Flatten', ({}, {}), '()', True, 'import base.basenetwork as BaseN\n'), ((13, 63, 13, 76), 'torch.nn.Softplus', 'nn.Softplus', ({}, {}), '()', False, 'from torch import nn\n'), ((14, 35, 14, 54), 'torch.nn.Linear', 'nn.Linear', ({(14, 45, 14, 49): '1024', (14, 50, 14, 53): '512'}, {}), '(1024, 512)', False, 'from torch import nn\n'), ((14, 55, 14, 64), 'torch.nn.Tanh', 'nn.Tanh', ({}, {}), '()', False, 'from torch import nn\n'), ((15, 35, 15, 53), 'torch.nn.Linear', 'nn.Linear', ({(15, 45, 15, 48): '512', (15, 49, 15, 52): '256'}, {}), '(512, 256)', False, 'from torch import nn\n'), ((16, 35, 16, 77), 'base.basenetwork.EigenLayer', 'BaseN.EigenLayer', ({(16, 52, 16, 55): '256', (16, 56, 16, 76): 'self.output_shape[0]'}, {}), '(256, self.output_shape[0])', True, 'import base.basenetwork as BaseN\n'), ((25, 35, 25, 50), 'base.basenetwork.Flatten', 'BaseN.Flatten', ({}, {}), '()', True, 'import base.basenetwork as BaseN\n'), ((26, 63, 26, 83), 'base.basenetwork.AdaptiveTanh', 'BaseN.AdaptiveTanh', ({}, {}), '()', True, 'import base.basenetwork as BaseN\n'), ((27, 35, 27, 55), 'torch.nn.Linear', 'nn.Linear', ({(27, 45, 27, 49): '1024', (27, 50, 27, 54): '1024'}, {}), '(1024, 1024)', False, 'from torch import nn\n'), ((27, 56, 27, 76), 'base.basenetwork.AdaptiveTanh', 'BaseN.AdaptiveTanh', ({}, {}), '()', True, 'import base.basenetwork as BaseN\n'), ((28, 35, 28, 54), 'torch.nn.Linear', 'nn.Linear', ({(28, 45, 28, 49): '1024', (28, 50, 28, 53): '512'}, {}), '(1024, 512)', False, 'from torch import nn\n'), ((28, 55, 28, 75), 'base.basenetwork.AdaptiveTanh', 'BaseN.AdaptiveTanh', ({}, {}), '()', True, 'import base.basenetwork as BaseN\n'), ((29, 35, 29, 77), 'base.basenetwork.EigenLayer', 'BaseN.EigenLayer', ({(29, 52, 29, 55): '512', (29, 56, 29, 76): 'self.output_shape[0]'}, {}), '(512, self.output_shape[0])', True, 'import base.basenetwork as BaseN\n'), ((37, 35, 37, 50), 'base.basenetwork.Flatten', 'BaseN.Flatten', ({}, {}), '()', True, 'import base.basenetwork as BaseN\n'), ((38, 63, 38, 72), 'torch.nn.ReLU', 'nn.ReLU', ({}, {}), '()', False, 'from torch import nn\n'), ((39, 35, 39, 55), 'torch.nn.Linear', 'nn.Linear', ({(39, 45, 39, 49): '1024', (39, 50, 39, 54): '1024'}, {}), '(1024, 1024)', False, 'from torch import nn\n'), ((39, 56, 39, 65), 'torch.nn.ReLU', 'nn.ReLU', ({}, {}), '()', False, 'from torch import nn\n'), ((40, 35, 40, 54), 'torch.nn.Linear', 'nn.Linear', ({(40, 45, 40, 49): '1024', (40, 50, 40, 53): '512'}, {}), '(1024, 512)', False, 'from torch import nn\n'), ((40, 55, 40, 64), 'torch.nn.ReLU', 'nn.ReLU', ({}, {}), '()', False, 'from torch import nn\n'), ((41, 35, 41, 72), 'torch.nn.Linear', 'nn.Linear', ({(41, 45, 41, 48): '512', (41, 49, 41, 71): 'self.output_shape[0] - 1'}, {}), '(512, self.output_shape[0] - 1)', False, 'from torch import nn\n'), ((41, 73, 41, 82), 'torch.nn.Tanh', 'nn.Tanh', ({}, {}), '()', False, 'from torch import nn\n'), ((42, 35, 42, 53), 'base.basenetwork.EigenLayer', 'BaseN.EigenLayer', ({}, {}), '()', True, 'import base.basenetwork as BaseN\n'), ((49, 16, 49, 36), 'numpy.prod', 'np.prod', ({(49, 24, 49, 35): 'input_shape'}, {}), '(input_shape)', True, 'import numpy as np\n'), ((50, 35, 50, 50), 'base.basenetwork.Flatten', 'BaseN.Flatten', ({}, {}), '()', True, 'import base.basenetwork as BaseN\n'), ((51, 35, 51, 49), 'torch.nn.Linear', 'nn.Linear', ({(51, 45, 51, 46): 'x', (51, 47, 51, 48): 'x'}, {}), '(x, x)', False, 'from torch import nn\n'), ((51, 50, 51, 59), 'torch.nn.Tanh', 'nn.Tanh', ({}, {}), '()', False, 'from torch import nn\n'), ((52, 35, 52, 68), 'torch.nn.Linear', 'nn.Linear', ({(52, 45, 52, 46): 'x', (52, 47, 52, 67): 'self.output_shape[0]'}, {}), '(x, self.output_shape[0])', False, 'from torch import nn\n'), ((66, 35, 66, 50), 'base.basenetwork.Flatten', 'BaseN.Flatten', ({}, {}), '()', True, 'import base.basenetwork as BaseN\n'), ((67, 62, 67, 82), 'base.basenetwork.AdaptiveTanh', 'BaseN.AdaptiveTanh', ({}, {}), '()', True, 'import base.basenetwork as BaseN\n'), ((68, 35, 68, 53), 'torch.nn.Linear', 'nn.Linear', ({(68, 45, 68, 48): '512', (68, 49, 68, 52): '256'}, {}), '(512, 256)', False, 'from torch import nn\n'), ((69, 35, 69, 88), 'base.basenetwork.EigenLayer', 'BaseN.EigenLayer', (), '', True, 'import base.basenetwork as BaseN\n'), ((82, 35, 82, 50), 'base.basenetwork.Flatten', 'BaseN.Flatten', ({}, {}), '()', True, 'import base.basenetwork as BaseN\n'), ((83, 62, 83, 82), 'base.basenetwork.AdaptiveTanh', 'BaseN.AdaptiveTanh', ({}, {}), '()', True, 'import base.basenetwork as BaseN\n'), ((84, 35, 84, 53), 'torch.nn.Linear', 'nn.Linear', ({(84, 45, 84, 48): '512', (84, 49, 84, 52): '256'}, {}), '(512, 256)', False, 'from torch import nn\n'), ((85, 35, 85, 88), 'base.basenetwork.EigenLayer', 'BaseN.EigenLayer', (), '', True, 'import base.basenetwork as BaseN\n'), ((97, 35, 97, 50), 'base.basenetwork.Flatten', 'BaseN.Flatten', ({}, {}), '()', True, 'import base.basenetwork as BaseN\n'), ((99, 35, 99, 53), 'torch.nn.Linear', 'nn.Linear', ({(99, 45, 99, 48): '512', (99, 49, 99, 52): '512'}, {}), '(512, 512)', False, 'from torch import nn\n'), ((99, 54, 99, 74), 'base.basenetwork.AdaptiveTanh', 'BaseN.AdaptiveTanh', ({}, {}), '()', True, 'import base.basenetwork as BaseN\n'), ((100, 35, 100, 53), 'torch.nn.Linear', 'nn.Linear', ({(100, 45, 100, 48): '512', (100, 49, 100, 52): '256'}, {}), '(512, 256)', False, 'from torch import nn\n'), ((101, 35, 101, 88), 'base.basenetwork.EigenLayer', 'BaseN.EigenLayer', (), '', True, 'import base.basenetwork as BaseN\n'), ((113, 35, 113, 50), 'base.basenetwork.Flatten', 'BaseN.Flatten', ({}, {}), '()', True, 'import base.basenetwork as BaseN\n'), ((115, 35, 115, 54), 'torch.nn.Linear', 'nn.Linear', ({(115, 45, 115, 48): '512', (115, 49, 115, 53): '1024'}, {}), '(512, 1024)', False, 'from torch import nn\n'), ((115, 55, 115, 75), 'base.basenetwork.AdaptiveTanh', 'BaseN.AdaptiveTanh', ({}, {}), '()', True, 'import base.basenetwork as BaseN\n'), ((116, 35, 116, 54), 'torch.nn.Linear', 'nn.Linear', ({(116, 45, 116, 49): '1024', (116, 50, 116, 53): '256'}, {}), '(1024, 256)', False, 'from torch import nn\n'), ((117, 35, 117, 88), 'base.basenetwork.EigenLayer', 'BaseN.EigenLayer', (), '', True, 'import base.basenetwork as BaseN\n'), ((129, 35, 129, 50), 'base.basenetwork.Flatten', 'BaseN.Flatten', ({}, {}), '()', True, 'import base.basenetwork as BaseN\n'), ((131, 35, 131, 53), 'torch.nn.Linear', 'nn.Linear', ({(131, 45, 131, 48): '512', (131, 49, 131, 52): '256'}, {}), '(512, 256)', False, 'from torch import nn\n'), ((131, 54, 131, 63), 'torch.nn.Tanh', 'nn.Tanh', ({}, {}), '()', False, 'from torch import nn\n'), ((132, 35, 132, 53), 'torch.nn.Linear', 'nn.Linear', ({(132, 45, 132, 48): '256', (132, 49, 132, 52): '512'}, {}), '(256, 512)', False, 'from torch import nn\n'), ((133, 35, 133, 54), 'torch.nn.Linear', 'nn.Linear', ({(133, 45, 133, 48): '512', (133, 49, 133, 53): '1024'}, {}), '(512, 1024)', False, 'from torch import nn\n'), ((133, 55, 133, 64), 'torch.nn.Tanh', 'nn.Tanh', ({}, {}), '()', False, 'from torch import nn\n'), ((134, 35, 134, 54), 'torch.nn.Linear', 'nn.Linear', ({(134, 45, 134, 49): '1024', (134, 50, 134, 53): '512'}, {}), '(1024, 512)', False, 'from torch import nn\n'), ((135, 35, 135, 53), 'torch.nn.Linear', 'nn.Linear', ({(135, 45, 135, 48): '512', (135, 49, 135, 52): '256'}, {}), '(512, 256)', False, 'from torch import nn\n'), ((135, 54, 135, 63), 'torch.nn.Tanh', 'nn.Tanh', ({}, {}), '()', False, 'from torch import nn\n'), ((136, 35, 136, 53), 'torch.nn.Linear', 'nn.Linear', ({(136, 45, 136, 48): '256', (136, 49, 136, 52): '256'}, {}), '(256, 256)', False, 'from torch import nn\n'), ((137, 35, 137, 77), 'base.basenetwork.EigenLayer', 'BaseN.EigenLayer', ({(137, 52, 137, 55): '256', (137, 56, 137, 76): 'self.output_shape[0]'}, {}), '(256, self.output_shape[0])', True, 'import base.basenetwork as BaseN\n'), ((149, 35, 149, 50), 'base.basenetwork.Flatten', 'BaseN.Flatten', ({}, {}), '()', True, 'import base.basenetwork as BaseN\n'), ((151, 35, 151, 53), 'torch.nn.Linear', 'nn.Linear', ({(151, 45, 151, 48): '512', (151, 49, 151, 52): '256'}, {}), '(512, 256)', False, 'from torch import nn\n'), ((151, 54, 151, 63), 'torch.nn.Tanh', 'nn.Tanh', ({}, {}), '()', False, 'from torch import nn\n'), ((152, 35, 152, 53), 'torch.nn.Linear', 'nn.Linear', ({(152, 45, 152, 48): '256', (152, 49, 152, 52): '512'}, {}), '(256, 512)', False, 'from torch import nn\n'), ((153, 35, 153, 77), 'base.basenetwork.EigenLayer', 'BaseN.EigenLayer', ({(153, 52, 153, 55): '512', (153, 56, 153, 76): 'self.output_shape[0]'}, {}), '(512, self.output_shape[0])', True, 'import base.basenetwork as BaseN\n'), ((165, 35, 165, 50), 'base.basenetwork.Flatten', 'BaseN.Flatten', ({}, {}), '()', True, 'import base.basenetwork as BaseN\n'), ((167, 35, 167, 53), 'torch.nn.Linear', 'nn.Linear', ({(167, 45, 167, 48): '512', (167, 49, 167, 52): '256'}, {}), '(512, 256)', False, 'from torch import nn\n'), ((167, 54, 167, 63), 'torch.nn.Tanh', 'nn.Tanh', ({}, {}), '()', False, 'from torch import nn\n'), ((168, 35, 168, 53), 'torch.nn.Linear', 'nn.Linear', ({(168, 45, 168, 48): '256', (168, 49, 168, 52): '512'}, {}), '(256, 512)', False, 'from torch import nn\n'), ((169, 35, 169, 88), 'base.basenetwork.EigenLayer', 'BaseN.EigenLayer', (), '', True, 'import base.basenetwork as BaseN\n'), ((181, 35, 181, 50), 'base.basenetwork.Flatten', 'BaseN.Flatten', ({}, {}), '()', True, 'import base.basenetwork as BaseN\n'), ((183, 35, 183, 53), 'torch.nn.Linear', 'nn.Linear', ({(183, 45, 183, 48): '512', (183, 49, 183, 52): '512'}, {}), '(512, 512)', False, 'from torch import nn\n'), ((183, 54, 183, 63), 'torch.nn.Tanh', 'nn.Tanh', ({}, {}), '()', False, 'from torch import nn\n'), ((184, 35, 184, 54), 'torch.nn.Linear', 'nn.Linear', ({(184, 45, 184, 48): '512', (184, 49, 184, 53): '1024'}, {}), '(512, 1024)', False, 'from torch import nn\n'), ((185, 35, 185, 78), 'base.basenetwork.EigenLayer', 'BaseN.EigenLayer', ({(185, 52, 185, 56): '1024', (185, 57, 185, 77): 'self.output_shape[0]'}, {}), '(1024, self.output_shape[0])', True, 'import base.basenetwork as BaseN\n'), ((198, 35, 198, 50), 'base.basenetwork.Flatten', 'BaseN.Flatten', ({}, {}), '()', True, 'import base.basenetwork as BaseN\n'), ((200, 35, 200, 53), 'torch.nn.Linear', 'nn.Linear', ({(200, 45, 200, 48): '512', (200, 49, 200, 52): '256'}, {}), '(512, 256)', False, 'from torch import nn\n'), ((200, 54, 200, 63), 'torch.nn.Tanh', 'nn.Tanh', ({}, {}), '()', False, 'from torch import nn\n'), ((201, 35, 201, 53), 'torch.nn.Linear', 'nn.Linear', ({(201, 45, 201, 48): '256', (201, 49, 201, 52): '512'}, {}), '(256, 512)', False, 'from torch import nn\n'), ((202, 35, 202, 70), 'torch.nn.Linear', 'nn.Linear', ({(202, 45, 202, 48): '512', (202, 49, 202, 69): 'self.output_shape[0]'}, {}), '(512, self.output_shape[0])', False, 'from torch import nn\n'), ((210, 21, 210, 44), 'base.basenetwork.ResNetBlock', 'BaseN.ResNetBlock', ({(210, 39, 210, 40): '(1)', (210, 41, 210, 43): '(32)'}, {}), '(1, 32)', True, 'import base.basenetwork as BaseN\n'), ((211, 25, 211, 45), 'base.basenetwork.conv3_2', 'BaseN.conv3_2', ({(211, 39, 211, 41): '(32)', (211, 42, 211, 44): '(64)'}, {}), '(32, 64)', True, 'import base.basenetwork as BaseN\n'), ((213, 49, 213, 62), 'torch.nn.Softplus', 'nn.Softplus', ({}, {}), '()', False, 'from torch import nn\n'), ((214, 36, 214, 51), 'base.basenetwork.Flatten', 'BaseN.Flatten', ({}, {}), '()', True, 'import base.basenetwork as BaseN\n'), ((216, 36, 216, 54), 'torch.nn.Linear', 'nn.Linear', ({(216, 46, 216, 49): '512', (216, 50, 216, 53): '256'}, {}), '(512, 256)', False, 'from torch import nn\n'), ((216, 55, 216, 64), 'torch.nn.Tanh', 'nn.Tanh', ({}, {}), '()', False, 'from torch import nn\n'), ((217, 36, 217, 78), 'base.basenetwork.EigenLayer', 'BaseN.EigenLayer', ({(217, 53, 217, 56): '256', (217, 57, 217, 77): 'self.output_shape[0]'}, {}), '(256, self.output_shape[0])', True, 'import base.basenetwork as BaseN\n'), ((228, 35, 228, 50), 'base.basenetwork.Flatten', 'BaseN.Flatten', ({}, {}), '()', True, 'import base.basenetwork as BaseN\n'), ((230, 35, 230, 53), 'torch.nn.Linear', 'nn.Linear', ({(230, 45, 230, 48): '512', (230, 49, 230, 52): '256'}, {}), '(512, 256)', False, 'from torch import nn\n'), ((230, 54, 230, 63), 'torch.nn.Tanh', 'nn.Tanh', ({}, {}), '()', False, 'from torch import nn\n'), ((231, 35, 231, 70), 'torch.nn.Linear', 'nn.Linear', ({(231, 45, 231, 48): '256', (231, 49, 231, 69): 'self.output_shape[0]'}, {}), '(256, self.output_shape[0])', False, 'from torch import nn\n'), ((240, 35, 240, 50), 'base.basenetwork.Flatten', 'BaseN.Flatten', ({}, {}), '()', True, 'import base.basenetwork as BaseN\n'), ((241, 63, 241, 76), 'torch.nn.Softplus', 'nn.Softplus', ({}, {}), '()', False, 'from torch import nn\n'), ((242, 35, 242, 54), 'torch.nn.Linear', 'nn.Linear', ({(242, 45, 242, 49): '1024', (242, 50, 242, 53): '512'}, {}), '(1024, 512)', False, 'from torch import nn\n'), ((243, 35, 243, 53), 'torch.nn.Linear', 'nn.Linear', ({(243, 45, 243, 48): '512', (243, 49, 243, 52): '256'}, {}), '(512, 256)', False, 'from torch import nn\n'), ((243, 54, 243, 63), 'torch.nn.Tanh', 'nn.Tanh', ({}, {}), '()', False, 'from torch import nn\n'), ((244, 35, 244, 70), 'torch.nn.Linear', 'nn.Linear', ({(244, 45, 244, 48): '256', (244, 49, 244, 69): 'self.output_shape[0]'}, {}), '(256, self.output_shape[0])', False, 'from torch import nn\n'), ((13, 45, 13, 55), 'numpy.prod', 'np.prod', ({(13, 53, 13, 54): 'x'}, {}), '(x)', True, 'import numpy as np\n'), ((26, 45, 26, 55), 'numpy.prod', 'np.prod', ({(26, 53, 26, 54): 'x'}, {}), '(x)', True, 'import numpy as np\n'), ((38, 45, 38, 55), 'numpy.prod', 'np.prod', ({(38, 53, 38, 54): 'x'}, {}), '(x)', True, 'import numpy as np\n'), ((61, 35, 61, 67), 'base.basenetwork.conv3_2', 'BaseN.conv3_2', ({(61, 49, 61, 63): 'input_shape[0]', (61, 65, 61, 66): '(8)'}, {}), '(input_shape[0], 8)', True, 'import base.basenetwork as BaseN\n'), ((61, 68, 61, 77), 'torch.nn.ReLU', 'nn.ReLU', ({}, {}), '()', False, 'from torch import nn\n'), ((62, 49, 62, 69), 'base.basenetwork.conv3_2', 'BaseN.conv3_2', ({(62, 63, 62, 64): '(8)', (62, 66, 62, 68): '(16)'}, {}), '(8, 16)', True, 'import base.basenetwork as BaseN\n'), ((62, 70, 62, 79), 'torch.nn.ReLU', 'nn.ReLU', ({}, {}), '()', False, 'from torch import nn\n'), ((63, 49, 63, 68), 'base.basenetwork.conv3_2', 'BaseN.conv3_2', ({(63, 63, 63, 64): '(8)', (63, 66, 63, 67): '(8)'}, {}), '(8, 8)', True, 'import base.basenetwork as BaseN\n'), ((67, 45, 67, 55), 'numpy.prod', 'np.prod', ({(67, 53, 67, 54): 'x'}, {}), '(x)', True, 'import numpy as np\n'), ((76, 35, 76, 67), 'base.basenetwork.conv3_2', 'BaseN.conv3_2', ({(76, 49, 76, 63): 'input_shape[0]', (76, 65, 76, 66): '(8)'}, {}), '(input_shape[0], 8)', True, 'import base.basenetwork as BaseN\n'), ((76, 68, 76, 81), 'torch.nn.Softplus', 'nn.Softplus', ({}, {}), '()', False, 'from torch import nn\n'), ((77, 35, 77, 55), 'base.basenetwork.conv3_2', 'BaseN.conv3_2', ({(77, 49, 77, 50): '(8)', (77, 52, 77, 54): '(12)'}, {}), '(8, 12)', True, 'import base.basenetwork as BaseN\n'), ((77, 56, 77, 76), 'base.basenetwork.AdaptiveTanh', 'BaseN.AdaptiveTanh', ({}, {}), '()', True, 'import base.basenetwork as BaseN\n'), ((78, 35, 78, 56), 'base.basenetwork.conv3_2', 'BaseN.conv3_2', ({(78, 49, 78, 51): '(12)', (78, 53, 78, 55): '(16)'}, {}), '(12, 16)', True, 'import base.basenetwork as BaseN\n'), ((79, 35, 79, 56), 'base.basenetwork.conv3_2', 'BaseN.conv3_2', ({(79, 49, 79, 51): '(16)', (79, 53, 79, 55): '(20)'}, {}), '(16, 20)', True, 'import base.basenetwork as BaseN\n'), ((83, 45, 83, 55), 'numpy.prod', 'np.prod', ({(83, 53, 83, 54): 'x'}, {}), '(x)', True, 'import numpy as np\n'), ((93, 35, 93, 67), 'base.basenetwork.conv3_2', 'BaseN.conv3_2', ({(93, 49, 93, 63): 'input_shape[0]', (93, 65, 93, 66): '(4)'}, {}), '(input_shape[0], 4)', True, 'import base.basenetwork as BaseN\n'), ((93, 68, 93, 81), 'torch.nn.Softplus', 'nn.Softplus', ({}, {}), '()', False, 'from torch import nn\n'), ((94, 35, 94, 54), 'base.basenetwork.conv3_2', 'BaseN.conv3_2', ({(94, 49, 94, 50): '(4)', (94, 52, 94, 53): '(8)'}, {}), '(4, 8)', True, 'import base.basenetwork as BaseN\n'), ((94, 55, 94, 75), 'base.basenetwork.AdaptiveTanh', 'BaseN.AdaptiveTanh', ({}, {}), '()', True, 'import base.basenetwork as BaseN\n'), ((98, 45, 98, 55), 'numpy.prod', 'np.prod', ({(98, 53, 98, 54): 'x'}, {}), '(x)', True, 'import numpy as np\n'), ((109, 35, 109, 67), 'base.basenetwork.conv3_2', 'BaseN.conv3_2', ({(109, 49, 109, 63): 'input_shape[0]', (109, 65, 109, 66): '(4)'}, {}), '(input_shape[0], 4)', True, 'import base.basenetwork as BaseN\n'), ((109, 68, 109, 77), 'torch.nn.ReLU', 'nn.ReLU', ({}, {}), '()', False, 'from torch import nn\n'), ((110, 35, 110, 54), 'base.basenetwork.conv3_2', 'BaseN.conv3_2', ({(110, 49, 110, 50): '(4)', (110, 52, 110, 53): '(4)'}, {}), '(4, 4)', True, 'import base.basenetwork as BaseN\n'), ((110, 55, 110, 75), 'base.basenetwork.AdaptiveTanh', 'BaseN.AdaptiveTanh', ({}, {}), '()', True, 'import base.basenetwork as BaseN\n'), ((114, 45, 114, 55), 'numpy.prod', 'np.prod', ({(114, 53, 114, 54): 'x'}, {}), '(x)', True, 'import numpy as np\n'), ((125, 35, 125, 67), 'base.basenetwork.conv3_2', 'BaseN.conv3_2', ({(125, 49, 125, 63): 'input_shape[0]', (125, 65, 125, 66): '(3)'}, {}), '(input_shape[0], 3)', True, 'import base.basenetwork as BaseN\n'), ((125, 68, 125, 81), 'torch.nn.Softplus', 'nn.Softplus', ({}, {}), '()', False, 'from torch import nn\n'), ((126, 49, 126, 68), 'base.basenetwork.conv3_2', 'BaseN.conv3_2', ({(126, 63, 126, 64): '(3)', (126, 66, 126, 67): '(6)'}, {}), '(3, 6)', True, 'import base.basenetwork as BaseN\n'), ((126, 69, 126, 89), 'base.basenetwork.conv3_2', 'BaseN.conv3_2', ({(126, 83, 126, 84): '(6)', (126, 86, 126, 88): '(12)'}, {}), '(6, 12)', True, 'import base.basenetwork as BaseN\n'), ((130, 45, 130, 55), 'numpy.prod', 'np.prod', ({(130, 53, 130, 54): 'x'}, {}), '(x)', True, 'import numpy as np\n'), ((144, 35, 144, 67), 'base.basenetwork.conv3_2', 'BaseN.conv3_2', ({(144, 49, 144, 63): 'input_shape[0]', (144, 65, 144, 66): '(8)'}, {}), '(input_shape[0], 8)', True, 'import base.basenetwork as BaseN\n'), ((144, 68, 144, 81), 'torch.nn.Softplus', 'nn.Softplus', ({}, {}), '()', False, 'from torch import nn\n'), ((145, 49, 145, 69), 'base.basenetwork.conv3_2', 'BaseN.conv3_2', ({(145, 63, 145, 64): '(8)', (145, 66, 145, 68): '(16)'}, {}), '(8, 16)', True, 'import base.basenetwork as BaseN\n'), ((145, 70, 145, 83), 'torch.nn.Softplus', 'nn.Softplus', ({}, {}), '()', False, 'from torch import nn\n'), ((146, 49, 146, 70), 'base.basenetwork.conv3_2', 'BaseN.conv3_2', ({(146, 63, 146, 65): '(16)', (146, 67, 146, 69): '(32)'}, {}), '(16, 32)', True, 'import base.basenetwork as BaseN\n'), ((150, 45, 150, 55), 'numpy.prod', 'np.prod', ({(150, 53, 150, 54): 'x'}, {}), '(x)', True, 'import numpy as np\n'), ((161, 35, 161, 67), 'base.basenetwork.conv3_2', 'BaseN.conv3_2', ({(161, 49, 161, 63): 'input_shape[0]', (161, 65, 161, 66): '(4)'}, {}), '(input_shape[0], 4)', True, 'import base.basenetwork as BaseN\n'), ((161, 68, 161, 81), 'torch.nn.Softplus', 'nn.Softplus', ({}, {}), '()', False, 'from torch import nn\n'), ((162, 49, 162, 68), 'base.basenetwork.conv3_2', 'BaseN.conv3_2', ({(162, 63, 162, 64): '(4)', (162, 66, 162, 67): '(4)'}, {}), '(4, 4)', True, 'import base.basenetwork as BaseN\n'), ((162, 69, 162, 89), 'base.basenetwork.AdaptiveTanh', 'BaseN.AdaptiveTanh', ({}, {}), '()', True, 'import base.basenetwork as BaseN\n'), ((166, 45, 166, 55), 'numpy.prod', 'np.prod', ({(166, 53, 166, 54): 'x'}, {}), '(x)', True, 'import numpy as np\n'), ((176, 35, 176, 67), 'base.basenetwork.conv3_2', 'BaseN.conv3_2', ({(176, 49, 176, 63): 'input_shape[0]', (176, 65, 176, 66): '(8)'}, {}), '(input_shape[0], 8)', True, 'import base.basenetwork as BaseN\n'), ((176, 68, 176, 81), 'torch.nn.Softplus', 'nn.Softplus', ({}, {}), '()', False, 'from torch import nn\n'), ((177, 49, 177, 69), 'base.basenetwork.conv3_2', 'BaseN.conv3_2', ({(177, 63, 177, 64): '(8)', (177, 66, 177, 68): '(16)'}, {}), '(8, 16)', True, 'import base.basenetwork as BaseN\n'), ((178, 49, 178, 70), 'base.basenetwork.conv3_2', 'BaseN.conv3_2', ({(178, 63, 178, 65): '(16)', (178, 67, 178, 69): '(32)'}, {}), '(16, 32)', True, 'import base.basenetwork as BaseN\n'), ((182, 45, 182, 55), 'numpy.prod', 'np.prod', ({(182, 53, 182, 54): 'x'}, {}), '(x)', True, 'import numpy as np\n'), ((193, 35, 193, 67), 'base.basenetwork.conv3_2', 'BaseN.conv3_2', ({(193, 49, 193, 63): 'input_shape[0]', (193, 65, 193, 66): '(8)'}, {}), '(input_shape[0], 8)', True, 'import base.basenetwork as BaseN\n'), ((193, 68, 193, 81), 'torch.nn.Softplus', 'nn.Softplus', ({}, {}), '()', False, 'from torch import nn\n'), ((194, 49, 194, 69), 'base.basenetwork.conv3_2', 'BaseN.conv3_2', ({(194, 63, 194, 64): '(8)', (194, 66, 194, 68): '(16)'}, {}), '(8, 16)', True, 'import base.basenetwork as BaseN\n'), ((195, 49, 195, 70), 'base.basenetwork.conv3_2', 'BaseN.conv3_2', ({(195, 63, 195, 65): '(16)', (195, 67, 195, 69): '(32)'}, {}), '(16, 32)', True, 'import base.basenetwork as BaseN\n'), ((199, 45, 199, 55), 'numpy.prod', 'np.prod', ({(199, 53, 199, 54): 'x'}, {}), '(x)', True, 'import numpy as np\n'), ((215, 46, 215, 56), 'numpy.prod', 'np.prod', ({(215, 54, 215, 55): 'x'}, {}), '(x)', True, 'import numpy as np\n'), ((225, 35, 225, 67), 'base.basenetwork.conv3_2', 'BaseN.conv3_2', ({(225, 49, 225, 63): 'input_shape[0]', (225, 65, 225, 66): '(4)'}, {}), '(input_shape[0], 4)', True, 'import base.basenetwork as BaseN\n'), ((225, 68, 225, 81), 'torch.nn.Softplus', 'nn.Softplus', ({}, {}), '()', False, 'from torch import nn\n'), ((229, 45, 229, 55), 'numpy.prod', 'np.prod', ({(229, 53, 229, 54): 'x'}, {}), '(x)', True, 'import numpy as np\n'), ((241, 45, 241, 55), 'numpy.prod', 'np.prod', ({(241, 53, 241, 54): 'x'}, {}), '(x)', True, 'import numpy as np\n')]