import torch
from torch import nn


class NonSquareError(ValueError):

    def __init__(self, name, size):
        super().__init__(
            'The {} parametrization can just be applied to square matrices. Got a tensor of size {}'
            .format(name, size))


class VectorError(ValueError):

    def __init__(self, name, size):
        super().__init__(
            'Cannot instantiate {} on a tensor of less than 2 dimensions. Got a tensor of size {}'
            .format(name, size))


class Symmetric(nn.Module):

    def __init__(self, lower=True):
        """
        Vector space of symmetric matrices, parametrized in terms of the upper or lower
        triangular part of a matrix.

        Args:
            size (torch.size): Size of the tensor to be parametrized
            lower (bool): Optional. Uses the lower triangular part of the matrix to
                parametrize the matrix. Default: ``True``
        """
        super().__init__()
        self.lower = lower

    @staticmethod
    def frame(X, lower):
        if lower:
            return X.tril(0) + X.tril(-1).transpose(-2, -1)
        else:
            return X.triu(0) + X.triu(1).transpose(-2, -1)

    def forward(self, X):
        if len(X.size()) < 2:
            raise VectorError(type(self).__name__, X.size())
        if X.size(-2) != X.size(-1):
            raise NonSquareError(type(self).__name__, X.size())
        return self.frame(X, self.lower)

    @staticmethod
    def in_manifold(X, eps=1e-06):
        return X.dim() >= 2 and X.size(-2) == X.size(-1) and torch.allclose(X,
            X.transpose(-2, -1), atol=eps)


def get_inputs():
    return [torch.rand([4, 4, 4, 4])]


def get_init_inputs():
    return [[], {}]

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 import nn
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_add_tril_0(in_ptr0, out_ptr0, ynumel, xnumel, YBLOCK:
    tl.constexpr, XBLOCK: tl.constexpr):
    ynumel = 64
    xnumel = 4
    yoffset = tl.program_id(1) * YBLOCK
    yindex = yoffset + tl.arange(0, YBLOCK)[None, :]
    ymask = yindex < ynumel
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:, None]
    xmask = xindex < xnumel
    x2 = xindex
    y0 = yindex % 4
    y3 = yindex
    y1 = yindex // 4
    tmp3 = tl.load(in_ptr0 + (x2 + 4 * y3), xmask & ymask, eviction_policy=
        'evict_last')
    tmp9 = tl.load(in_ptr0 + (y0 + 4 * x2 + 16 * y1), xmask & ymask,
        eviction_policy='evict_last')
    tmp0 = x2 + -1 * y0
    tmp1 = tl.full([1, 1], 0, tl.int64)
    tmp2 = tmp0 <= tmp1
    tmp4 = 0.0
    tmp5 = tl.where(tmp2, tmp3, tmp4)
    tmp6 = y0 + -1 * x2
    tmp7 = tl.full([1, 1], -1, tl.int64)
    tmp8 = tmp6 <= tmp7
    tmp10 = tl.where(tmp8, tmp9, tmp4)
    tmp11 = tmp5 + tmp10
    tl.store(out_ptr0 + (x2 + 4 * y3), tmp11, xmask & ymask)


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, 4, 4, 4), (64, 16, 4, 1), torch.float32)
        get_raw_stream(0)
        triton_poi_fused_add_tril_0[grid(64, 4)](arg0_1, buf0, 64, 4,
            XBLOCK=4, YBLOCK=32, num_warps=4, num_stages=1)
        del arg0_1
    return buf0,


class NonSquareError(ValueError):

    def __init__(self, name, size):
        super().__init__(
            'The {} parametrization can just be applied to square matrices. Got a tensor of size {}'
            .format(name, size))


class VectorError(ValueError):

    def __init__(self, name, size):
        super().__init__(
            'Cannot instantiate {} on a tensor of less than 2 dimensions. Got a tensor of size {}'
            .format(name, size))


class SymmetricNew(nn.Module):

    def __init__(self, lower=True):
        """
        Vector space of symmetric matrices, parametrized in terms of the upper or lower
        triangular part of a matrix.

        Args:
            size (torch.size): Size of the tensor to be parametrized
            lower (bool): Optional. Uses the lower triangular part of the matrix to
                parametrize the matrix. Default: ``True``
        """
        super().__init__()
        self.lower = lower

    @staticmethod
    def frame(X, lower):
        if lower:
            return X.tril(0) + X.tril(-1).transpose(-2, -1)
        else:
            return X.triu(0) + X.triu(1).transpose(-2, -1)

    @staticmethod
    def in_manifold(X, eps=1e-06):
        return X.dim() >= 2 and X.size(-2) == X.size(-1) and torch.allclose(X,
            X.transpose(-2, -1), atol=eps)

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