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

[((2834, 2857), 'torch.norm', 'torch.norm', (['(x - self.x0)'], {}), '(x - self.x0)\n', (2844, 2857), False, 'import torch\n'), ((4545, 4577), 'torch.cat', 'torch.cat', (['[lower, upper]'], {'dim': '(1)'}), '([lower, upper], dim=1)\n', (4554, 4577), False, 'import torch\n'), ((5114, 5146), 'torch.cat', 'torch.cat', (['[lower, upper]'], {'dim': '(1)'}), '([lower, upper], dim=1)\n', (5123, 5146), False, 'import torch\n'), ((5522, 5537), 'torch.nn.ModuleList', 'nn.ModuleList', ([], {}), '()\n', (5535, 5537), True, 'import torch.nn as nn\n'), ((5868, 5887), 'torch.zeros_like', 'torch.zeros_like', (['x'], {}), '(x)\n', (5884, 5887), False, 'import torch\n'), ((5906, 5929), 'torch.zeros', 'torch.zeros', (['z.shape[0]'], {}), '(z.shape[0])\n', (5917, 5929), False, 'import torch\n'), ((6345, 6364), 'torch.zeros_like', 'torch.zeros_like', (['z'], {}), '(z)\n', (6361, 6364), False, 'import torch\n'), ((6383, 6406), 'torch.zeros', 'torch.zeros', (['z.shape[0]'], {}), '(z.shape[0])\n', (6394, 6406), False, 'import torch\n'), ((7214, 7239), 'torch.sum', 'torch.sum', (['self.log_sigma'], {}), '(self.log_sigma)\n', (7223, 7239), False, 'import torch\n'), ((7686, 7701), 'scipy.linalg.lu', 'sp.linalg.lu', (['W'], {}), '(W)\n', (7698, 7701), True, 'import scipy as sp\n'), ((7719, 7753), 'torch.tensor', 'torch.tensor', (['P'], {'dtype': 'torch.float'}), '(P, dtype=torch.float)\n', (7731, 7753), False, 'import torch\n'), ((8171, 8201), 'torch.triu', 'torch.triu', (['self.U'], {'diagonal': '(1)'}), '(self.U, diagonal=1)\n', (8181, 8201), False, 'import torch\n'), ((9057, 9072), 'torch.nn.ModuleList', 'nn.ModuleList', ([], {}), '()\n', (9070, 9072), True, 'import torch.nn as nn\n'), ((9310, 9355), 'torch.nn.init.uniform_', 'init.uniform_', (['self.init_param', '(-1 / 2)', '(1 / 2)'], {}), '(self.init_param, -1 / 2, 1 / 2)\n', (9323, 9355), True, 'import torch.nn.init as init\n'), ((9396, 9415), 'torch.zeros_like', 'torch.zeros_like', (['x'], {}), '(x)\n', (9412, 9415), False, 'import torch\n'), ((9434, 9457), 'torch.zeros', 'torch.zeros', (['z.shape[0]'], {}), '(z.shape[0])\n', (9445, 9457), False, 'import torch\n'), ((10148, 10167), 'torch.zeros_like', 'torch.zeros_like', (['z'], {}), '(z)\n', (10164, 10167), False, 'import torch\n'), ((10186, 10209), 'torch.zeros', 'torch.zeros', (['x.shape[0]'], {}), '(x.shape[0])\n', (10197, 10209), False, 'import torch\n'), ((11339, 11362), 'torch.zeros', 'torch.zeros', (['x.shape[0]'], {}), '(x.shape[0])\n', (11350, 11362), False, 'import torch\n'), ((11519, 11550), 'torch.split', 'torch.split', (['out', 'self.K'], {'dim': '(2)'}), '(out, self.K, dim=2)\n', (11530, 11550), False, 'import torch\n'), ((11679, 11692), 'torch.nn.functional.softplus', 'F.softplus', (['D'], {}), '(D)\n', (11689, 11692), True, 'import torch.nn.functional as F\n'), ((11713, 11780), 'nf.utils.unconstrained_RQS', 'unconstrained_RQS', (['upper', 'W', 'H', 'D'], {'inverse': '(False)', 'tail_bound': 'self.B'}), '(upper, W, H, D, inverse=False, tail_bound=self.B)\n', (11730, 11780), False, 'from nf.utils import unconstrained_RQS\n'), ((11813, 11833), 'torch.sum', 'torch.sum', (['ld'], {'dim': '(1)'}), '(ld, dim=1)\n', (11822, 11833), False, 'import torch\n'), ((11926, 11957), 'torch.split', 'torch.split', (['out', 'self.K'], {'dim': '(2)'}), '(out, self.K, dim=2)\n', (11937, 11957), False, 'import torch\n'), ((12086, 12099), 'torch.nn.functional.softplus', 'F.softplus', (['D'], {}), '(D)\n', (12096, 12099), True, 'import torch.nn.functional as F\n'), ((12120, 12187), 'nf.utils.unconstrained_RQS', 'unconstrained_RQS', (['lower', 'W', 'H', 'D'], {'inverse': '(False)', 'tail_bound': 'self.B'}), '(lower, W, H, D, inverse=False, tail_bound=self.B)\n', (12137, 12187), False, 'from nf.utils import unconstrained_RQS\n'), ((12220, 12240), 'torch.sum', 'torch.sum', (['ld'], {'dim': '(1)'}), '(ld, dim=1)\n', (12229, 12240), False, 'import torch\n'), ((12348, 12371), 'torch.zeros', 'torch.zeros', (['z.shape[0]'], {}), '(z.shape[0])\n', (12359, 12371), False, 'import torch\n'), ((12528, 12559), 'torch.split', 'torch.split', (['out', 'self.K'], {'dim': '(2)'}), '(out, self.K, dim=2)\n', (12539, 12559), False, 'import torch\n'), ((12688, 12701), 'torch.nn.functional.softplus', 'F.softplus', (['D'], {}), '(D)\n', (12698, 12701), True, 'import torch.nn.functional as F\n'), ((12722, 12788), 'nf.utils.unconstrained_RQS', 'unconstrained_RQS', (['lower', 'W', 'H', 'D'], {'inverse': '(True)', 'tail_bound': 'self.B'}), '(lower, W, H, D, inverse=True, tail_bound=self.B)\n', (12739, 12788), False, 'from nf.utils import unconstrained_RQS\n'), ((12821, 12841), 'torch.sum', 'torch.sum', (['ld'], {'dim': '(1)'}), '(ld, dim=1)\n', (12830, 12841), False, 'import torch\n'), ((12934, 12965), 'torch.split', 'torch.split', (['out', 'self.K'], {'dim': '(2)'}), '(out, self.K, dim=2)\n', (12945, 12965), False, 'import torch\n'), ((13094, 13107), 'torch.nn.functional.softplus', 'F.softplus', (['D'], {}), '(D)\n', (13104, 13107), True, 'import torch.nn.functional as F\n'), ((13128, 13194), 'nf.utils.unconstrained_RQS', 'unconstrained_RQS', (['upper', 'W', 'H', 'D'], {'inverse': '(True)', 'tail_bound': 'self.B'}), '(upper, W, H, D, inverse=True, tail_bound=self.B)\n', (13145, 13194), False, 'from nf.utils import unconstrained_RQS\n'), ((13231, 13251), 'torch.sum', 'torch.sum', (['ld'], {'dim': '(1)'}), '(ld, dim=1)\n', (13240, 13251), False, 'import torch\n'), ((894, 911), 'torch.Tensor', 'torch.Tensor', (['dim'], {}), '(dim)\n', (906, 911), False, 'import torch\n'), ((943, 960), 'torch.Tensor', 'torch.Tensor', (['dim'], {}), '(dim)\n', (955, 960), False, 'import torch\n'), ((992, 1007), 'torch.Tensor', 'torch.Tensor', (['(1)'], {}), '(1)\n', (1004, 1007), False, 'import torch\n'), ((1131, 1149), 'math.sqrt', 'math.sqrt', (['(1 / dim)'], {}), '(1 / dim)\n', (1140, 1149), False, 'import math\n'), ((1198, 1216), 'math.sqrt', 'math.sqrt', (['(1 / dim)'], {}), '(1 / dim)\n', (1207, 1216), False, 'import math\n'), ((1265, 1283), 'math.sqrt', 'math.sqrt', (['(1 / dim)'], {}), '(1 / dim)\n', (1274, 1283), False, 'import math\n'), ((1774, 1804), 'torch.unsqueeze', 'torch.unsqueeze', (['(x @ self.w)', '(1)'], {}), '(x @ self.w, 1)\n', (1789, 1804), False, 'import torch\n'), ((2316, 2333), 'torch.Tensor', 'torch.Tensor', (['dim'], {}), '(dim)\n', (2328, 2333), False, 'import torch\n'), ((2373, 2388), 'torch.Tensor', 'torch.Tensor', (['(1)'], {}), '(1)\n', (2385, 2388), False, 'import torch\n'), ((2423, 2438), 'torch.Tensor', 'torch.Tensor', (['(1)'], {}), '(1)\n', (2435, 2438), False, 'import torch\n'), ((2522, 2540), 'math.sqrt', 'math.sqrt', (['(1 / dim)'], {}), '(1 / dim)\n', (2531, 2540), False, 'import math\n'), ((2597, 2615), 'math.sqrt', 'math.sqrt', (['(1 / dim)'], {}), '(1 / dim)\n', (2606, 2615), False, 'import math\n'), ((2667, 2685), 'math.sqrt', 'math.sqrt', (['(1 / dim)'], {}), '(1 / dim)\n', (2676, 2685), False, 'import math\n'), ((3445, 3474), 'torch.nn.Linear', 'nn.Linear', (['in_dim', 'hidden_dim'], {}), '(in_dim, hidden_dim)\n', (3454, 3474), True, 'import torch.nn as nn\n'), ((3488, 3497), 'torch.nn.Tanh', 'nn.Tanh', ([], {}), '()\n', (3495, 3497), True, 'import torch.nn as nn\n'), ((3511, 3544), 'torch.nn.Linear', 'nn.Linear', (['hidden_dim', 'hidden_dim'], {}), '(hidden_dim, hidden_dim)\n', (3520, 3544), True, 'import torch.nn as nn\n'), ((3558, 3567), 'torch.nn.Tanh', 'nn.Tanh', ([], {}), '()\n', (3565, 3567), True, 'import torch.nn as nn\n'), ((3581, 3611), 'torch.nn.Linear', 'nn.Linear', (['hidden_dim', 'out_dim'], {}), '(hidden_dim, out_dim)\n', (3590, 3611), True, 'import torch.nn as nn\n'), ((4596, 4628), 'torch.sum', 'torch.sum', (['s1_transformed'], {'dim': '(1)'}), '(s1_transformed, dim=1)\n', (4605, 4628), False, 'import torch\n'), ((4651, 4683), 'torch.sum', 'torch.sum', (['s2_transformed'], {'dim': '(1)'}), '(s2_transformed, dim=1)\n', (4660, 4683), False, 'import torch\n'), ((4925, 4951), 'torch.exp', 'torch.exp', (['(-s2_transformed)'], {}), '(-s2_transformed)\n', (4934, 4951), False, 'import torch\n'), ((5075, 5101), 'torch.exp', 'torch.exp', (['(-s1_transformed)'], {}), '(-s1_transformed)\n', (5084, 5101), False, 'import torch\n'), ((5165, 5198), 'torch.sum', 'torch.sum', (['(-s1_transformed)'], {'dim': '(1)'}), '(-s1_transformed, dim=1)\n', (5174, 5198), False, 'import torch\n'), ((5221, 5254), 'torch.sum', 'torch.sum', (['(-s2_transformed)'], {'dim': '(1)'}), '(-s2_transformed, dim=1)\n', (5230, 5254), False, 'import torch\n'), ((5580, 5595), 'torch.Tensor', 'torch.Tensor', (['(2)'], {}), '(2)\n', (5592, 5595), False, 'import torch\n'), ((5813, 5827), 'math.sqrt', 'math.sqrt', (['(0.5)'], {}), '(0.5)\n', (5822, 5827), False, 'import math\n'), ((7001, 7036), 'torch.zeros', 'torch.zeros', (['dim'], {'dtype': 'torch.float'}), '(dim, dtype=torch.float)\n', (7012, 7036), False, 'import torch\n'), ((7078, 7113), 'torch.zeros', 'torch.zeros', (['dim'], {'dtype': 'torch.float'}), '(dim, dtype=torch.float)\n', (7089, 7113), False, 'import torch\n'), ((7322, 7347), 'torch.exp', 'torch.exp', (['self.log_sigma'], {}), '(self.log_sigma)\n', (7331, 7347), False, 'import torch\n'), ((7367, 7392), 'torch.sum', 'torch.sum', (['self.log_sigma'], {}), '(self.log_sigma)\n', (7376, 7392), False, 'import torch\n'), ((7641, 7666), 'numpy.random.randn', 'np.random.randn', (['dim', 'dim'], {}), '(dim, dim)\n', (7656, 7666), True, 'import numpy as np\n'), ((7786, 7820), 'torch.tensor', 'torch.tensor', (['L'], {'dtype': 'torch.float'}), '(L, dtype=torch.float)\n', (7798, 7820), False, 'import torch\n'), ((8090, 8121), 'torch.tril', 'torch.tril', (['self.L'], {'diagonal': '(-1)'}), '(self.L, diagonal=-1)\n', (8100, 8121), False, 'import torch\n'), ((8516, 8546), 'torch.triu', 'torch.triu', (['self.U'], {'diagonal': '(1)'}), '(self.U, diagonal=1)\n', (8526, 8546), False, 'import torch\n'), ((8628, 8644), 'torch.inverse', 'torch.inverse', (['W'], {}), '(W)\n', (8641, 8644), False, 'import torch\n'), ((9112, 9135), 'torch.Tensor', 'torch.Tensor', (['(3 * K - 1)'], {}), '(3 * K - 1)\n', (9124, 9135), False, 'import torch\n'), ((9929, 9942), 'torch.nn.functional.softplus', 'F.softplus', (['D'], {}), '(D)\n', (9939, 9942), True, 'import torch.nn.functional as F\n'), ((9969, 10040), 'nf.utils.unconstrained_RQS', 'unconstrained_RQS', (['x[:, (i)]', 'W', 'H', 'D'], {'inverse': '(False)', 'tail_bound': 'self.B'}), '(x[:, (i)], W, H, D, inverse=False, tail_bound=self.B)\n', (9986, 10040), False, 'from nf.utils import unconstrained_RQS\n'), ((10681, 10694), 'torch.nn.functional.softplus', 'F.softplus', (['D'], {}), '(D)\n', (10691, 10694), True, 'import torch.nn.functional as F\n'), ((10721, 10791), 'nf.utils.unconstrained_RQS', 'unconstrained_RQS', (['z[:, (i)]', 'W', 'H', 'D'], {'inverse': '(True)', 'tail_bound': 'self.B'}), '(z[:, (i)], W, H, D, inverse=True, tail_bound=self.B)\n', (10738, 10791), False, 'from nf.utils import unconstrained_RQS\n'), ((11568, 11591), 'torch.softmax', 'torch.softmax', (['W'], {'dim': '(2)'}), '(W, dim=2)\n', (11581, 11591), False, 'import torch\n'), ((11595, 11618), 'torch.softmax', 'torch.softmax', (['H'], {'dim': '(2)'}), '(H, dim=2)\n', (11608, 11618), False, 'import torch\n'), ((11975, 11998), 'torch.softmax', 'torch.softmax', (['W'], {'dim': '(2)'}), '(W, dim=2)\n', (11988, 11998), False, 'import torch\n'), ((12002, 12025), 'torch.softmax', 'torch.softmax', (['H'], {'dim': '(2)'}), '(H, dim=2)\n', (12015, 12025), False, 'import torch\n'), ((12258, 12290), 'torch.cat', 'torch.cat', (['[lower, upper]'], {'dim': '(1)'}), '([lower, upper], dim=1)\n', (12267, 12290), False, 'import torch\n'), ((12577, 12600), 'torch.softmax', 'torch.softmax', (['W'], {'dim': '(2)'}), '(W, dim=2)\n', (12590, 12600), False, 'import torch\n'), ((12604, 12627), 'torch.softmax', 'torch.softmax', (['H'], {'dim': '(2)'}), '(H, dim=2)\n', (12617, 12627), False, 'import torch\n'), ((12983, 13006), 'torch.softmax', 'torch.softmax', (['W'], {'dim': '(2)'}), '(W, dim=2)\n', (12996, 13006), False, 'import torch\n'), ((13010, 13033), 'torch.softmax', 'torch.softmax', (['H'], {'dim': '(2)'}), '(H, dim=2)\n', (13023, 13033), False, 'import torch\n'), ((13269, 13301), 'torch.cat', 'torch.cat', (['[lower, upper]'], {'dim': '(1)'}), '([lower, upper], dim=1)\n', (13278, 13301), False, 'import torch\n'), ((344, 357), 'torch.tanh', 'torch.tanh', (['x'], {}), '(x)\n', (354, 357), False, 'import torch\n'), ((608, 620), 'torch.exp', 'torch.exp', (['x'], {}), '(x)\n', (617, 620), False, 'import torch\n'), ((1113, 1131), 'math.sqrt', 'math.sqrt', (['(1 / dim)'], {}), '(1 / dim)\n', (1122, 1131), False, 'import math\n'), ((1180, 1198), 'math.sqrt', 'math.sqrt', (['(1 / dim)'], {}), '(1 / dim)\n', (1189, 1198), False, 'import math\n'), ((1247, 1265), 'math.sqrt', 'math.sqrt', (['(1 / dim)'], {}), '(1 / dim)\n', (1256, 1265), False, 'import math\n'), ((1933, 1955), 'torch.abs', 'torch.abs', (['(1 + phi @ u)'], {}), '(1 + phi @ u)\n', (1942, 1955), False, 'import torch\n'), ((2504, 2522), 'math.sqrt', 'math.sqrt', (['(1 / dim)'], {}), '(1 / dim)\n', (2513, 2522), False, 'import math\n'), ((2579, 2597), 'math.sqrt', 'math.sqrt', (['(1 / dim)'], {}), '(1 / dim)\n', (2588, 2597), False, 'import math\n'), ((2649, 2667), 'math.sqrt', 'math.sqrt', (['(1 / dim)'], {}), '(1 / dim)\n', (2658, 2667), False, 'import math\n'), ((2875, 2900), 'torch.exp', 'torch.exp', (['self.log_alpha'], {}), '(self.log_alpha)\n', (2884, 2900), False, 'import torch\n'), ((2922, 2947), 'torch.exp', 'torch.exp', (['self.log_alpha'], {}), '(self.log_alpha)\n', (2931, 2947), False, 'import torch\n'), ((3055, 3078), 'torch.log', 'torch.log', (['(1 + beta * h)'], {}), '(1 + beta * h)\n', (3064, 3078), False, 'import torch\n'), ((4360, 4385), 'torch.exp', 'torch.exp', (['s1_transformed'], {}), '(s1_transformed)\n', (4369, 4385), False, 'import torch\n'), ((4507, 4532), 'torch.exp', 'torch.exp', (['s2_transformed'], {}), '(s2_transformed)\n', (4516, 4532), False, 'import torch\n'), ((5797, 5811), 'math.sqrt', 'math.sqrt', (['(0.5)'], {}), '(0.5)\n', (5806, 5811), False, 'import math\n'), ((6217, 6233), 'torch.exp', 'torch.exp', (['alpha'], {}), '(alpha)\n', (6226, 6233), False, 'import torch\n'), ((7160, 7185), 'torch.exp', 'torch.exp', (['self.log_sigma'], {}), '(self.log_sigma)\n', (7169, 7185), False, 'import torch\n'), ((7867, 7877), 'numpy.diag', 'np.diag', (['U'], {}), '(U)\n', (7874, 7877), True, 'import numpy as np\n'), ((7942, 7976), 'torch.tensor', 'torch.tensor', (['U'], {'dtype': 'torch.float'}), '(U, dtype=torch.float)\n', (7954, 7976), False, 'import torch\n'), ((8137, 8157), 'torch.ones', 'torch.ones', (['self.dim'], {}), '(self.dim)\n', (8147, 8157), False, 'import torch\n'), ((8238, 8256), 'torch.diag', 'torch.diag', (['self.S'], {}), '(self.S)\n', (8248, 8256), False, 'import torch\n'), ((8296, 8313), 'torch.abs', 'torch.abs', (['self.S'], {}), '(self.S)\n', (8305, 8313), False, 'import torch\n'), ((8413, 8444), 'torch.tril', 'torch.tril', (['self.L'], {'diagonal': '(-1)'}), '(self.L, diagonal=-1)\n', (8423, 8444), False, 'import torch\n'), ((9621, 9659), 'torch.split', 'torch.split', (['init_param', 'self.K'], {'dim': '(1)'}), '(init_param, self.K, dim=1)\n', (9632, 9659), False, 'import torch\n'), ((9757, 9788), 'torch.split', 'torch.split', (['out', 'self.K'], {'dim': '(1)'}), '(out, self.K, dim=1)\n', (9768, 9788), False, 'import torch\n'), ((9810, 9833), 'torch.softmax', 'torch.softmax', (['W'], {'dim': '(1)'}), '(W, dim=1)\n', (9823, 9833), False, 'import torch\n'), ((9837, 9860), 'torch.softmax', 'torch.softmax', (['H'], {'dim': '(1)'}), '(H, dim=1)\n', (9850, 9860), False, 'import torch\n'), ((10373, 10411), 'torch.split', 'torch.split', (['init_param', 'self.K'], {'dim': '(1)'}), '(init_param, self.K, dim=1)\n', (10384, 10411), False, 'import torch\n'), ((10509, 10540), 'torch.split', 'torch.split', (['out', 'self.K'], {'dim': '(1)'}), '(out, self.K, dim=1)\n', (10520, 10540), False, 'import torch\n'), ((10562, 10585), 'torch.softmax', 'torch.softmax', (['W'], {'dim': '(1)'}), '(W, dim=1)\n', (10575, 10585), False, 'import torch\n'), ((10589, 10612), 'torch.softmax', 'torch.softmax', (['H'], {'dim': '(1)'}), '(H, dim=1)\n', (10602, 10612), False, 'import torch\n'), ((2964, 2984), 'torch.exp', 'torch.exp', (['self.beta'], {}), '(self.beta)\n', (2973, 2984), False, 'import torch\n'), ((6712, 6728), 'torch.exp', 'torch.exp', (['alpha'], {}), '(alpha)\n', (6721, 6728), False, 'import torch\n'), ((8478, 8498), 'torch.ones', 'torch.ones', (['self.dim'], {}), '(self.dim)\n', (8488, 8498), False, 'import torch\n'), ((8583, 8601), 'torch.diag', 'torch.diag', (['self.S'], {}), '(self.S)\n', (8593, 8601), False, 'import torch\n'), ((8711, 8728), 'torch.abs', 'torch.abs', (['self.S'], {}), '(self.S)\n', (8720, 8728), False, 'import torch\n'), ((1654, 1672), 'torch.norm', 'torch.norm', (['self.w'], {}), '(self.w)\n', (1664, 1672), False, 'import torch\n'), ((1563, 1589), 'torch.exp', 'torch.exp', (['(self.w @ self.u)'], {}), '(self.w @ self.u)\n', (1572, 1589), False, 'import torch\n'), ((3168, 3193), 'torch.exp', 'torch.exp', (['self.log_alpha'], {}), '(self.log_alpha)\n', (3177, 3193), False, 'import torch\n')]