import torch
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 import triton_helpers
from torch._inductor.runtime.triton_helpers import libdevice, math as tl_math
import torch.nn as nn
import torch.utils.data
assert_size_stride = torch._C._dynamo.guards.assert_size_stride
empty_strided_cuda = torch._C._dynamo.guards._empty_strided_cuda


@triton.jit
def triton_poi_fused__to_copy_0(out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 64
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex % 4
    x2 = xindex
    tmp0 = x0
    tmp1 = tmp0.to(tl.float32)
    tmp2 = 0.3333333333333333
    tmp3 = tmp1 * tmp2
    tmp4 = 1.0
    tmp5 = triton_helpers.minimum(tmp3, tmp4)
    tmp6 = 4.0
    tmp7 = tmp5 * tmp6
    tmp8 = tmp7.to(tl.int32)
    tl.store(out_ptr0 + x2, tmp8, xmask)


@triton.jit
def triton_poi_fused_add_clamp_1(out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 64
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex % 4
    x2 = xindex
    tmp0 = x0
    tmp1 = tmp0.to(tl.float32)
    tmp2 = 0.3333333333333333
    tmp3 = tmp1 * tmp2
    tmp4 = 1.0
    tmp5 = triton_helpers.minimum(tmp3, tmp4)
    tmp6 = 4.0
    tmp7 = tmp5 * tmp6
    tmp8 = tmp7.to(tl.int32)
    tmp9 = tl.full([1], 1, tl.int64)
    tmp10 = tmp8 + tmp9
    tmp11 = tl.full([1], 4, tl.int64)
    tmp12 = triton_helpers.minimum(tmp10, tmp11)
    tl.store(out_ptr0 + x2, tmp12, xmask)


@triton.jit
def triton_per_fused_add_mul_rsub_sort_2(in_ptr0, out_ptr0, out_ptr1,
    xnumel, rnumel, XBLOCK: tl.constexpr):
    xnumel = 64
    RBLOCK: tl.constexpr = 4
    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)
    r2 = rindex
    x0 = xindex % 4
    x1 = xindex // 4
    tmp0 = tl.load(in_ptr0 + (x0 + 4 * r2 + 16 * x1), xmask, other=0.0)
    tmp1 = x0
    tmp2 = tmp1.to(tl.float32)
    tmp3 = 0.3333333333333333
    tmp4 = tmp2 * tmp3
    tmp5 = 1.0
    tmp6 = tmp4 <= tmp5
    tmp7 = tmp6.to(tl.float32)
    tmp8 = tmp5 - tmp7
    tmp9 = -99999.0
    tmp10 = tmp8 * tmp9
    tmp11 = tmp0 + tmp10
    tmp12 = r2
    tmp13 = tmp12.to(tl.int16)
    tmp14 = tl.broadcast_to(tmp11, [XBLOCK, RBLOCK])
    tmp15 = tl.broadcast_to(tmp13, [XBLOCK, RBLOCK])
    tmp16, tmp17 = triton_helpers.sort_with_index(tmp14, tmp15, None, 1,
        stable=False, descending=True)
    tl.store(out_ptr0 + (x0 + 4 * r2 + 16 * x1), tmp16, xmask)
    tl.store(out_ptr1 + (x0 + 4 * r2 + 16 * x1), tmp17, xmask)


@triton.jit
def triton_poi_fused_sort_3(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 256
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex
    tmp0 = tl.load(in_ptr0 + x0, xmask)
    tmp1 = tmp0.to(tl.int64)
    tl.store(out_ptr0 + x0, tmp1, xmask)


@triton.jit
def triton_poi_fused_add_frac_gather_mul_rsub_4(in_ptr0, in_ptr1, out_ptr0,
    xnumel, XBLOCK: tl.constexpr):
    xnumel = 256
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x3 = xindex
    x0 = xindex % 4
    x1 = xindex // 4 % 4
    tmp0 = tl.load(in_ptr0 + x3, xmask)
    tmp1 = x0
    tmp2 = tmp1.to(tl.float32)
    tmp3 = 0.3333333333333333
    tmp4 = tmp2 * tmp3
    tmp5 = 1.0
    tmp6 = triton_helpers.minimum(tmp4, tmp5)
    tmp7 = 4.0
    tmp8 = tmp6 * tmp7
    tmp9 = tl_math.abs(tmp8)
    tmp10 = libdevice.floor(tmp9)
    tmp11 = tl.full([1], 0, tl.int32)
    tmp12 = tmp11 < tmp8
    tmp13 = tmp12.to(tl.int8)
    tmp14 = tmp8 < tmp11
    tmp15 = tmp14.to(tl.int8)
    tmp16 = tmp13 - tmp15
    tmp17 = tmp16.to(tmp8.dtype)
    tmp18 = tmp10 * tmp17
    tmp19 = tmp8 - tmp18
    tmp20 = tmp5 - tmp19
    tmp21 = tmp8.to(tl.int32)
    tmp22 = tl.load(in_ptr1 + (tmp21 + 5 * x1), xmask, eviction_policy=
        'evict_last')
    tmp23 = tmp20 * tmp22
    tmp24 = tl.full([1], 1, tl.int64)
    tmp25 = tmp21 + tmp24
    tmp26 = tl.full([1], 4, tl.int64)
    tmp27 = triton_helpers.minimum(tmp25, tmp26)
    tmp28 = tl.load(in_ptr1 + (tmp27 + 5 * x1), xmask, eviction_policy=
        'evict_last')
    tmp29 = tmp19 * tmp28
    tmp30 = tmp23 + tmp29
    tmp31 = tmp0 * tmp30
    tmp32 = tmp4 <= tmp5
    tmp33 = tmp32.to(tl.float32)
    tmp34 = tmp31 * tmp33
    tl.store(out_ptr0 + x3, tmp34, xmask)


@triton.jit
def triton_poi_fused_sum_5(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 64
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex % 4
    x1 = xindex // 4
    x2 = xindex
    tmp0 = tl.load(in_ptr0 + (x0 + 16 * x1), xmask)
    tmp1 = tl.load(in_ptr0 + (4 + x0 + 16 * x1), xmask)
    tmp3 = tl.load(in_ptr0 + (8 + x0 + 16 * x1), xmask)
    tmp5 = tl.load(in_ptr0 + (12 + x0 + 16 * x1), xmask)
    tmp2 = tmp0 + tmp1
    tmp4 = tmp2 + tmp3
    tmp6 = tmp4 + tmp5
    tl.store(out_ptr0 + x2, tmp6, xmask)


def call(args):
    primals_1, primals_2 = args
    args.clear()
    assert_size_stride(primals_1, (4, 4, 4, 4), (64, 16, 4, 1))
    assert_size_stride(primals_2, (4, 5), (5, 1))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf0 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.int64)
        get_raw_stream(0)
        triton_poi_fused__to_copy_0[grid(64)](buf0, 64, XBLOCK=64,
            num_warps=1, num_stages=1)
        buf1 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.int64)
        triton_poi_fused_add_clamp_1[grid(64)](buf1, 64, XBLOCK=64,
            num_warps=1, num_stages=1)
        buf2 = empty_strided_cuda((4, 4, 4, 4), (64, 16, 4, 1), torch.float32)
        buf3 = empty_strided_cuda((4, 4, 4, 4), (64, 16, 4, 1), torch.int16)
        triton_per_fused_add_mul_rsub_sort_2[grid(64)](primals_1, buf2,
            buf3, 64, 4, XBLOCK=8, num_warps=2, num_stages=1)
        del primals_1
        buf4 = empty_strided_cuda((4, 4, 4, 4), (64, 16, 4, 1), torch.int64)
        triton_poi_fused_sort_3[grid(256)](buf3, buf4, 256, XBLOCK=256,
            num_warps=4, num_stages=1)
        del buf3
        buf5 = empty_strided_cuda((4, 4, 4, 4), (64, 16, 4, 1), torch.float32)
        triton_poi_fused_add_frac_gather_mul_rsub_4[grid(256)](buf2,
            primals_2, buf5, 256, XBLOCK=256, num_warps=4, num_stages=1)
        del primals_2
        buf6 = empty_strided_cuda((4, 4, 4), (16, 4, 1), torch.float32)
        triton_poi_fused_sum_5[grid(64)](buf5, buf6, 64, XBLOCK=64,
            num_warps=1, num_stages=1)
        del buf5
    return buf6, buf4, buf0, buf1, buf2


def deterministic_sort(s, tau):
    """
    "Stochastic Optimization of Sorting Networks via Continuous Relaxations" https://openreview.net/forum?id=H1eSS3CcKX

    Aditya Grover, Eric Wang, Aaron Zweig, Stefano Ermon

    s: input elements to be sorted. Shape: batch_size x n x 1
    tau: temperature for relaxation. Scalar.
    """
    n = s.size()[1]
    one = torch.ones((n, 1), dtype=torch.float32, device=s.device)
    A_s = torch.abs(s - s.permute(0, 2, 1))
    B = torch.matmul(A_s, torch.matmul(one, one.transpose(0, 1)))
    scaling = (n + 1 - 2 * (torch.arange(n, device=s.device) + 1)).type(torch
        .float32)
    C = torch.matmul(s, scaling.unsqueeze(0))
    P_max = (C - B).permute(0, 2, 1)
    sm = torch.nn.Softmax(-1)
    P_hat = sm(P_max / tau)
    return P_hat


def cont_sort(x, perm=None, temp=1):
    """ Helper function that calls deterministic_sort with the right shape.
    Since it assumes a shape of (batch_size, n, 1) while the input x is of shape (batch_size, channels, n),
    we can get this to the right shape by merging the first two dimensions.
    If an existing perm is passed in, we compute the "inverse" (transpose of perm) and just use that to unsort x.
    """
    original_size = x.size()
    x = x.view(-1, x.size(2), 1)
    if perm is None:
        perm = deterministic_sort(x, temp)
    else:
        perm = perm.transpose(1, 2)
    x = perm.matmul(x)
    x = x.view(original_size)
    return x, perm


def fill_sizes(sizes, x=None):
    """
        sizes is a LongTensor of size [batch_size], containing the set sizes.
        Each set size n is turned into [0/(n-1), 1/(n-1), ..., (n-2)/(n-1), 1, 0, 0, ..., 0, 0].
        These are the ratios r at which f is evaluated at.
        The 0s at the end are there for padding to the largest n in the batch.
        If the input set x is passed in, it guarantees that the mask is the correct size even when sizes.max()
        is less than x.size(), which can be a case if there is at least one padding element in each set in the batch.
    """
    if x is not None:
        max_size = x.size(2)
    else:
        max_size = sizes.max()
    size_tensor = sizes.new(sizes.size(0), max_size).float().fill_(-1)
    size_tensor = torch.arange(end=max_size, device=sizes.device, dtype=
        torch.float32)
    size_tensor = size_tensor.unsqueeze(0) / (sizes.float() - 1).clamp(min=1
        ).unsqueeze(1)
    mask = size_tensor <= 1
    mask = mask.unsqueeze(1)
    return size_tensor.clamp(max=1), mask.float()


class FSPoolNew(nn.Module):
    """
        Featurewise sort pooling. From:

        FSPool: Learning Set Representations with Featurewise Sort Pooling.
        Yan Zhang, Jonathon Hare, Adam Prügel-Bennett
        https://arxiv.org/abs/1906.02795
        https://github.com/Cyanogenoid/fspool
    """

    def __init__(self, in_channels, n_pieces, relaxed=False):
        """
        in_channels: Number of channels in input
        n_pieces: Number of pieces in piecewise linear
        relaxed: Use sorting networks relaxation instead of traditional sorting
        """
        super().__init__()
        self.n_pieces = n_pieces
        self.weight = nn.Parameter(torch.zeros(in_channels, n_pieces + 1))
        self.relaxed = relaxed
        self.reset_parameters()

    def reset_parameters(self):
        nn.init.normal_(self.weight)

    def forward_transpose(self, x, perm, n=None):
        """ FSUnpool 

        x: FloatTensor of shape (batch_size, in_channels)
        perm: Permutation matrix returned by forward function.
        n: LongTensor fo shape (batch_size)
        """
        if n is None:
            n = x.new(x.size(0)).fill_(perm.size(2)).long()
        sizes, mask = fill_sizes(n)
        mask = mask.expand(mask.size(0), x.size(1), mask.size(2))
        weight = self.determine_weight(sizes)
        x = x.unsqueeze(2) * weight * mask.float()
        if self.relaxed:
            x, _ = cont_sort(x, perm)
        else:
            x = x.scatter(2, perm, x)
        return x, mask

    def determine_weight(self, sizes):
        """
            Piecewise linear function. Evaluates f at the ratios in sizes.
            This should be a faster implementation than doing the sum over max terms, since we know that most terms in it are 0.
        """
        weight = self.weight.unsqueeze(0)
        weight = weight.expand(sizes.size(0), weight.size(1), weight.size(2))
        index = self.n_pieces * sizes
        index = index.unsqueeze(1)
        index = index.expand(index.size(0), weight.size(1), index.size(2))
        idx = index.long()
        frac = index.frac()
        left = weight.gather(2, idx)
        right = weight.gather(2, (idx + 1).clamp(max=self.n_pieces))
        return (1 - frac) * left + frac * right

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


Measured latency_ms: 0.078799