import math
import numpy as np
import scipy as sp
import scipy.linalg
import torch
import torch.nn as nn
import torch.nn.init as init
import torch.nn.functional as F
from nf.utils import unconstrained_RQS

# supported non-linearities: note that the function must be invertible
functional_derivatives = {
    torch.tanh: lambda x: 1 - torch.pow(torch.tanh(x), 2),
    F.leaky_relu: lambda x: (x > 0).type(torch.FloatTensor) + \
                            (x < 0).type(torch.FloatTensor) * -0.01,
    F.elu: lambda x: (x > 0).type(torch.FloatTensor) + \
                     (x < 0).type(torch.FloatTensor) * torch.exp(x)
}


class Planar(nn.Module):
    """
    Planar flow.

        z = f(x) = x + u h(wᵀx + b)

    [Rezende and Mohamed, 2015]
    """
    def __init__(self, dim, nonlinearity=torch.tanh):
        super().__init__()
        self.h = nonlinearity
        self.w = nn.Parameter(torch.Tensor(dim))
        self.u = nn.Parameter(torch.Tensor(dim))
        self.b = nn.Parameter(torch.Tensor(1))
        self.reset_parameters(dim)

    def reset_parameters(self, dim):
        init.uniform_(self.w, -math.sqrt(1/dim), math.sqrt(1/dim))
        init.uniform_(self.u, -math.sqrt(1/dim), math.sqrt(1/dim))
        init.uniform_(self.b, -math.sqrt(1/dim), math.sqrt(1/dim))

    def forward(self, x):
        """
        Given x, returns z and the log-determinant log|df/dx|.

        Returns
        -------
        """
        if self.h in (F.elu, F.leaky_relu):
            u = self.u
        elif self.h == torch.tanh:
            scal = torch.log(1+torch.exp(self.w @ self.u)) - self.w @ self.u - 1
            u = self.u + scal * self.w / torch.norm(self.w)
        else:
            raise NotImplementedError("Non-linearity is not supported.")
        lin = torch.unsqueeze(x @ self.w, 1) + self.b
        z = x + u * self.h(lin)
        phi = functional_derivatives[self.h](lin) * self.w
        log_det = torch.log(torch.abs(1 + phi @ u) + 1e-4)
        return z, log_det

    def backward(self, z):
        raise NotImplementedError("Planar flow has no algebraic inverse.")


class Radial(nn.Module):
    """
    Radial flow.

        z = f(x) = = x + β h(α, r)(z − z0)

    [Rezende and Mohamed 2015]
    """
    def __init__(self, dim):
        super().__init__()
        self.x0 = nn.Parameter(torch.Tensor(dim))
        self.log_alpha = nn.Parameter(torch.Tensor(1))
        self.beta = nn.Parameter(torch.Tensor(1))

    def reset_parameters(dim):
        init.uniform_(self.z0, -math.sqrt(1/dim), math.sqrt(1/dim))
        init.uniform_(self.log_alpha, -math.sqrt(1/dim), math.sqrt(1/dim))
        init.uniform_(self.beta, -math.sqrt(1/dim), math.sqrt(1/dim))

    def forward(self, x):
        """
        Given x, returns z and the log-determinant log|df/dx|.
        """
        m, n = x.shape
        r = torch.norm(x - self.x0)
        h = 1 / (torch.exp(self.log_alpha) + r)
        beta = -torch.exp(self.log_alpha) + torch.log(1 + torch.exp(self.beta))
        z = x + beta * h * (x - self.x0)
        log_det = (n - 1) * torch.log(1 + beta * h) + \
                  torch.log(1 + beta * h - \
                            beta * r / (torch.exp(self.log_alpha) + r) ** 2)
        return z, log_det


