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 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_div_linalg_vector_norm_mul_0(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
    x2 = xindex
    x1 = xindex // 4
    tmp0 = tl.load(in_ptr0 + x2, xmask)
    tmp1 = tl.load(in_ptr0 + 4 * x1, xmask, eviction_policy='evict_last')
    tmp3 = tl.load(in_ptr0 + (1 + 4 * x1), xmask, eviction_policy='evict_last')
    tmp6 = tl.load(in_ptr0 + (2 + 4 * x1), xmask, eviction_policy='evict_last')
    tmp9 = tl.load(in_ptr0 + (3 + 4 * x1), xmask, eviction_policy='evict_last')
    tmp2 = tmp1 * tmp1
    tmp4 = tmp3 * tmp3
    tmp5 = tmp2 + tmp4
    tmp7 = tmp6 * tmp6
    tmp8 = tmp5 + tmp7
    tmp10 = tmp9 * tmp9
    tmp11 = tmp8 + tmp10
    tmp12 = libdevice.sqrt(tmp11)
    tmp13 = tmp0 / tmp12
    tmp14 = 1.0
    tmp15 = tmp13 * tmp14
    tl.store(out_ptr0 + x2, tmp15, xmask)


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_div_linalg_vector_norm_mul_0[grid(256)](arg0_1,
            buf0, 256, XBLOCK=256, num_warps=4, num_stages=1)
        del arg0_1
    return buf0,


def _extra_repr(**kwargs):
    if 'n' in kwargs:
        ret = 'n={}'.format(kwargs['n'])
    elif 'dim' in kwargs:
        ret = 'dim={}'.format(kwargs['dim'])
    else:
        ret = ''
    if 'k' in kwargs:
        ret += ', k={}'.format(kwargs['k'])
    if 'rank' in kwargs:
        ret += ', rank={}'.format(kwargs['rank'])
    if 'radius' in kwargs:
        ret += ', radius={}'.format(kwargs['radius'])
    if 'lam' in kwargs:
        ret += ', lambda={}'.format(kwargs['lam'])
    if 'f' in kwargs:
        ret += ', f={}'.format(kwargs['f'].__name__)
    if 'tensorial_size' in kwargs:
        ts = kwargs['tensorial_size']
        if len(ts) != 0:
            ret += ', tensorial_size={}'.format(tuple(ts))
    if 'triv' in kwargs:
        ret += ', triv={}'.format(kwargs['triv'].__name__)
    if 'no_inv' in kwargs:
        if kwargs['no_inv']:
            ret += ', no inverse'
    if 'transposed' in kwargs:
        if kwargs['transposed']:
            ret += ', transposed'
    return ret


def _in_sphere(x, r, eps):
    norm = x.norm(dim=-1)
    rs = torch.full_like(norm, r)
    return (torch.norm(norm - rs, p=float('inf')) < eps).all()


def project(x):
    return x / x.norm(dim=-1, keepdim=True)


def uniform_init_sphere_(x, r=1.0):
    """Samples a point uniformly on the sphere into the tensor ``x``.
    If ``x`` has :math:`d > 1` dimensions, the first :math:`d-1` dimensions
    are treated as batch dimensions.
    """
    with torch.no_grad():
        x.normal_()
        x.data = r * project(x)
    return x


class InManifoldError(ValueError):

    def __init__(self, X, M):
        super().__init__('Tensor not contained in {}. Got\n{}'.format(M, X))


class SphereEmbeddedNew(nn.Module):

    def __init__(self, size, radius=1.0):
        """
        Sphere as the orthogonal projection from
        :math:`\\mathbb{R}^n` to :math:`\\mathbb{S}^{n-1}`, that is,
        :math:`x \\mapsto \\frac{x}{\\lVert x \\rVert}`.

        Args:
            size (torch.size): Size of the tensor to be parametrized
            radius (float): Optional.
                Radius of the sphere. A positive number. Default: ``1.``
        """
        super().__init__()
        self.n = size[-1]
        self.tensorial_size = size[:-1]
        self.radius = SphereEmbeddedNew.parse_radius(radius)

    @staticmethod
    def parse_radius(radius):
        if radius <= 0.0:
            raise ValueError(
                'The radius has to be a positive real number. Got {}'.
                format(radius))
        return radius

    def right_inverse(self, x, check_in_manifold=True):
        if check_in_manifold and not self.in_manifold(x):
            raise InManifoldError(x, self)
        return x / self.radius

    def in_manifold(self, x, eps=1e-05):
        """
        Checks that a vector is on the sphere.

        For tensors with more than 2 dimensions the first dimensions are
        treated as batch dimensions.

        Args:
            X (torch.Tensor): The vector to be checked.
            eps (float): Optional. Threshold at which the norm is considered
                    to be equal to ``1``. Default: ``1e-5``
        """
        return _in_sphere(x, self.radius, eps)

    def sample(self):
        """
        Returns a uniformly sampled vector on the sphere.
        """
        x = torch.empty(*(self.tensorial_size + (self.n,)))
        return uniform_init_sphere_(x, r=self.radius)

    def extra_repr(self):
        return _extra_repr(n=self.n, radius=self.radius, tensorial_size=
            self.tensorial_size)

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


Measured latency_ms: 0.01496