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