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.triton_helpers import libdevice
from typing import Tuple
from torch.nn.modules.loss import _Loss
from typing import List
from typing import Optional
assert_size_stride = torch._C._dynamo.guards.assert_size_stride
empty_strided_cuda = torch._C._dynamo.guards._empty_strided_cuda


@triton.jit
def triton_per_fused_pow_sub_sum_0(in_ptr0, out_ptr0, xnumel, rnumel,
    XBLOCK: tl.constexpr):
    xnumel = 4
    rnumel = 48
    RBLOCK: tl.constexpr = 64
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:, None]
    xmask = xindex < xnumel
    rindex = tl.arange(0, RBLOCK)[None, :]
    rmask = rindex < rnumel
    r1 = rindex % 12
    r2 = rindex // 12
    x0 = xindex
    tmp0 = tl.load(in_ptr0 + (4 + r1 + 16 * r2 + 64 * x0), rmask & xmask,
        other=0.0)
    tmp1 = tl.load(in_ptr0 + (r1 + 16 * r2 + 64 * x0), rmask & xmask, other=0.0
        )
    tmp2 = tmp0 - tmp1
    tmp3 = tmp2 * tmp2
    tmp4 = tl.broadcast_to(tmp3, [XBLOCK, RBLOCK])
    tmp6 = tl.where(rmask & xmask, tmp4, 0)
    tmp7 = tl.sum(tmp6, 1)[:, None]
    tl.store(out_ptr0 + x0, tmp7, xmask)


@triton.jit
def triton_per_fused_pow_sub_sum_1(in_ptr0, out_ptr0, xnumel, rnumel,
    XBLOCK: tl.constexpr):
    xnumel = 4
    rnumel = 48
    RBLOCK: tl.constexpr = 64
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:, None]
    xmask = xindex < xnumel
    rindex = tl.arange(0, RBLOCK)[None, :]
    rmask = rindex < rnumel
    r1 = rindex % 3
    r2 = rindex // 3
    x0 = xindex
    tmp0 = tl.load(in_ptr0 + (1 + r1 + 4 * r2 + 64 * x0), rmask & xmask,
        other=0.0)
    tmp1 = tl.load(in_ptr0 + (r1 + 4 * r2 + 64 * x0), rmask & xmask, other=0.0)
    tmp2 = tmp0 - tmp1
    tmp3 = tmp2 * tmp2
    tmp4 = tl.broadcast_to(tmp3, [XBLOCK, RBLOCK])
    tmp6 = tl.where(rmask & xmask, tmp4, 0)
    tmp7 = tl.sum(tmp6, 1)[:, None]
    tl.store(out_ptr0 + x0, tmp7, xmask)


@triton.jit
def triton_per_fused_add_mean_sqrt_2(in_out_ptr0, in_ptr0, in_ptr1, xnumel,
    rnumel, XBLOCK: tl.constexpr):
    RBLOCK: tl.constexpr = 4
    xoffset = tl.program_id(0) * XBLOCK
    xoffset + tl.arange(0, XBLOCK)[:, None]
    tl.full([XBLOCK, RBLOCK], True, tl.int1)
    rindex = tl.arange(0, RBLOCK)[None, :]
    tl.full([XBLOCK, RBLOCK], True, tl.int1)
    r0 = rindex
    tmp0 = tl.load(in_ptr0 + r0, None)
    tmp1 = tl.load(in_ptr1 + r0, None)
    tmp2 = tmp0 + tmp1
    tmp3 = libdevice.sqrt(tmp2)
    tmp4 = tl.broadcast_to(tmp3, [XBLOCK, RBLOCK])
    tmp6 = tl.sum(tmp4, 1)[:, None]
    tmp7 = 4.0
    tmp8 = tmp6 / tmp7
    tl.debug_barrier()
    tl.store(in_out_ptr0 + tl.full([XBLOCK, 1], 0, tl.int32), tmp8, None)


def call(args):
    arg0_1, = args
    args.clear()
    assert_size_stride(arg0_1, (4, 4, 4, 4), (64, 16, 4, 1))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf0 = empty_strided_cuda((4,), (1,), torch.float32)
        get_raw_stream(0)
        triton_per_fused_pow_sub_sum_0[grid(4)](arg0_1, buf0, 4, 48, XBLOCK
            =1, num_warps=2, num_stages=1)
        buf1 = empty_strided_cuda((4,), (1,), torch.float32)
        triton_per_fused_pow_sub_sum_1[grid(4)](arg0_1, buf1, 4, 48, XBLOCK
            =1, num_warps=2, num_stages=1)
        del arg0_1
        buf2 = empty_strided_cuda((), (), torch.float32)
        buf3 = buf2
        del buf2
        triton_per_fused_add_mean_sqrt_2[grid(1)](buf3, buf0, buf1, 1, 4,
            XBLOCK=1, num_warps=2, num_stages=1)
        del buf0
        del buf1
    return buf3,


def _reduce(x: 'torch.Tensor', reduction: 'str'='mean') ->torch.Tensor:
    """Reduce input in batch dimension if needed.

    Args:
        x: Tensor with shape (N, *).
        reduction: Specifies the reduction type:
            ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'mean'``
    """
    if reduction == 'none':
        return x
    elif reduction == 'mean':
        return x.mean(dim=0)
    elif reduction == 'sum':
        return x.sum(dim=0)
    else:
        raise ValueError(
            "Uknown reduction. Expected one of {'none', 'mean', 'sum'}")


