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, 12, 90, 35), 'torch.norm', 'torch.norm', ({(90, 23, 90, 34): 'x - self.x0'}, {}), '(x - self.x0)', False, 'import torch\n'), ((140, 12, 140, 44), 'torch.cat', 'torch.cat', (), '', False, 'import torch\n'), ((153, 12, 153, 44), 'torch.cat', 'torch.cat', (), '', False, 'import torch\n'), ((168, 22, 168, 37), 'torch.nn.ModuleList', 'nn.ModuleList', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((178, 12, 178, 31), 'torch.zeros_like', 'torch.zeros_like', ({(178, 29, 178, 30): 'x'}, {}), '(x)', False, 'import torch\n'), ((179, 18, 179, 41), 'torch.zeros', 'torch.zeros', ({(179, 30, 179, 40): 'z.shape[0]'}, {}), '(z.shape[0])', False, 'import torch\n'), ((191, 12, 191, 31), 'torch.zeros_like', 'torch.zeros_like', ({(191, 29, 191, 30): 'z'}, {}), '(z)', False, 'import torch\n'), ((192, 18, 192, 41), 'torch.zeros', 'torch.zeros', ({(192, 30, 192, 40): 'z.shape[0]'}, {}), '(z.shape[0])', False, 'import torch\n'), ((219, 18, 219, 43), 'torch.sum', 'torch.sum', ({(219, 28, 219, 42): 'self.log_sigma'}, {}), '(self.log_sigma)', False, 'import torch\n'), ((238, 18, 238, 33), 'scipy.linalg.lu', 'sp.linalg.lu', ({(238, 31, 238, 32): 'W'}, {}), '(W)', True, 'import scipy as sp\n'), ((239, 17, 239, 53), 'torch.tensor', 'torch.tensor', (), '', False, 'import torch\n'), ((248, 12, 248, 44), 'torch.triu', 'torch.triu', (), '', False, 'import torch\n'), ((276, 22, 276, 37), 'torch.nn.ModuleList', 'nn.ModuleList', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((283, 8, 283, 54), 'torch.nn.init.uniform_', 'init.uniform_', ({(283, 22, 283, 37): 'self.init_param', (283, 39, 283, 46): '(-1 / 2)', (283, 48, 283, 53): '(1 / 2)'}, {}), '(self.init_param, -1 / 2, 1 / 2)', True, 'import torch.nn.init as init\n'), ((286, 12, 286, 31), 'torch.zeros_like', 'torch.zeros_like', ({(286, 29, 286, 30): 'x'}, {}), '(x)', False, 'import torch\n'), ((287, 18, 287, 41), 'torch.zeros', 'torch.zeros', ({(287, 30, 287, 40): 'z.shape[0]'}, {}), '(z.shape[0])', False, 'import torch\n'), ((304, 12, 304, 31), 'torch.zeros_like', 'torch.zeros_like', ({(304, 29, 304, 30): 'z'}, {}), '(z)', False, 'import torch\n'), ((305, 18, 305, 41), 'torch.zeros', 'torch.zeros', ({(305, 30, 305, 40): 'x.shape[0]'}, {}), '(x.shape[0])', False, 'import torch\n'), ((337, 18, 337, 41), 'torch.zeros', 'torch.zeros', ({(337, 30, 337, 40): 'x.shape[0]'}, {}), '(x.shape[0])', False, 'import torch\n'), ((340, 18, 340, 51), 'torch.split', 'torch.split', (), '', False, 'import torch\n'), ((343, 12, 343, 25), 'torch.nn.functional.softplus', 'F.softplus', ({(343, 23, 343, 24): 'D'}, {}), '(D)', True, 'import torch.nn.functional as F\n'), ((344, 20, 345, 61), 'nf.utils.unconstrained_RQS', 'unconstrained_RQS', (), '', False, 'from nf.utils import unconstrained_RQS\n'), ((346, 19, 346, 41), 'torch.sum', 'torch.sum', (), '', False, 'import torch\n'), ((348, 18, 348, 51), 'torch.split', 'torch.split', (), '', False, 'import torch\n'), ((351, 12, 351, 25), 'torch.nn.functional.softplus', 'F.softplus', ({(351, 23, 351, 24): 'D'}, {}), '(D)', True, 'import torch.nn.functional as F\n'), ((352, 20, 353, 61), 'nf.utils.unconstrained_RQS', 'unconstrained_RQS', (), '', False, 'from nf.utils import unconstrained_RQS\n'), ((354, 19, 354, 41), 'torch.sum', 'torch.sum', (), '', False, 'import torch\n'), ((358, 18, 358, 41), 'torch.zeros', 'torch.zeros', ({(358, 30, 358, 40): 'z.shape[0]'}, {}), '(z.shape[0])', False, 'import torch\n'), ((361, 18, 361, 51), 'torch.split', 'torch.split', (), '', False, 'import torch\n'), ((364, 12, 364, 25), 'torch.nn.functional.softplus', 'F.softplus', ({(364, 23, 364, 24): 'D'}, {}), '(D)', True, 'import torch.nn.functional as F\n'), ((365, 20, 366, 60), 'nf.utils.unconstrained_RQS', 'unconstrained_RQS', (), '', False, 'from nf.utils import unconstrained_RQS\n'), ((367, 19, 367, 41), 'torch.sum', 'torch.sum', (), '', False, 'import torch\n'), ((369, 18, 369, 51), 'torch.split', 'torch.split', (), '', False, 'import torch\n'), ((372, 12, 372, 25), 'torch.nn.functional.softplus', 'F.softplus', ({(372, 23, 372, 24): 'D'}, {}), '(D)', True, 'import torch.nn.functional as F\n'), ((373, 20, 374, 64), 'nf.utils.unconstrained_RQS', 'unconstrained_RQS', (), '', False, 'from nf.utils import unconstrained_RQS\n'), ((375, 19, 375, 41), 'torch.sum', 'torch.sum', (), '', False, 'import torch\n'), ((32, 30, 32, 47), 'torch.Tensor', 'torch.Tensor', ({(32, 43, 32, 46): 'dim'}, {}), '(dim)', False, 'import torch\n'), ((33, 30, 33, 47), 'torch.Tensor', 'torch.Tensor', ({(33, 43, 33, 46): 'dim'}, {}), '(dim)', False, 'import torch\n'), ((34, 30, 34, 45), 'torch.Tensor', 'torch.Tensor', ({(34, 43, 34, 44): '1'}, {}), '(1)', False, 'import torch\n'), ((38, 49, 38, 65), 'math.sqrt', 'math.sqrt', ({(38, 59, 38, 64): '(1 / dim)'}, {}), '(1 / dim)', False, 'import math\n'), ((39, 49, 39, 65), 'math.sqrt', 'math.sqrt', ({(39, 59, 39, 64): '(1 / dim)'}, {}), '(1 / dim)', False, 'import math\n'), ((40, 49, 40, 65), 'math.sqrt', 'math.sqrt', ({(40, 59, 40, 64): '(1 / dim)'}, {}), '(1 / dim)', False, 'import math\n'), ((56, 14, 56, 44), 'torch.unsqueeze', 'torch.unsqueeze', ({(56, 30, 56, 40): '(x @ self.w)', (56, 42, 56, 43): '(1)'}, {}), '(x @ self.w, 1)', False, 'import torch\n'), ((76, 31, 76, 48), 'torch.Tensor', 'torch.Tensor', ({(76, 44, 76, 47): 'dim'}, {}), '(dim)', False, 'import torch\n'), ((77, 38, 77, 53), 'torch.Tensor', 'torch.Tensor', ({(77, 51, 77, 52): '1'}, {}), '(1)', False, 'import torch\n'), ((78, 33, 78, 48), 'torch.Tensor', 'torch.Tensor', ({(78, 46, 78, 47): '1'}, {}), '(1)', False, 'import torch\n'), ((81, 50, 81, 66), 'math.sqrt', 'math.sqrt', ({(81, 60, 81, 65): '(1 / dim)'}, {}), '(1 / dim)', False, 'import math\n'), ((82, 57, 82, 73), 'math.sqrt', 'math.sqrt', ({(82, 67, 82, 72): '(1 / dim)'}, {}), '(1 / dim)', False, 'import math\n'), ((83, 52, 83, 68), 'math.sqrt', 'math.sqrt', ({(83, 62, 83, 67): '(1 / dim)'}, {}), '(1 / dim)', False, 'import math\n'), ((107, 12, 107, 41), 'torch.nn.Linear', 'nn.Linear', ({(107, 22, 107, 28): 'in_dim', (107, 30, 107, 40): 'hidden_dim'}, {}), '(in_dim, hidden_dim)', True, 'import torch.nn as nn\n'), ((108, 12, 108, 21), 'torch.nn.Tanh', 'nn.Tanh', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((109, 12, 109, 45), 'torch.nn.Linear', 'nn.Linear', ({(109, 22, 109, 32): 'hidden_dim', (109, 34, 109, 44): 'hidden_dim'}, {}), '(hidden_dim, hidden_dim)', True, 'import torch.nn as nn\n'), ((110, 12, 110, 21), 'torch.nn.Tanh', 'nn.Tanh', ({}, {}), '()', True, 'import torch.nn as nn\n'), ((111, 12, 111, 42), 'torch.nn.Linear', 'nn.Linear', ({(111, 22, 111, 32): 'hidden_dim', (111, 34, 111, 41): 'out_dim'}, {}), '(hidden_dim, out_dim)', True, 'import torch.nn as nn\n'), ((141, 18, 141, 50), 'torch.sum', 'torch.sum', (), '', False, 'import torch\n'), ((142, 18, 142, 50), 'torch.sum', 'torch.sum', (), '', False, 'import torch\n'), ((149, 43, 149, 69), 'torch.exp', 'torch.exp', ({(149, 53, 149, 68): '(-s2_transformed)'}, {}), '(-s2_transformed)', False, 'import torch\n'), ((152, 43, 152, 69), 'torch.exp', 'torch.exp', ({(152, 53, 152, 68): '(-s1_transformed)'}, {}), '(-s1_transformed)', False, 'import torch\n'), ((154, 18, 154, 51), 'torch.sum', 'torch.sum', (), '', False, 'import torch\n'), ((155, 18, 155, 51), 'torch.sum', 'torch.sum', (), '', False, 'import torch\n'), ((169, 42, 169, 57), 'torch.Tensor', 'torch.Tensor', ({(169, 55, 169, 56): '2'}, {}), '(2)', False, 'import torch\n'), ((175, 59, 175, 73), 'math.sqrt', 'math.sqrt', ({(175, 69, 175, 72): '(0.5)'}, {}), '(0.5)', False, 'import math\n'), ((214, 31, 214, 68), 'torch.zeros', 'torch.zeros', (), '', False, 'import torch\n'), ((215, 38, 215, 75), 'torch.zeros', 'torch.zeros', (), '', False, 'import torch\n'), ((223, 28, 223, 53), 'torch.exp', 'torch.exp', ({(223, 38, 223, 52): 'self.log_sigma'}, {}), '(self.log_sigma)', False, 'import torch\n'), ((224, 19, 224, 44), 'torch.sum', 'torch.sum', ({(224, 29, 224, 43): 'self.log_sigma'}, {}), '(self.log_sigma)', False, 'import torch\n'), ((237, 28, 237, 53), 'numpy.random.randn', 'np.random.randn', ({(237, 44, 237, 47): 'dim', (237, 49, 237, 52): 'dim'}, {}), '(dim, dim)', True, 'import numpy as np\n'), ((240, 30, 240, 66), 'torch.tensor', 'torch.tensor', (), '', False, 'import torch\n'), ((247, 12, 247, 45), 'torch.tril', 'torch.tril', (), '', False, 'import torch\n'), ((257, 16, 257, 48), 'torch.triu', 'torch.triu', (), '', False, 'import torch\n'), ((259, 25, 259, 41), 'torch.inverse', 'torch.inverse', ({(259, 39, 259, 40): 'W'}, {}), '(W)', False, 'import torch\n'), ((277, 39, 277, 62), 'torch.Tensor', 'torch.Tensor', ({(277, 52, 277, 61): '3 * K - 1'}, {}), '(3 * K - 1)', False, 'import torch\n'), ((297, 16, 297, 29), 'torch.nn.functional.softplus', 'F.softplus', ({(297, 27, 297, 28): 'D'}, {}), '(D)', True, 'import torch.nn.functional as F\n'), ((298, 26, 299, 67), 'nf.utils.unconstrained_RQS', 'unconstrained_RQS', (), '', False, 'from nf.utils import unconstrained_RQS\n'), ((315, 16, 315, 29), 'torch.nn.functional.softplus', 'F.softplus', ({(315, 27, 315, 28): 'D'}, {}), '(D)', True, 'import torch.nn.functional as F\n'), ((316, 26, 317, 70), 'nf.utils.unconstrained_RQS', 'unconstrained_RQS', (), '', False, 'from nf.utils import unconstrained_RQS\n'), ((341, 15, 341, 40), 'torch.softmax', 'torch.softmax', (), '', False, 'import torch\n'), ((341, 42, 341, 67), 'torch.softmax', 'torch.softmax', (), '', False, 'import torch\n'), ((349, 15, 349, 40), 'torch.softmax', 'torch.softmax', (), '', False, 'import torch\n'), ((349, 42, 349, 67), 'torch.softmax', 'torch.softmax', (), '', False, 'import torch\n'), ((355, 15, 355, 49), 'torch.cat', 'torch.cat', (), '', False, 'import torch\n'), ((362, 15, 362, 40), 'torch.softmax', 'torch.softmax', (), '', False, 'import torch\n'), ((362, 42, 362, 67), 'torch.softmax', 'torch.softmax', (), '', False, 'import torch\n'), ((370, 15, 370, 40), 'torch.softmax', 'torch.softmax', (), '', False, 'import torch\n'), ((370, 42, 370, 67), 'torch.softmax', 'torch.softmax', (), '', False, 'import torch\n'), ((376, 15, 376, 49), 'torch.cat', 'torch.cat', (), '', False, 'import torch\n'), ((13, 40, 13, 53), 'torch.tanh', 'torch.tanh', ({(13, 51, 13, 52): 'x'}, {}), '(x)', False, 'import torch\n'), ((17, 55, 17, 67), 'torch.exp', 'torch.exp', ({(17, 65, 17, 66): 'x'}, {}), '(x)', False, 'import torch\n'), ((38, 31, 38, 47), 'math.sqrt', 'math.sqrt', ({(38, 41, 38, 46): '(1 / dim)'}, {}), '(1 / dim)', False, 'import math\n'), ((39, 31, 39, 47), 'math.sqrt', 'math.sqrt', ({(39, 41, 39, 46): '(1 / dim)'}, {}), '(1 / dim)', False, 'import math\n'), ((40, 31, 40, 47), 'math.sqrt', 'math.sqrt', ({(40, 41, 40, 46): '(1 / dim)'}, {}), '(1 / dim)', False, 'import math\n'), ((59, 28, 59, 50), 'torch.abs', 'torch.abs', ({(59, 38, 59, 49): '1 + phi @ u'}, {}), '(1 + phi @ u)', False, 'import torch\n'), ((81, 32, 81, 48), 'math.sqrt', 'math.sqrt', ({(81, 42, 81, 47): '(1 / dim)'}, {}), '(1 / dim)', False, 'import math\n'), ((82, 39, 82, 55), 'math.sqrt', 'math.sqrt', ({(82, 49, 82, 54): '(1 / dim)'}, {}), '(1 / dim)', False, 'import math\n'), ((83, 34, 83, 50), 'math.sqrt', 'math.sqrt', ({(83, 44, 83, 49): '(1 / dim)'}, {}), '(1 / dim)', False, 'import math\n'), ((91, 17, 91, 42), 'torch.exp', 'torch.exp', ({(91, 27, 91, 41): 'self.log_alpha'}, {}), '(self.log_alpha)', False, 'import torch\n'), ((92, 16, 92, 41), 'torch.exp', 'torch.exp', ({(92, 26, 92, 40): 'self.log_alpha'}, {}), '(self.log_alpha)', False, 'import torch\n'), ((94, 28, 94, 51), 'torch.log', 'torch.log', ({(94, 38, 94, 50): '(1 + beta * h)'}, {}), '(1 + beta * h)', False, 'import torch\n'), ((136, 41, 136, 66), 'torch.exp', 'torch.exp', ({(136, 51, 136, 65): 's1_transformed'}, {}), '(s1_transformed)', False, 'import torch\n'), ((139, 41, 139, 66), 'torch.exp', 'torch.exp', ({(139, 51, 139, 65): 's2_transformed'}, {}), '(s2_transformed)', False, 'import torch\n'), ((175, 43, 175, 57), 'math.sqrt', 'math.sqrt', ({(175, 53, 175, 56): '(0.5)'}, {}), '(0.5)', False, 'import math\n'), ((186, 39, 186, 55), 'torch.exp', 'torch.exp', ({(186, 49, 186, 54): 'alpha'}, {}), '(alpha)', False, 'import torch\n'), ((218, 16, 218, 41), 'torch.exp', 'torch.exp', ({(218, 26, 218, 40): 'self.log_sigma'}, {}), '(self.log_sigma)', False, 'import torch\n'), ((241, 43, 241, 53), 'numpy.diag', 'np.diag', ({(241, 51, 241, 52): 'U'}, {}), '(U)', True, 'import numpy as np\n'), ((242, 41, 242, 77), 'torch.tensor', 'torch.tensor', (), '', False, 'import torch\n'), ((247, 59, 247, 79), 'torch.ones', 'torch.ones', ({(247, 70, 247, 78): 'self.dim'}, {}), '(self.dim)', False, 'import torch\n'), ((249, 34, 249, 52), 'torch.diag', 'torch.diag', ({(249, 45, 249, 51): 'self.S'}, {}), '(self.S)', False, 'import torch\n'), ((250, 38, 250, 55), 'torch.abs', 'torch.abs', ({(250, 48, 250, 54): 'self.S'}, {}), '(self.S)', False, 'import torch\n'), ((255, 16, 255, 49), 'torch.tril', 'torch.tril', (), '', False, 'import torch\n'), ((291, 26, 291, 66), 'torch.split', 'torch.split', (), '', False, 'import torch\n'), ((294, 26, 294, 59), 'torch.split', 'torch.split', (), '', False, 'import torch\n'), ((295, 19, 295, 44), 'torch.softmax', 'torch.softmax', (), '', False, 'import torch\n'), ((295, 46, 295, 71), 'torch.softmax', 'torch.softmax', (), '', False, 'import torch\n'), ((309, 26, 309, 66), 'torch.split', 'torch.split', (), '', False, 'import torch\n'), ((312, 26, 312, 59), 'torch.split', 'torch.split', (), '', False, 'import torch\n'), ((313, 19, 313, 44), 'torch.softmax', 'torch.softmax', (), '', False, 'import torch\n'), ((313, 46, 313, 71), 'torch.softmax', 'torch.softmax', (), '', False, 'import torch\n'), ((92, 58, 92, 78), 'torch.exp', 'torch.exp', ({(92, 68, 92, 77): 'self.beta'}, {}), '(self.beta)', False, 'import torch\n'), ((200, 27, 200, 43), 'torch.exp', 'torch.exp', ({(200, 37, 200, 42): 'alpha'}, {}), '(alpha)', False, 'import torch\n'), ((256, 27, 256, 47), 'torch.ones', 'torch.ones', ({(256, 38, 256, 46): 'self.dim'}, {}), '(self.dim)', False, 'import torch\n'), ((258, 34, 258, 52), 'torch.diag', 'torch.diag', ({(258, 45, 258, 51): 'self.S'}, {}), '(self.S)', False, 'import torch\n'), ((261, 39, 261, 56), 'torch.abs', 'torch.abs', ({(261, 49, 261, 55): 'self.S'}, {}), '(self.S)', False, 'import torch\n'), ((53, 41, 53, 59), 'torch.norm', 'torch.norm', ({(53, 52, 53, 58): 'self.w'}, {}), '(self.w)', False, 'import torch\n'), ((52, 31, 52, 57), 'torch.exp', 'torch.exp', ({(52, 41, 52, 56): '(self.w @ self.u)'}, {}), '(self.w @ self.u)', False, 'import torch\n'), ((96, 40, 96, 65), 'torch.exp', 'torch.exp', ({(96, 50, 96, 64): 'self.log_alpha'}, {}), '(self.log_alpha)', False, 'import torch\n')]