class FCNN(nn.Module):
    """
    Simple fully connected neural network.
    """
    def __init__(self, in_dim, out_dim, hidden_dim):
        super().__init__()
        self.network = nn.Sequential(
            nn.Linear(in_dim, hidden_dim),
            nn.Tanh(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.Tanh(),
            nn.Linear(hidden_dim, out_dim),
        )

    def forward(self, x):
        return self.network(x)


class RealNVP(nn.Module):
    """
    Non-volume preserving flow.

    [Dinh et. al. 2017]
    """
    def __init__(self, dim, hidden_dim = 8, base_network=FCNN):
        super().__init__()
        self.dim = dim
        self.t1 = base_network(dim // 2, dim // 2, hidden_dim)
        self.s1 = base_network(dim // 2, dim // 2, hidden_dim)
        self.t2 = base_network(dim // 2, dim // 2, hidden_dim)
        self.s2 = base_network(dim // 2, dim // 2, hidden_dim)

    def forward(self, x):
        lower, upper = x[:,:self.dim // 2], x[:,self.dim // 2:]
        t1_transformed = self.t1(lower)
        s1_transformed = self.s1(lower)
        upper = t1_transformed + upper * torch.exp(s1_transformed)
        t2_transformed = self.t2(upper)
        s2_transformed = self.s2(upper)
        lower = t2_transformed + lower * torch.exp(s2_transformed)
        z = torch.cat([lower, upper], dim=1)
        log_det = torch.sum(s1_transformed, dim=1) + \
                  torch.sum(s2_transformed, dim=1)
        return z, log_det

    def backward(self, z):
        lower, upper = z[:,:self.dim // 2], z[:,self.dim // 2:]
        t2_transformed = self.t2(upper)
        s2_transformed = self.s2(upper)
        lower = (lower - t2_transformed) * torch.exp(-s2_transformed)
        t1_transformed = self.t1(lower)
        s1_transformed = self.s1(lower)
        upper = (upper - t1_transformed) * torch.exp(-s1_transformed)
        x = torch.cat([lower, upper], dim=1)
        log_det = torch.sum(-s1_transformed, dim=1) + \
                  torch.sum(-s2_transformed, dim=1)
        return x, log_det


class MAF(nn.Module):
    """
    Masked auto-regressive flow.

    [Papamakarios et al. 2018]
    """
    def __init__(self, dim, hidden_dim = 8, base_network=FCNN):
        super().__init__()
        self.dim = dim
        self.layers = nn.ModuleList()
        self.initial_param = nn.Parameter(torch.Tensor(2))
        for i in range(1, dim):
            self.layers += [base_network(i, 2, hidden_dim)]
        self.reset_parameters()

    def reset_parameters(self):
        init.uniform_(self.initial_param, -math.sqrt(0.5), math.sqrt(0.5))

    def forward(self, x):
        z = torch.zeros_like(x)
        log_det = torch.zeros(z.shape[0])
        for i in range(self.dim):
            if i == 0:
                mu, alpha = self.initial_param[0], self.initial_param[1]
            else:
                out = self.layers[i - 1](x[:, :i])
                mu, alpha = out[:, 0], out[:, 1]
            z[:, i] = (x[:, i] - mu) / torch.exp(alpha)
            log_det -= alpha
        return z.flip(dims=(1,)), log_det

    def backward(self, z):
        x = torch.zeros_like(z)
        log_det = torch.zeros(z.shape[0])
        z = z.flip(dims=(1,))
        for i in range(self.dim):
            if i == 0:
                mu, alpha = self.initial_param[0], self.initial_param[1]
            else:
                out = self.layers[i - 1](x[:, :i])
                mu, alpha = out[:, 0], out[:, 1]
            x[:, i] = mu + torch.exp(alpha) * z[:, i]
            log_det += alpha
        return x, log_det


class ActNorm(nn.Module):
    """
    ActNorm layer.

    [Kingma and Dhariwal, 2018.]
    """
    def __init__(self, dim):
        super().__init__()
        self.dim = dim
        self.mu = nn.Parameter(torch.zeros(dim, dtype = torch.float))
        self.log_sigma = nn.Parameter(torch.zeros(dim, dtype = torch.float))

    def forward(self, x):
        z = x * torch.exp(self.log_sigma) + self.mu
        log_det = torch.sum(self.log_sigma)
        return z, log_det

    def backward(self, z):
        x = (z - self.mu) / torch.exp(self.log_sigma)
        log_det = -torch.sum(self.log_sigma)
        return x, log_det


class OneByOneConv(nn.Module):
    """
    Invertible 1x1 convolution.

    [Kingma and Dhariwal, 2018.]
    """
    def __init__(self, dim):
        super().__init__()
        self.dim = dim
        W, _ = sp.linalg.qr(np.random.randn(dim, dim))
        P, L, U = sp.linalg.lu(W)
        self.P = torch.tensor(P, dtype = torch.float)
        self.L = nn.Parameter(torch.tensor(L, dtype = torch.float))
        self.S = nn.Parameter(torch.tensor(np.diag(U), dtype = torch.float))
        self.U = nn.Parameter(torch.triu(torch.tensor(U, dtype = torch.float),
                              diagonal = 1))
        self.W_inv = None

    def forward(self, x):
        L = torch.tril(self.L, diagonal = -1) + torch.diag(torch.ones(self.dim))
        U = torch.triu(self.U, diagonal = 1)
        z = x @ self.P @ L @ (U + torch.diag(self.S))
        log_det = torch.sum(torch.log(torch.abs(self.S)))
        return z, log_det

    def backward(self, z):
        if not self.W_inv:
            L = torch.tril(self.L, diagonal = -1) + \
                torch.diag(torch.ones(self.dim))
            U = torch.triu(self.U, diagonal = 1)
            W = self.P @ L @ (U + torch.diag(self.S))
            self.W_inv = torch.inverse(W)
        x = z @ self.W_inv
        log_det = -torch.sum(torch.log(torch.abs(self.S)))
        return x, log_det


class NSF_AR(nn.Module):
    """
    Neural spline flow, auto-regressive.

    [Durkan et al. 2019]
    """
    def __init__(self, dim, K = 5, B = 3, hidden_dim = 8, base_network = FCNN):
        super().__init__()
        self.dim = dim
        self.K = K
        self.B = B
        self.layers = nn.ModuleList()
        self.init_param = nn.Parameter(torch.Tensor(3 * K - 1))
        for i in range(1, dim):
            self.layers += [base_network(i, 3 * K - 1, hidden_dim)]
        self.reset_parameters()

    def reset_parameters(self):
        init.uniform_(self.init_param, - 1 / 2, 1 / 2)

    def forward(self, x):
        z = torch.zeros_like(x)
        log_det = torch.zeros(z.shape[0])
        for i in range(self.dim):
            if i == 0:
                init_param = self.init_param.expand(x.shape[0], 3 * self.K - 1)
                W, H, D = torch.split(init_param, self.K, dim = 1)
            else:
                out = self.layers[i - 1](x[:, :i])
                W, H, D = torch.split(out, self.K, dim = 1)
            W, H = torch.softmax(W, dim = 1), torch.softmax(H, dim = 1)
            W, H = 2 * self.B * W, 2 * self.B * H
            D = F.softplus(D)
            z[:, i], ld = unconstrained_RQS(
                x[:, i], W, H, D, inverse=False, tail_bound=self.B)
            log_det += ld
        return z, log_det

    def backward(self, z):
        x = torch.zeros_like(z)
        log_det = torch.zeros(x.shape[0])
        for i in range(self.dim):
            if i == 0:
                init_param = self.init_param.expand(x.shape[0], 3 * self.K - 1)
                W, H, D = torch.split(init_param, self.K, dim = 1)
            else:
                out = self.layers[i - 1](x[:, :i])
                W, H, D = torch.split(out, self.K, dim = 1)
            W, H = torch.softmax(W, dim = 1), torch.softmax(H, dim = 1)
            W, H = 2 * self.B * W, 2 * self.B * H
            D = F.softplus(D)
            x[:, i], ld = unconstrained_RQS(
                z[:, i], W, H, D, inverse = True, tail_bound = self.B)
            log_det += ld
        return x, log_det


class NSF_CL(nn.Module):
    """
    Neural spline flow, coupling layer.

    [Durkan et al. 2019]
    """
    def __init__(self, dim, K = 5, B = 3, hidden_dim = 8, base_network = FCNN):
        super().__init__()
        self.dim = dim
        self.K = K
        self.B = B
        self.f1 = base_network(dim // 2, (3 * K - 1) * dim // 2, hidden_dim)
        self.f2 = base_network(dim // 2, (3 * K - 1) * dim // 2, hidden_dim)

    def forward(self, x):
        log_det = torch.zeros(x.shape[0])
        lower, upper = x[:, :self.dim // 2], x[:, self.dim // 2:]
        out = self.f1(lower).reshape(-1, self.dim // 2, 3 * self.K - 1)
        W, H, D = torch.split(out, self.K, dim = 2)
        W, H = torch.softmax(W, dim = 2), torch.softmax(H, dim = 2)
        W, H = 2 * self.B * W, 2 * self.B * H
        D = F.softplus(D)
        upper, ld = unconstrained_RQS(
            upper, W, H, D, inverse=False, tail_bound=self.B)
        log_det += torch.sum(ld, dim = 1)
        out = self.f2(upper).reshape(-1, self.dim // 2, 3 * self.K - 1)
        W, H, D = torch.split(out, self.K, dim = 2)
        W, H = torch.softmax(W, dim = 2), torch.softmax(H, dim = 2)
        W, H = 2 * self.B * W, 2 * self.B * H
        D = F.softplus(D)
        lower, ld = unconstrained_RQS(
            lower, W, H, D, inverse=False, tail_bound=self.B)
        log_det += torch.sum(ld, dim = 1)
        return torch.cat([lower, upper], dim = 1), log_det

    def backward(self, z):
        log_det = torch.zeros(z.shape[0])
        lower, upper = z[:, :self.dim // 2], z[:, self.dim // 2:]
        out = self.f2(upper).reshape(-1, self.dim // 2, 3 * self.K - 1)
        W, H, D = torch.split(out, self.K, dim = 2)
        W, H = torch.softmax(W, dim = 2), torch.softmax(H, dim = 2)
        W, H = 2 * self.B * W, 2 * self.B * H
        D = F.softplus(D)
        lower, ld = unconstrained_RQS(
            lower, W, H, D, inverse=True, tail_bound=self.B)
        log_det += torch.sum(ld, dim = 1)
        out = self.f1(lower).reshape(-1, self.dim // 2, 3 * self.K - 1)
        W, H, D = torch.split(out, self.K, dim = 2)
        W, H = torch.softmax(W, dim = 2), torch.softmax(H, dim = 2)
        W, H = 2 * self.B * W, 2 * self.B * H
        D = F.softplus(D)
        upper, ld = unconstrained_RQS(
            upper, W, H, D, inverse = True, tail_bound = self.B)
        log_det += torch.sum(ld, dim = 1)
        return torch.cat([lower, upper], dim = 1), log_det

[(90, 'torch.norm', 'torch.norm', (['(x - self.x0)'], {}), False, 'import torch\n'), (140, 'torch.cat', 'torch.cat', (['[lower, upper]'], {'dim': '(1)'}), False, 'import torch\n'), (153, 'torch.cat', 'torch.cat', (['[lower, upper]'], {'dim': '(1)'}), False, 'import torch\n'), (168, 'torch.nn.ModuleList', 'nn.ModuleList', ([], {}), True, 'import torch.nn as nn\n'), (178, 'torch.zeros_like', 'torch.zeros_like', (['x'], {}), False, 'import torch\n'), (179, 'torch.zeros', 'torch.zeros', (['z.shape[0]'], {}), False, 'import torch\n'), (191, 'torch.zeros_like', 'torch.zeros_like', (['z'], {}), False, 'import torch\n'), (192, 'torch.zeros', 'torch.zeros', (['z.shape[0]'], {}), False, 'import torch\n'), (219, 'torch.sum', 'torch.sum', (['self.log_sigma'], {}), False, 'import torch\n'), (238, 'scipy.linalg.lu', 'sp.linalg.lu', (['W'], {}), True, 'import scipy as sp\n'), (239, 'torch.tensor', 'torch.tensor', (['P'], {'dtype': 'torch.float'}), False, 'import torch\n'), (248, 'torch.triu', 'torch.triu', (['self.U'], {'diagonal': '(1)'}), False, 'import torch\n'), (276, 'torch.nn.ModuleList', 'nn.ModuleList', ([], {}), True, 'import torch.nn as nn\n'), (283, 'torch.nn.init.uniform_', 'init.uniform_', (['self.init_param', '(-1 / 2)', '(1 / 2)'], {}), True, 'import torch.nn.init as init\n'), (286, 'torch.zeros_like', 'torch.zeros_like', (['x'], {}), False, 'import torch\n'), (287, 'torch.zeros', 'torch.zeros', (['z.shape[0]'], {}), False, 'import torch\n'), (304, 'torch.zeros_like', 'torch.zeros_like', (['z'], {}), False, 'import torch\n'), (305, 'torch.zeros', 'torch.zeros', (['x.shape[0]'], {}), False, 'import torch\n'), (337, 'torch.zeros', 'torch.zeros', (['x.shape[0]'], {}), False, 'import torch\n'), (340, 'torch.split', 'torch.split', (['out', 'self.K'], {'dim': '(2)'}), False, 'import torch\n'), (343, 'torch.nn.functional.softplus', 'F.softplus', (['D'], {}), True, 'import torch.nn.functional as F\n'), (344, 'nf.utils.unconstrained_RQS', 'unconstrained_RQS', (['upper', 'W', 'H', 'D'], {'inverse': '(False)', 'tail_bound': 'self.B'}), False, 'from nf.utils import unconstrained_RQS\n'), (346, 'torch.sum', 'torch.sum', (['ld'], {'dim': '(1)'}), False, 'import torch\n'), (348, 'torch.split', 'torch.split', (['out', 'self.K'], {'dim': '(2)'}), False, 'import torch\n'), (351, 'torch.nn.functional.softplus', 'F.softplus', (['D'], {}), True, 'import torch.nn.functional as F\n'), (352, 'nf.utils.unconstrained_RQS', 'unconstrained_RQS', (['lower', 'W', 'H', 'D'], {'inverse': '(False)', 'tail_bound': 'self.B'}), False, 'from nf.utils import unconstrained_RQS\n'), (354, 'torch.sum', 'torch.sum', (['ld'], {'dim': '(1)'}), False, 'import torch\n'), (358, 'torch.zeros', 'torch.zeros', (['z.shape[0]'], {}), False, 'import torch\n'), (361, 'torch.split', 'torch.split', (['out', 'self.K'], {'dim': '(2)'}), False, 'import torch\n'), (364, 'torch.nn.functional.softplus', 'F.softplus', (['D'], {}), True, 'import torch.nn.functional as F\n'), (365, 'nf.utils.unconstrained_RQS', 'unconstrained_RQS', (['lower', 'W', 'H', 'D'], {'inverse': '(True)', 'tail_bound': 'self.B'}), False, 'from nf.utils import unconstrained_RQS\n'), (367, 'torch.sum', 'torch.sum', (['ld'], {'dim': '(1)'}), False, 'import torch\n'), (369, 'torch.split', 'torch.split', (['out', 'self.K'], {'dim': '(2)'}), False, 'import torch\n'), (372, 'torch.nn.functional.softplus', 'F.softplus', (['D'], {}), True, 'import torch.nn.functional as F\n'), (373, 'nf.utils.unconstrained_RQS', 'unconstrained_RQS', (['upper', 'W', 'H', 'D'], {'inverse': '(True)', 'tail_bound': 'self.B'}), False, 'from nf.utils import unconstrained_RQS\n'), (375, 'torch.sum', 'torch.sum', (['ld'], {'dim': '(1)'}), False, 'import torch\n'), (32, 'torch.Tensor', 'torch.Tensor', (['dim'], {}), False, 'import torch\n'), (33, 'torch.Tensor', 'torch.Tensor', (['dim'], {}), False, 'import torch\n'), (34, 'torch.Tensor', 'torch.Tensor', (['(1)'], {}), False, 'import torch\n'), (38, 'math.sqrt', 'math.sqrt', (['(1 / dim)'], {}), False, 'import math\n'), (39, 'math.sqrt', 'math.sqrt', (['(1 / dim)'], {}), False, 'import math\n'), (40, 'math.sqrt', 'math.sqrt', (['(1 / dim)'], {}), False, 'import math\n'), (56, 'torch.unsqueeze', 'torch.unsqueeze', (['(x @ self.w)', '(1)'], {}), False, 'import torch\n'), (76, 'torch.Tensor', 'torch.Tensor', (['dim'], {}), False, 'import torch\n'), (77, 'torch.Tensor', 'torch.Tensor', (['(1)'], {}), False, 'import torch\n'), (78, 'torch.Tensor', 'torch.Tensor', (['(1)'], {}), False, 'import torch\n'), (81, 'math.sqrt', 'math.sqrt', (['(1 / dim)'], {}), False, 'import math\n'), (82, 'math.sqrt', 'math.sqrt', (['(1 / dim)'], {}), False, 'import math\n'), (83, 'math.sqrt', 'math.sqrt', (['(1 / dim)'], {}), False, 'import math\n'), (107, 'torch.nn.Linear', 'nn.Linear', (['in_dim', 'hidden_dim'], {}), True, 'import torch.nn as nn\n'), (108, 'torch.nn.Tanh', 'nn.Tanh', ([], {}), True, 'import torch.nn as nn\n'), (109, 'torch.nn.Linear', 'nn.Linear', (['hidden_dim', 'hidden_dim'], {}), True, 'import torch.nn as nn\n'), (110, 'torch.nn.Tanh', 'nn.Tanh', ([], {}), True, 'import torch.nn as nn\n'), (111, 'torch.nn.Linear', 'nn.Linear', (['hidden_dim', 'out_dim'], {}), True, 'import torch.nn as nn\n'), (141, 'torch.sum', 'torch.sum', (['s1_transformed'], {'dim': '(1)'}), False, 'import torch\n'), (142, 'torch.sum', 'torch.sum', (['s2_transformed'], {'dim': '(1)'}), False, 'import torch\n'), (149, 'torch.exp', 'torch.exp', (['(-s2_transformed)'], {}), False, 'import torch\n'), (152, 'torch.exp', 'torch.exp', (['(-s1_transformed)'], {}), False, 'import torch\n'), (154, 'torch.sum', 'torch.sum', (['(-s1_transformed)'], {'dim': '(1)'}), False, 'import torch\n'), (155, 'torch.sum', 'torch.sum', (['(-s2_transformed)'], {'dim': '(1)'}), False, 'import torch\n'), (169, 'torch.Tensor', 'torch.Tensor', (['(2)'], {}), False, 'import torch\n'), (175, 'math.sqrt', 'math.sqrt', (['(0.5)'], {}), False, 'import math\n'), (214, 'torch.zeros', 'torch.zeros', (['dim'], {'dtype': 'torch.float'}), False, 'import torch\n'), (215, 'torch.zeros', 'torch.zeros', (['dim'], {'dtype': 'torch.float'}), False, 'import torch\n'), (223, 'torch.exp', 'torch.exp', (['self.log_sigma'], {}), False, 'import torch\n'), (224, 'torch.sum', 'torch.sum', (['self.log_sigma'], {}), False, 'import torch\n'), (237, 'numpy.random.randn', 'np.random.randn', (['dim', 'dim'], {}), True, 'import numpy as np\n'), (240, 'torch.tensor', 'torch.tensor', (['L'], {'dtype': 'torch.float'}), False, 'import torch\n'), (247, 'torch.tril', 'torch.tril', (['self.L'], {'diagonal': '(-1)'}), False, 'import torch\n'), (257, 'torch.triu', 'torch.triu', (['self.U'], {'diagonal': '(1)'}), False, 'import torch\n'), (259, 'torch.inverse', 'torch.inverse', (['W'], {}), False, 'import torch\n'), (277, 'torch.Tensor', 'torch.Tensor', (['(3 * K - 1)'], {}), False, 'import torch\n'), (297, 'torch.nn.functional.softplus', 'F.softplus', (['D'], {}), True, 'import torch.nn.functional as F\n'), (298, 'nf.utils.unconstrained_RQS', 'unconstrained_RQS', (['x[:, (i)]', 'W', 'H', 'D'], {'inverse': '(False)', 'tail_bound': 'self.B'}), False, 'from nf.utils import unconstrained_RQS\n'), (315, 'torch.nn.functional.softplus', 'F.softplus', (['D'], {}), True, 'import torch.nn.functional as F\n'), (316, 'nf.utils.unconstrained_RQS', 'unconstrained_RQS', (['z[:, (i)]', 'W', 'H', 'D'], {'inverse': '(True)', 'tail_bound': 'self.B'}), False, 'from nf.utils import unconstrained_RQS\n'), (341, 'torch.softmax', 'torch.softmax', (['W'], {'dim': '(2)'}), False, 'import torch\n'), (341, 'torch.softmax', 'torch.softmax', (['H'], {'dim': '(2)'}), False, 'import torch\n'), (349, 'torch.softmax', 'torch.softmax', (['W'], {'dim': '(2)'}), False, 'import torch\n'), (349, 'torch.softmax', 'torch.softmax', (['H'], {'dim': '(2)'}), False, 'import torch\n'), (355, 'torch.cat', 'torch.cat', (['[lower, upper]'], {'dim': '(1)'}), False, 'import torch\n'), (362, 'torch.softmax', 'torch.softmax', (['W'], {'dim': '(2)'}), False, 'import torch\n'), (362, 'torch.softmax', 'torch.softmax', (['H'], {'dim': '(2)'}), False, 'import torch\n'), (370, 'torch.softmax', 'torch.softmax', (['W'], {'dim': '(2)'}), False, 'import torch\n'), (370, 'torch.softmax', 'torch.softmax', (['H'], {'dim': '(2)'}), False, 'import torch\n'), (376, 'torch.cat', 'torch.cat', (['[lower, upper]'], {'dim': '(1)'}), False, 'import torch\n'), (13, 'torch.tanh', 'torch.tanh', (['x'], {}), False, 'import torch\n'), (17, 'torch.exp', 'torch.exp', (['x'], {}), False, 'import torch\n'), (38, 'math.sqrt', 'math.sqrt', (['(1 / dim)'], {}), False, 'import math\n'), (39, 'math.sqrt', 'math.sqrt', (['(1 / dim)'], {}), False, 'import math\n'), (40, 'math.sqrt', 'math.sqrt', (['(1 / dim)'], {}), False, 'import math\n'), (59, 'torch.abs', 'torch.abs', (['(1 + phi @ u)'], {}), False, 'import torch\n'), (81, 'math.sqrt', 'math.sqrt', (['(1 / dim)'], {}), False, 'import math\n'), (82, 'math.sqrt', 'math.sqrt', (['(1 / dim)'], {}), False, 'import math\n'), (83, 'math.sqrt', 'math.sqrt', (['(1 / dim)'], {}), False, 'import math\n'), (91, 'torch.exp', 'torch.exp', (['self.log_alpha'], {}), False, 'import torch\n'), (92, 'torch.exp', 'torch.exp', (['self.log_alpha'], {}), False, 'import torch\n'), (94, 'torch.log', 'torch.log', (['(1 + beta * h)'], {}), False, 'import torch\n'), (136, 'torch.exp', 'torch.exp', (['s1_transformed'], {}), False, 'import torch\n'), (139, 'torch.exp', 'torch.exp', (['s2_transformed'], {}), False, 'import torch\n'), (175, 'math.sqrt', 'math.sqrt', (['(0.5)'], {}), False, 'import math\n'), (186, 'torch.exp', 'torch.exp', (['alpha'], {}), False, 'import torch\n'), (218, 'torch.exp', 'torch.exp', (['self.log_sigma'], {}), False, 'import torch\n'), (241, 'numpy.diag', 'np.diag', (['U'], {}), True, 'import numpy as np\n'), (242, 'torch.tensor', 'torch.tensor', (['U'], {'dtype': 'torch.float'}), False, 'import torch\n'), (247, 'torch.ones', 'torch.ones', (['self.dim'], {}), False, 'import torch\n'), (249, 'torch.diag', 'torch.diag', (['self.S'], {}), False, 'import torch\n'), (250, 'torch.abs', 'torch.abs', (['self.S'], {}), False, 'import torch\n'), (255, 'torch.tril', 'torch.tril', (['self.L'], {'diagonal': '(-1)'}), False, 'import torch\n'), (291, 'torch.split', 'torch.split', (['init_param', 'self.K'], {'dim': '(1)'}), False, 'import torch\n'), (294, 'torch.split', 'torch.split', (['out', 'self.K'], {'dim': '(1)'}), False, 'import torch\n'), (295, 'torch.softmax', 'torch.softmax', (['W'], {'dim': '(1)'}), False, 'import torch\n'), (295, 'torch.softmax', 'torch.softmax', (['H'], {'dim': '(1)'}), False, 'import torch\n'), (309, 'torch.split', 'torch.split', (['init_param', 'self.K'], {'dim': '(1)'}), False, 'import torch\n'), (312, 'torch.split', 'torch.split', (['out', 'self.K'], {'dim': '(1)'}), False, 'import torch\n'), (313, 'torch.softmax', 'torch.softmax', (['W'], {'dim': '(1)'}), False, 'import torch\n'), (313, 'torch.softmax', 'torch.softmax', (['H'], {'dim': '(1)'}), False, 'import torch\n'), (92, 'torch.exp', 'torch.exp', (['self.beta'], {}), False, 'import torch\n'), (200, 'torch.exp', 'torch.exp', (['alpha'], {}), False, 'import torch\n'), (256, 'torch.ones', 'torch.ones', (['self.dim'], {}), False, 'import torch\n'), (258, 'torch.diag', 'torch.diag', (['self.S'], {}), False, 'import torch\n'), (261, 'torch.abs', 'torch.abs', (['self.S'], {}), False, 'import torch\n'), (53, 'torch.norm', 'torch.norm', (['self.w'], {}), False, 'import torch\n'), (52, 'torch.exp', 'torch.exp', (['(self.w @ self.u)'], {}), False, 'import torch\n'), (96, 'torch.exp', 'torch.exp', (['self.log_alpha'], {}), False, 'import torch\n')]