import torch
from torch._inductor.select_algorithm import extern_kernels
import triton
import triton.language as tl
from torch._inductor.runtime.triton_heuristics import grid
from torch._C import _cuda_getCurrentRawStream as get_raw_stream
from torch._inductor.runtime.triton_helpers import math as tl_math
import torch.nn as nn
import torch.autograd
import torch.utils.data
assert_size_stride = torch._C._dynamo.guards.assert_size_stride
empty_strided_cuda = torch._C._dynamo.guards._empty_strided_cuda
reinterpret_tensor = torch._C._dynamo.guards._reinterpret_tensor


@triton.jit
def triton_per_fused_abs_exp_neg_sub_sum_0(in_out_ptr0, in_ptr0, xnumel,
    rnumel, XBLOCK: tl.constexpr):
    xnumel = 64
    RBLOCK: tl.constexpr = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:, None]
    xmask = xindex < xnumel
    rindex = tl.arange(0, RBLOCK)[None, :]
    tl.full([XBLOCK, RBLOCK], True, tl.int1)
    r3 = rindex
    x5 = xindex % 16
    x0 = xindex % 4
    x2 = xindex // 16
    x4 = xindex
    tmp0 = tl.load(in_ptr0 + (r3 + 16 * x5), xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp1 = tl.load(in_ptr0 + (r3 + 16 * x0 + 64 * x2), xmask,
        eviction_policy='evict_last', other=0.0)
    tmp2 = tmp0 - tmp1
    tmp3 = tl_math.abs(tmp2)
    tmp4 = tl.broadcast_to(tmp3, [XBLOCK, RBLOCK])
    tmp6 = tl.where(xmask, tmp4, 0)
    tmp7 = tl.sum(tmp6, 1)[:, None]
    tmp8 = -tmp7
    tmp9 = tl_math.exp(tmp8)
    tl.debug_barrier()
    tl.store(in_out_ptr0 + x4, tmp9, xmask)


@triton.jit
def triton_poi_fused_cat_1(in_ptr0, in_ptr1, out_ptr0, xnumel, XBLOCK: tl.
    constexpr):
    xnumel = 32
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex % 8
    x1 = xindex // 8
    x2 = xindex
    tmp0 = x0
    tl.full([1], 0, tl.int64)
    tmp3 = tl.full([1], 4, tl.int64)
    tmp4 = tmp0 < tmp3
    tmp5 = tl.load(in_ptr0 + (4 * x1 + x0), tmp4 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp6 = tmp0 >= tmp3
    tl.full([1], 8, tl.int64)
    tmp9 = tl.load(in_ptr1 + (4 * x1 + (-4 + x0)), tmp6 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp10 = tl.load(in_ptr1 + (16 + 4 * x1 + (-4 + x0)), tmp6 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp11 = tmp9 + tmp10
    tmp12 = tl.load(in_ptr1 + (32 + 4 * x1 + (-4 + x0)), tmp6 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp13 = tmp11 + tmp12
    tmp14 = tl.load(in_ptr1 + (48 + 4 * x1 + (-4 + x0)), tmp6 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp15 = tmp13 + tmp14
    tmp16 = 1.0
    tmp17 = tmp15 - tmp16
    tmp18 = tl.full(tmp17.shape, 0.0, tmp17.dtype)
    tmp19 = tl.where(tmp6, tmp17, tmp18)
    tmp20 = tl.where(tmp4, tmp5, tmp19)
    tl.store(out_ptr0 + x2, tmp20, xmask)


def call(args):
    primals_1, primals_2 = args
    args.clear()
    assert_size_stride(primals_1, (4, 4, 16), (64, 16, 1))
    assert_size_stride(primals_2, (4, 4), (4, 1))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf0 = empty_strided_cuda((4, 64), (64, 1), torch.float32)
        extern_kernels.mm(primals_2, reinterpret_tensor(primals_1, (4, 64),
            (64, 1), 0), out=buf0)
        del primals_1
        buf1 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        buf2 = buf1
        del buf1
        get_raw_stream(0)
        triton_per_fused_abs_exp_neg_sub_sum_0[grid(64)](buf2, buf0, 64, 16,
            XBLOCK=8, num_warps=2, num_stages=1)
        buf3 = empty_strided_cuda((4, 8), (8, 1), torch.float32)
        triton_poi_fused_cat_1[grid(32)](primals_2, buf2, buf3, 32, XBLOCK=
            32, num_warps=1, num_stages=1)
    return buf3, buf0, buf2, reinterpret_tensor(primals_2, (4, 4), (1, 4), 0)


class MinibatchDiscrimination1dNew(nn.Module):
    """1D Minibatch Discrimination Module as proposed in the paper `"Improved Techniques for
    Training GANs by Salimans et. al." <https://arxiv.org/abs/1805.08318>`_

    Allows the Discriminator to easily detect mode collapse by augmenting the activations to the succeeding
    layer with side information that allows it to determine the 'closeness' of the minibatch examples
    with each other

    .. math :: M_i = T * f(x_{i})
    .. math :: c_b(x_{i}, x_{j}) = \\exp(-||M_{i, b} - M_{j, b}||_1) \\in \\mathbb{R}.
    .. math :: o(x_{i})_b &= \\sum_{j=1}^{n} c_b(x_{i},x_{j}) \\in \\mathbb{R} \\\\
    .. math :: o(x_{i}) &= \\Big[ o(x_{i})_1, o(x_{i})_2, \\dots, o(x_{i})_B \\Big] \\in \\mathbb{R}^B \\\\
    .. math :: o(X) \\in \\mathbb{R}^{n \\times B}

    This is followed by concatenating :math:`o(x_{i})` and :math:`f(x_{i})`

    where

    - :math:`f(x_{i}) \\in \\mathbb{R}^A` : Activations from an intermediate layer
    - :math:`f(x_{i}) \\in \\mathbb{R}^A` : Parameter Tensor for generating minibatch discrimination matrix


    Args:
        in_features (int): Features input corresponding to dimension :math:`A`
        out_features (int): Number of output features that are to be concatenated corresponding to dimension :math:`B`
        intermediate_features (int): Intermediate number of features corresponding to dimension :math:`C`

    Returns:
        A Tensor of size :math:`(N, in_features + out_features)` where :math:`N` is the batch size
    """

    def __init__(self, in_features, out_features, intermediate_features=16):
        super(MinibatchDiscrimination1dNew, self).__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.intermediate_features = intermediate_features
        self.T = nn.Parameter(torch.Tensor(in_features, out_features,
            intermediate_features))
        nn.init.normal_(self.T)

    def forward(self, input_0):
        primals_1 = self.T
        primals_2 = input_0
        output = call([primals_1, primals_2])
        return output[0]


Measured latency_ms: 0.014136