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 import triton_helpers
from torch._inductor.runtime.triton_helpers import libdevice
from torch import nn as nn
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_poi_fused_clamp_div_linalg_vector_norm_mul_sum_0(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
    tmp0 = tl.load(in_ptr0 + 6 * x0, xmask, eviction_policy='evict_last')
    tmp3 = tl.load(in_ptr0 + (1 + 6 * x0), xmask, eviction_policy='evict_last')
    tmp7 = tl.load(in_ptr0 + (2 + 6 * x0), xmask, eviction_policy='evict_last')
    tmp15 = tl.load(in_ptr0 + (3 + 6 * x0), xmask, eviction_policy='evict_last'
        )
    tmp19 = tl.load(in_ptr0 + (4 + 6 * x0), xmask, eviction_policy='evict_last'
        )
    tmp24 = tl.load(in_ptr0 + (5 + 6 * x0), xmask, eviction_policy='evict_last'
        )
    tmp1 = tmp0.to(tl.float64)
    tmp2 = tmp1 * tmp1
    tmp4 = tmp3.to(tl.float64)
    tmp5 = tmp4 * tmp4
    tmp6 = tmp2 + tmp5
    tmp8 = tmp7.to(tl.float64)
    tmp9 = tmp8 * tmp8
    tmp10 = tmp6 + tmp9
    tmp11 = libdevice.sqrt(tmp10)
    tmp12 = tl.full([1], 1e-05, tl.float64)
    tmp13 = triton_helpers.maximum(tmp11, tmp12)
    tmp14 = tmp1 / tmp13
    tmp16 = tmp15.to(tl.float64)
    tmp17 = tmp14 * tmp16
    tmp18 = tmp4 / tmp13
    tmp20 = tmp19.to(tl.float64)
    tmp21 = tmp18 * tmp20
    tmp22 = tmp17 + tmp21
    tmp23 = tmp8 / tmp13
    tmp25 = tmp24.to(tl.float64)
    tmp26 = tmp23 * tmp25
    tmp27 = tmp22 + tmp26
    tl.store(out_ptr0 + x0, tmp27, xmask)


@triton.jit
def triton_poi_fused_clamp_div_linalg_vector_norm_mul_sub_sum_1(in_ptr0,
    in_ptr1, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 192
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex % 3
    x1 = xindex // 3
    x2 = xindex
    tmp0 = tl.load(in_ptr0 + (3 + x0 + 6 * x1), xmask)
    tmp2 = tl.load(in_ptr1 + x1, xmask, eviction_policy='evict_last')
    tmp3 = tl.load(in_ptr0 + (x0 + 6 * x1), xmask)
    tmp5 = tl.load(in_ptr0 + 6 * x1, xmask, eviction_policy='evict_last')
    tmp8 = tl.load(in_ptr0 + (1 + 6 * x1), xmask, eviction_policy='evict_last')
    tmp12 = tl.load(in_ptr0 + (2 + 6 * x1), xmask, eviction_policy='evict_last'
        )
    tmp1 = tmp0.to(tl.float64)
    tmp4 = tmp3.to(tl.float64)
    tmp6 = tmp5.to(tl.float64)
    tmp7 = tmp6 * tmp6
    tmp9 = tmp8.to(tl.float64)
    tmp10 = tmp9 * tmp9
    tmp11 = tmp7 + tmp10
    tmp13 = tmp12.to(tl.float64)
    tmp14 = tmp13 * tmp13
    tmp15 = tmp11 + tmp14
    tmp16 = libdevice.sqrt(tmp15)
    tmp17 = tl.full([1], 1e-05, tl.float64)
    tmp18 = triton_helpers.maximum(tmp16, tmp17)
    tmp19 = tmp4 / tmp18
    tmp20 = tmp2 * tmp19
    tmp21 = tmp1 - tmp20
    tl.store(out_ptr0 + x2, tmp21, xmask)


@triton.jit
def triton_poi_fused_clamp_div_linalg_cross_linalg_vector_norm_2(in_ptr0,
    in_ptr1, out_ptr0, out_ptr1, xnumel, XBLOCK: tl.constexpr):
    xnumel = 192
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex % 3
    x1 = xindex // 3
    x2 = xindex
    tmp0 = tl.load(in_ptr0 + (6 * x1 + (1 + x0) % 3), xmask)
    tmp2 = tl.load(in_ptr0 + 6 * x1, xmask, eviction_policy='evict_last')
    tmp5 = tl.load(in_ptr0 + (1 + 6 * x1), xmask, eviction_policy='evict_last')
    tmp9 = tl.load(in_ptr0 + (2 + 6 * x1), xmask, eviction_policy='evict_last')
    tmp17 = tl.load(in_ptr1 + (3 * x1 + (2 + x0) % 3), xmask,
        eviction_policy='evict_last')
    tmp18 = tl.load(in_ptr1 + 3 * x1, xmask, eviction_policy='evict_last')
    tmp20 = tl.load(in_ptr1 + (1 + 3 * x1), xmask, eviction_policy='evict_last'
        )
    tmp23 = tl.load(in_ptr1 + (2 + 3 * x1), xmask, eviction_policy='evict_last'
        )
    tmp30 = tl.load(in_ptr0 + (6 * x1 + (2 + x0) % 3), xmask,
        eviction_policy='evict_last')
    tmp33 = tl.load(in_ptr1 + (3 * x1 + (1 + x0) % 3), xmask)
    tmp1 = tmp0.to(tl.float64)
    tmp3 = tmp2.to(tl.float64)
    tmp4 = tmp3 * tmp3
    tmp6 = tmp5.to(tl.float64)
    tmp7 = tmp6 * tmp6
    tmp8 = tmp4 + tmp7
    tmp10 = tmp9.to(tl.float64)
    tmp11 = tmp10 * tmp10
    tmp12 = tmp8 + tmp11
    tmp13 = libdevice.sqrt(tmp12)
    tmp14 = tl.full([1], 1e-05, tl.float64)
    tmp15 = triton_helpers.maximum(tmp13, tmp14)
    tmp16 = tmp1 / tmp15
    tmp19 = tmp18 * tmp18
    tmp21 = tmp20 * tmp20
    tmp22 = tmp19 + tmp21
    tmp24 = tmp23 * tmp23
    tmp25 = tmp22 + tmp24
    tmp26 = libdevice.sqrt(tmp25)
    tmp27 = triton_helpers.maximum(tmp26, tmp14)
    tmp28 = tmp17 / tmp27
    tmp29 = tmp16 * tmp28
    tmp31 = tmp30.to(tl.float64)
    tmp32 = tmp31 / tmp15
    tmp34 = tmp33 / tmp27
    tmp35 = tmp32 * tmp34
    tl.store(out_ptr0 + x2, tmp29, xmask)
    tl.store(out_ptr1 + x2, tmp35, xmask)


@triton.jit
def triton_poi_fused__to_copy_stack_3(in_ptr0, in_ptr1, in_ptr2, in_ptr3,
    out_ptr1, xnumel, XBLOCK: tl.constexpr):
    xnumel = 576
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex % 9
    x1 = xindex // 9
    x2 = xindex
    tmp0 = x0
    tl.full([1], 0, tl.int64)
    tmp3 = tl.full([1], 3, tl.int64)
    tmp4 = tmp0 < tmp3
    tmp5 = tl.load(in_ptr0 + (6 * x1 + x0), tmp4 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp6 = tmp5.to(tl.float64)
    tmp7 = tl.load(in_ptr0 + 6 * x1, tmp4 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp8 = tmp7.to(tl.float64)
    tmp9 = tmp8 * tmp8
    tmp10 = tl.load(in_ptr0 + (1 + 6 * x1), tmp4 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp11 = tmp10.to(tl.float64)
    tmp12 = tmp11 * tmp11
    tmp13 = tmp9 + tmp12
    tmp14 = tl.load(in_ptr0 + (2 + 6 * x1), tmp4 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp15 = tmp14.to(tl.float64)
    tmp16 = tmp15 * tmp15
    tmp17 = tmp13 + tmp16
    tmp18 = libdevice.sqrt(tmp17)
    tmp19 = tl.full([1], 1e-05, tl.float64)
    tmp20 = triton_helpers.maximum(tmp18, tmp19)
    tmp21 = tmp6 / tmp20
    tmp22 = tl.full(tmp21.shape, 0.0, tmp21.dtype)
    tmp23 = tl.where(tmp4, tmp21, tmp22)
    tmp24 = tmp0 >= tmp3
    tmp25 = tl.full([1], 6, tl.int64)
    tmp26 = tmp0 < tmp25
    tmp27 = tmp24 & tmp26
    tmp28 = tl.load(in_ptr1 + (3 * x1 + (-3 + x0)), tmp27 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp29 = tl.load(in_ptr1 + 3 * x1, tmp27 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp30 = tmp29 * tmp29
    tmp31 = tl.load(in_ptr1 + (1 + 3 * x1), tmp27 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp32 = tmp31 * tmp31
    tmp33 = tmp30 + tmp32
    tmp34 = tl.load(in_ptr1 + (2 + 3 * x1), tmp27 & xmask, eviction_policy=
        'evict_last', other=0.0)
    tmp35 = tmp34 * tmp34
    tmp36 = tmp33 + tmp35
    tmp37 = libdevice.sqrt(tmp36)
    tmp38 = triton_helpers.maximum(tmp37, tmp19)
    tmp39 = tmp28 / tmp38
    tmp40 = tl.full(tmp39.shape, 0.0, tmp39.dtype)
    tmp41 = tl.where(tmp27, tmp39, tmp40)
    tmp42 = tmp0 >= tmp25
    tl.full([1], 9, tl.int64)
    tmp45 = tl.load(in_ptr2 + (3 * x1 + (-6 + x0)), tmp42 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp46 = tl.load(in_ptr3 + (3 * x1 + (-6 + x0)), tmp42 & xmask,
        eviction_policy='evict_last', other=0.0)
    tmp47 = tmp45 - tmp46
    tmp48 = tl.full(tmp47.shape, 0.0, tmp47.dtype)
    tmp49 = tl.where(tmp42, tmp47, tmp48)
    tmp50 = tl.where(tmp27, tmp41, tmp49)
    tmp51 = tl.where(tmp4, tmp23, tmp50)
    tmp52 = tmp51.to(tl.float32)
    tl.store(out_ptr1 + x2, tmp52, xmask)


def call(args):
    primals_1, primals_2, primals_3 = args
    args.clear()
    assert_size_stride(primals_1, (6, 4), (4, 1))
    assert_size_stride(primals_2, (6,), (1,))
    assert_size_stride(primals_3, (4, 4, 4, 4), (64, 16, 4, 1))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf0 = empty_strided_cuda((64, 6), (6, 1), torch.float32)
        extern_kernels.addmm(primals_2, reinterpret_tensor(primals_3, (64, 
            4), (4, 1), 0), reinterpret_tensor(primals_1, (4, 6), (1, 4), 0
            ), alpha=1, beta=1, out=buf0)
        del primals_1
        del primals_2
        buf1 = empty_strided_cuda((64, 1), (1, 64), torch.float64)
        get_raw_stream(0)
        triton_poi_fused_clamp_div_linalg_vector_norm_mul_sum_0[grid(64)](buf0,
            buf1, 64, XBLOCK=64, num_warps=1, num_stages=1)
        buf2 = empty_strided_cuda((64, 3), (3, 1), torch.float64)
        triton_poi_fused_clamp_div_linalg_vector_norm_mul_sub_sum_1[grid(192)](
            buf0, buf1, buf2, 192, XBLOCK=128, num_warps=4, num_stages=1)
        del buf1
        buf3 = empty_strided_cuda((64, 3), (3, 1), torch.float64)
        buf4 = empty_strided_cuda((64, 3), (3, 1), torch.float64)
        triton_poi_fused_clamp_div_linalg_cross_linalg_vector_norm_2[grid(192)
            ](buf0, buf2, buf3, buf4, 192, XBLOCK=128, num_warps=4,
            num_stages=1)
        buf6 = empty_strided_cuda((64, 3, 3), (9, 3, 1), torch.float32)
        triton_poi_fused__to_copy_stack_3[grid(576)](buf0, buf2, buf3, buf4,
            buf6, 576, XBLOCK=128, num_warps=4, num_stages=1)
        del buf2
        del buf3
        del buf4
    return buf6, reinterpret_tensor(primals_3, (64, 4), (4, 1), 0), buf0


def s2s2_gram_schmidt(v1, v2):
    """Normalise 2 3-vectors. Project second to orthogonal component.
    Take cross product for third. Stack to form SO matrix."""
    u1 = v1
    e1 = u1 / u1.norm(p=2, dim=-1, keepdim=True).clamp(min=1e-05)
    u2 = v2 - (e1 * v2).sum(-1, keepdim=True) * e1
    e2 = u2 / u2.norm(p=2, dim=-1, keepdim=True).clamp(min=1e-05)
    e3 = torch.cross(e1, e2)
    return torch.stack([e1, e2, e3], 1)


class S2S2MeanNew(nn.Module):
    """Module to map R^6 -> SO(3) with S2S2 method."""

    def __init__(self, input_dims):
        super().__init__()
        self.map = nn.Linear(input_dims, 6)
        self.map.weight.data.uniform_(-10, 10)
        self.map.bias.data.uniform_(-10, 10)

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


Measured latency_ms: 0.068291