def _validate_input(tensors: 'List[torch.Tensor]', dim_range:
    'Tuple[int, int]'=(0, -1), data_range: 'Tuple[float, float]'=(0.0, -1.0
    ), size_range: 'Optional[Tuple[int, int]]'=None) ->None:
    """Check that input(-s)  satisfies the requirements
    Args:
        tensors: Tensors to check
        dim_range: Allowed number of dimensions. (min, max)
        data_range: Allowed range of values in tensors. (min, max)
        size_range: Dimensions to include in size comparison. (start_dim, end_dim + 1)
    """
    if not __debug__:
        return
    x = tensors[0]
    for t in tensors:
        assert torch.is_tensor(t), f'Expected torch.Tensor, got {type(t)}'
        assert t.device == x.device, f'Expected tensors to be on {x.device}, got {t.device}'
        if size_range is None:
            assert t.size() == x.size(
                ), f'Expected tensors with same size, got {t.size()} and {x.size()}'
        else:
            assert t.size()[size_range[0]:size_range[1]] == x.size()[size_range
                [0]:size_range[1]
                ], f'Expected tensors with same size at given dimensions, got {t.size()} and {x.size()}'
        if dim_range[0] == dim_range[1]:
            assert t.dim() == dim_range[0
                ], f'Expected number of dimensions to be {dim_range[0]}, got {t.dim()}'
        elif dim_range[0] < dim_range[1]:
            assert dim_range[0] <= t.dim() <= dim_range[1
                ], f'Expected number of dimensions to be between {dim_range[0]} and {dim_range[1]}, got {t.dim()}'
        if data_range[0] < data_range[1]:
            assert data_range[0] <= t.min(
                ), f'Expected values to be greater or equal to {data_range[0]}, got {t.min()}'
            assert t.max() <= data_range[1
                ], f'Expected values to be lower or equal to {data_range[1]}, got {t.max()}'


def total_variation(x: 'torch.Tensor', reduction: 'str'='mean', norm_type:
    'str'='l2') ->torch.Tensor:
    """Compute Total Variation metric

    Args:
        x: Tensor. Shape :math:`(N, C, H, W)`.
        reduction: Specifies the reduction type:
            ``'none'`` | ``'mean'`` | ``'sum'``. Default:``'mean'``
        norm_type: ``'l1'`` | ``'l2'`` | ``'l2_squared'``,
            defines which type of norm to implement, isotropic  or anisotropic.

    Returns:
        Total variation of a given tensor

    References:
        https://www.wikiwand.com/en/Total_variation_denoising

        https://remi.flamary.com/demos/proxtv.html
    """
    _validate_input([x], dim_range=(4, 4), data_range=(0, -1))
    if norm_type == 'l1':
        w_variance = torch.sum(torch.abs(x[:, :, :, 1:] - x[:, :, :, :-1]),
            dim=[1, 2, 3])
        h_variance = torch.sum(torch.abs(x[:, :, 1:, :] - x[:, :, :-1, :]),
            dim=[1, 2, 3])
        score = h_variance + w_variance
    elif norm_type == 'l2':
        w_variance = torch.sum(torch.pow(x[:, :, :, 1:] - x[:, :, :, :-1], 
            2), dim=[1, 2, 3])
        h_variance = torch.sum(torch.pow(x[:, :, 1:, :] - x[:, :, :-1, :], 
            2), dim=[1, 2, 3])
        score = torch.sqrt(h_variance + w_variance)
    elif norm_type == 'l2_squared':
        w_variance = torch.sum(torch.pow(x[:, :, :, 1:] - x[:, :, :, :-1], 
            2), dim=[1, 2, 3])
        h_variance = torch.sum(torch.pow(x[:, :, 1:, :] - x[:, :, :-1, :], 
            2), dim=[1, 2, 3])
        score = h_variance + w_variance
    else:
        raise ValueError(
            "Incorrect norm type, should be one of {'l1', 'l2', 'l2_squared'}")
    return _reduce(score, reduction)


class TVLossNew(_Loss):
    """Creates a criterion that measures the total variation of the
    the given input :math:`x`.


    If :attr:`norm_type` set to ``'l2'`` the loss can be described as:

    .. math::
        TV(x) = \\sum_{N}\\sqrt{\\sum_{H, W, C}(|x_{:, :, i+1, j} - x_{:, :, i, j}|^2 +
        |x_{:, :, i, j+1} - x_{:, :, i, j}|^2)}

    Else if :attr:`norm_type` set to ``'l1'``:

    .. math::
        TV(x) = \\sum_{N}\\sum_{H, W, C}(|x_{:, :, i+1, j} - x_{:, :, i, j}| +
        |x_{:, :, i, j+1} - x_{:, :, i, j}|)

    where :math:`N` is the batch size, `C` is the channel size.

    Args:
        norm_type: one of ``'l1'`` | ``'l2'`` | ``'l2_squared'``
        reduction: Specifies the reduction type:
            ``'none'`` | ``'mean'`` | ``'sum'``. Default:``'mean'``

    Examples:

        >>> loss = TVLoss()
        >>> x = torch.rand(3, 3, 256, 256, requires_grad=True)
        >>> output = loss(x)
        >>> output.backward()

    References:
        https://www.wikiwand.com/en/Total_variation_denoising

        https://remi.flamary.com/demos/proxtv.html
    """

    def __init__(self, norm_type: 'str'='l2', reduction: 'str'='mean'):
        super().__init__()
        self.norm_type = norm_type
        self.reduction = reduction

    def forward(self, input_0):
        arg0_1 = input_0
        output = call([arg0_1])
        return output[0]


Measured latency_ms: 0.040228