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.nn import Module
import numpy as np
from numpy.random import RandomState
from torch.autograd import Variable
from torch.nn.parameter import Parameter
from scipy.stats import chi
import torch.fx
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_cat_0(in_ptr0, in_ptr1, in_ptr2, in_ptr3, out_ptr0,
    xnumel, XBLOCK: tl.constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex % 4
    x1 = xindex // 4
    x2 = xindex
    tmp9 = tl.load(in_ptr0 + 0)
    tmp10 = tl.broadcast_to(tmp9, [XBLOCK])
    tmp16 = tl.load(in_ptr1 + 0)
    tmp17 = tl.broadcast_to(tmp16, [XBLOCK])
    tmp26 = tl.load(in_ptr2 + 0)
    tmp27 = tl.broadcast_to(tmp26, [XBLOCK])
    tmp35 = tl.load(in_ptr3 + 0)
    tmp36 = tl.broadcast_to(tmp35, [XBLOCK])
    tmp0 = x0
    tl.full([1], 0, tl.int64)
    tmp3 = tl.full([1], 1, tl.int64)
    tmp4 = tmp0 < tmp3
    tmp5 = x1
    tmp7 = tmp5 < tmp3
    tmp7 & tmp4
    tmp11 = tmp5 >= tmp3
    tmp12 = tl.full([1], 2, tl.int64)
    tmp13 = tmp5 < tmp12
    tmp14 = tmp11 & tmp13
    tmp15 = tmp14 & tmp4
    tmp18 = -tmp17
    tmp19 = tl.full(tmp18.shape, 0.0, tmp18.dtype)
    tmp20 = tl.where(tmp15, tmp18, tmp19)
    tmp21 = tmp5 >= tmp12
    tmp22 = tl.full([1], 3, tl.int64)
    tmp23 = tmp5 < tmp22
    tmp24 = tmp21 & tmp23
    tmp25 = tmp24 & tmp4
    tmp28 = -tmp27
    tmp29 = tl.full(tmp28.shape, 0.0, tmp28.dtype)
    tmp30 = tl.where(tmp25, tmp28, tmp29)
    tmp31 = tmp5 >= tmp22
    tl.full([1], 4, tl.int64)
    tmp34 = tmp31 & tmp4
    tmp37 = -tmp36
    tmp38 = tl.full(tmp37.shape, 0.0, tmp37.dtype)
    tmp39 = tl.where(tmp34, tmp37, tmp38)
    tmp40 = tl.where(tmp24, tmp30, tmp39)
    tmp41 = tl.where(tmp14, tmp20, tmp40)
    tmp42 = tl.where(tmp7, tmp10, tmp41)
    tmp43 = tl.full(tmp42.shape, 0.0, tmp42.dtype)
    tmp44 = tl.where(tmp4, tmp42, tmp43)
    tmp45 = tmp0 >= tmp3
    tmp46 = tmp0 < tmp12
    tmp47 = tmp45 & tmp46
    tmp7 & tmp47
    tmp14 & tmp47
    tmp50 = tmp24 & tmp47
    tmp51 = tl.where(tmp50, tmp37, tmp38)
    tmp31 & tmp47
    tmp53 = tl.where(tmp24, tmp51, tmp27)
    tmp54 = tl.where(tmp14, tmp10, tmp53)
    tmp55 = tl.where(tmp7, tmp17, tmp54)
    tmp56 = tl.full(tmp55.shape, 0.0, tmp55.dtype)
    tmp57 = tl.where(tmp47, tmp55, tmp56)
    tmp58 = tmp0 >= tmp12
    tmp59 = tmp0 < tmp22
    tmp60 = tmp58 & tmp59
    tmp7 & tmp60
    tmp14 & tmp60
    tmp24 & tmp60
    tmp64 = tmp31 & tmp60
    tmp65 = tl.where(tmp64, tmp18, tmp19)
    tmp66 = tl.where(tmp24, tmp10, tmp65)
    tmp67 = tl.where(tmp14, tmp36, tmp66)
    tmp68 = tl.where(tmp7, tmp27, tmp67)
    tmp69 = tl.full(tmp68.shape, 0.0, tmp68.dtype)
    tmp70 = tl.where(tmp60, tmp68, tmp69)
    tmp71 = tmp0 >= tmp22
    tmp7 & tmp71
    tmp74 = tmp14 & tmp71
    tmp75 = tl.where(tmp74, tmp28, tmp29)
    tmp24 & tmp71
    tmp31 & tmp71
    tmp78 = tl.where(tmp24, tmp17, tmp10)
    tmp79 = tl.where(tmp14, tmp75, tmp78)
    tmp80 = tl.where(tmp7, tmp36, tmp79)
    tmp81 = tl.full(tmp80.shape, 0.0, tmp80.dtype)
    tmp82 = tl.where(tmp71, tmp80, tmp81)
    tmp83 = tl.where(tmp60, tmp70, tmp82)
    tmp84 = tl.where(tmp47, tmp57, tmp83)
    tmp85 = tl.where(tmp4, tmp44, tmp84)
    tl.store(out_ptr0 + x2, tmp85, xmask)


def call(args):
    arg0_1, arg1_1, arg2_1, arg3_1, arg4_1, arg5_1 = args
    args.clear()
    assert_size_stride(arg0_1, (4, 4), (4, 1))
    assert_size_stride(arg1_1, (1, 1), (1, 1))
    assert_size_stride(arg2_1, (1, 1), (1, 1))
    assert_size_stride(arg3_1, (1, 1), (1, 1))
    assert_size_stride(arg4_1, (1, 1), (1, 1))
    assert_size_stride(arg5_1, (4,), (1,))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf0 = empty_strided_cuda((4, 4), (4, 1), torch.float32)
        get_raw_stream(0)
        triton_poi_fused_cat_0[grid(16)](arg4_1, arg1_1, arg2_1, arg3_1,
            buf0, 16, XBLOCK=16, num_warps=1, num_stages=1)
        del arg1_1
        del arg2_1
        del arg3_1
        del arg4_1
        buf1 = empty_strided_cuda((4, 4), (4, 1), torch.float32)
        extern_kernels.addmm(arg5_1, arg0_1, buf0, alpha=1, beta=1, out=buf1)
        del arg0_1
        del arg5_1
        del buf0
    return buf1,


def quaternion_init(in_features, out_features, rng, kernel_size=None,
    criterion='glorot'):
    if kernel_size is not None:
        receptive_field = np.prod(kernel_size)
        fan_in = in_features * receptive_field
        fan_out = out_features * receptive_field
    else:
        fan_in = in_features
        fan_out = out_features
    if criterion == 'glorot':
        s = 1.0 / np.sqrt(2 * (fan_in + fan_out))
    elif criterion == 'he':
        s = 1.0 / np.sqrt(2 * fan_in)
    else:
        raise ValueError('Invalid criterion: ' + criterion)
    rng = RandomState(np.random.randint(1, 1234))
    if kernel_size is None:
        kernel_shape = in_features, out_features
    elif type(kernel_size) is int:
        kernel_shape = (out_features, in_features) + tuple((kernel_size,))
    else:
        kernel_shape = (out_features, in_features) + (*kernel_size,)
    modulus = chi.rvs(4, loc=0, scale=s, size=kernel_shape)
    number_of_weights = np.prod(kernel_shape)
    v_i = np.random.uniform(-1.0, 1.0, number_of_weights)
    v_j = np.random.uniform(-1.0, 1.0, number_of_weights)
    v_k = np.random.uniform(-1.0, 1.0, number_of_weights)
    for i in range(0, number_of_weights):
        norm = np.sqrt(v_i[i] ** 2 + v_j[i] ** 2 + v_k[i] ** 2 + 0.0001)
        v_i[i] /= norm
        v_j[i] /= norm
        v_k[i] /= norm
    v_i = v_i.reshape(kernel_shape)
    v_j = v_j.reshape(kernel_shape)
    v_k = v_k.reshape(kernel_shape)
    phase = rng.uniform(low=-np.pi, high=np.pi, size=kernel_shape)
    weight_r = modulus * np.cos(phase)
    weight_i = modulus * v_i * np.sin(phase)
    weight_j = modulus * v_j * np.sin(phase)
    weight_k = modulus * v_k * np.sin(phase)
    return weight_r, weight_i, weight_j, weight_k


def unitary_init(in_features, out_features, rng, kernel_size=None,
    criterion='he'):
    if kernel_size is not None:
        receptive_field = np.prod(kernel_size)
        in_features * receptive_field
        out_features * receptive_field
    else:
        pass
    if kernel_size is None:
        kernel_shape = in_features, out_features
    elif type(kernel_size) is int:
        kernel_shape = (out_features, in_features) + tuple((kernel_size,))
    else:
        kernel_shape = (out_features, in_features) + (*kernel_size,)
    number_of_weights = np.prod(kernel_shape)
    v_r = np.random.uniform(-1.0, 1.0, number_of_weights)
    v_i = np.random.uniform(-1.0, 1.0, number_of_weights)
    v_j = np.random.uniform(-1.0, 1.0, number_of_weights)
    v_k = np.random.uniform(-1.0, 1.0, number_of_weights)
    for i in range(0, number_of_weights):
        norm = np.sqrt(v_r[i] ** 2 + v_i[i] ** 2 + v_j[i] ** 2 + v_k[i] ** 2
            ) + 0.0001
        v_r[i] /= norm
        v_i[i] /= norm
        v_j[i] /= norm
        v_k[i] /= norm
    v_r = v_r.reshape(kernel_shape)
    v_i = v_i.reshape(kernel_shape)
    v_j = v_j.reshape(kernel_shape)
    v_k = v_k.reshape(kernel_shape)
    return v_r, v_i, v_j, v_k


def affect_init(r_weight, i_weight, j_weight, k_weight, init_func, rng,
    init_criterion):
    if r_weight.size() != i_weight.size() or r_weight.size() != j_weight.size(
        ) or r_weight.size() != k_weight.size():
        raise ValueError(
            'The real and imaginary weights should have the same size . Found: r:'
             + str(r_weight.size()) + ' i:' + str(i_weight.size()) + ' j:' +
            str(j_weight.size()) + ' k:' + str(k_weight.size()))
    elif r_weight.dim() != 2:
        raise Exception(
            'affect_init accepts only matrices. Found dimension = ' + str(
            r_weight.dim()))
    kernel_size = None
    r, i, j, k = init_func(r_weight.size(0), r_weight.size(1), rng,
        kernel_size, init_criterion)
    r, i, j, k = torch.from_numpy(r), torch.from_numpy(i), torch.from_numpy(j
        ), torch.from_numpy(k)
    r_weight.data = r.type_as(r_weight.data)
    i_weight.data = i.type_as(i_weight.data)
    j_weight.data = j.type_as(j_weight.data)
    k_weight.data = k.type_as(k_weight.data)


def check_input(input):
    if input.dim() not in {2, 3, 4, 5}:
        raise RuntimeError(
            'Quaternion linear accepts only input of dimension 2 or 3. Quaternion conv accepts up to 5 dim  input.dim = '
             + str(input.dim()))
    if input.dim() < 4:
        nb_hidden = input.size()[-1]
    else:
        nb_hidden = input.size()[1]
    if nb_hidden % 4 != 0:
        raise RuntimeError(
            'Quaternion Tensors must be divisible by 4. input.size()[1] = ' +
            str(nb_hidden))


def get_i(input):
    if input.dim() < 4:
        nb_hidden = input.size()[-1]
    else:
        nb_hidden = input.size()[1]
    if input.dim() == 2:
        return input.narrow(1, nb_hidden // 4, nb_hidden // 4)
    if input.dim() == 3:
        return input.narrow(2, nb_hidden // 4, nb_hidden // 4)
    if input.dim() >= 4:
        return input.narrow(1, nb_hidden // 4, nb_hidden // 4)


def get_j(input):
    check_input(input)
    if input.dim() < 4:
        nb_hidden = input.size()[-1]
    else:
        nb_hidden = input.size()[1]
    if input.dim() == 2:
        return input.narrow(1, nb_hidden // 2, nb_hidden // 4)
    if input.dim() == 3:
        return input.narrow(2, nb_hidden // 2, nb_hidden // 4)
    if input.dim() >= 4:
        return input.narrow(1, nb_hidden // 2, nb_hidden // 4)


def get_k(input):
    check_input(input)
    if input.dim() < 4:
        nb_hidden = input.size()[-1]
    else:
        nb_hidden = input.size()[1]
    if input.dim() == 2:
        return input.narrow(1, nb_hidden - nb_hidden // 4, nb_hidden // 4)
    if input.dim() == 3:
        return input.narrow(2, nb_hidden - nb_hidden // 4, nb_hidden // 4)
    if input.dim() >= 4:
        return input.narrow(1, nb_hidden - nb_hidden // 4, nb_hidden // 4)


def get_r(input):
    check_input(input)
    if input.dim() < 4:
        nb_hidden = input.size()[-1]
    else:
        nb_hidden = input.size()[1]
    if input.dim() == 2:
        return input.narrow(1, 0, nb_hidden // 4)
    if input.dim() == 3:
        return input.narrow(2, 0, nb_hidden // 4)
    if input.dim() >= 4:
        return input.narrow(1, 0, nb_hidden // 4)


class QuaternionLinearFunction(torch.autograd.Function):

    @staticmethod
    def forward(ctx, input, r_weight, i_weight, j_weight, k_weight, bias=None):
        ctx.save_for_backward(input, r_weight, i_weight, j_weight, k_weight,
            bias)
        check_input(input)
        cat_kernels_4_r = torch.cat([r_weight, -i_weight, -j_weight, -
            k_weight], dim=0)
        cat_kernels_4_i = torch.cat([i_weight, r_weight, -k_weight,
            j_weight], dim=0)
        cat_kernels_4_j = torch.cat([j_weight, k_weight, r_weight, -
            i_weight], dim=0)
        cat_kernels_4_k = torch.cat([k_weight, -j_weight, i_weight,
            r_weight], dim=0)
        cat_kernels_4_quaternion = torch.cat([cat_kernels_4_r,
            cat_kernels_4_i, cat_kernels_4_j, cat_kernels_4_k], dim=1)
        if input.dim() == 2:
            if bias is not None:
                return torch.addmm(bias, input, cat_kernels_4_quaternion)
            else:
                return torch.mm(input, cat_kernels_4_quaternion)
        else:
            output = torch.matmul(input, cat_kernels_4_quaternion)
            if bias is not None:
                return output + bias
            else:
                return output

    @staticmethod
    def backward(ctx, grad_output):
        input, r_weight, i_weight, j_weight, k_weight, _bias = (ctx.
            saved_tensors)
        (grad_input) = (grad_weight_r) = (grad_weight_i) = (grad_weight_j) = (
            grad_weight_k) = (grad_bias) = None
        input_r = torch.cat([r_weight, -i_weight, -j_weight, -k_weight], dim=0)
        input_i = torch.cat([i_weight, r_weight, -k_weight, j_weight], dim=0)
        input_j = torch.cat([j_weight, k_weight, r_weight, -i_weight], dim=0)
        input_k = torch.cat([k_weight, -j_weight, i_weight, r_weight], dim=0)
        cat_kernels_4_quaternion_T = Variable(torch.cat([input_r, input_i,
            input_j, input_k], dim=1).permute(1, 0), requires_grad=False)
        r = get_r(input)
        i = get_i(input)
        j = get_j(input)
        k = get_k(input)
        input_r = torch.cat([r, -i, -j, -k], dim=0)
        input_i = torch.cat([i, r, -k, j], dim=0)
        input_j = torch.cat([j, k, r, -i], dim=0)
        input_k = torch.cat([k, -j, i, r], dim=0)
        input_mat = Variable(torch.cat([input_r, input_i, input_j, input_k],
            dim=1), requires_grad=False)
        r = get_r(grad_output)
        i = get_i(grad_output)
        j = get_j(grad_output)
        k = get_k(grad_output)
        input_r = torch.cat([r, i, j, k], dim=1)
        input_i = torch.cat([-i, r, k, -j], dim=1)
        input_j = torch.cat([-j, -k, r, i], dim=1)
        input_k = torch.cat([-k, j, -i, r], dim=1)
        grad_mat = torch.cat([input_r, input_i, input_j, input_k], dim=0)
        if ctx.needs_input_grad[0]:
            grad_input = grad_output.mm(cat_kernels_4_quaternion_T)
        if ctx.needs_input_grad[1]:
            grad_weight = grad_mat.permute(1, 0).mm(input_mat).permute(1, 0)
            unit_size_x = r_weight.size(0)
            unit_size_y = r_weight.size(1)
            grad_weight_r = grad_weight.narrow(0, 0, unit_size_x).narrow(1,
                0, unit_size_y)
            grad_weight_i = grad_weight.narrow(0, 0, unit_size_x).narrow(1,
                unit_size_y, unit_size_y)
            grad_weight_j = grad_weight.narrow(0, 0, unit_size_x).narrow(1,
                unit_size_y * 2, unit_size_y)
            grad_weight_k = grad_weight.narrow(0, 0, unit_size_x).narrow(1,
                unit_size_y * 3, unit_size_y)
        if ctx.needs_input_grad[5]:
            grad_bias = grad_output.sum(0).squeeze(0)
        return (grad_input, grad_weight_r, grad_weight_i, grad_weight_j,
            grad_weight_k, grad_bias)


class QuaternionLinearNew(Module):
    """Applies a quaternion linear transformation to the incoming data.
    """

    def __init__(self, in_features, out_features, bias=True, init_criterion
        ='he', weight_init='quaternion', seed=None):
        super(QuaternionLinearNew, self).__init__()
        self.in_features = in_features // 4
        self.out_features = out_features // 4
        self.r_weight = Parameter(torch.Tensor(self.in_features, self.
            out_features))
        self.i_weight = Parameter(torch.Tensor(self.in_features, self.
            out_features))
        self.j_weight = Parameter(torch.Tensor(self.in_features, self.
            out_features))
        self.k_weight = Parameter(torch.Tensor(self.in_features, self.
            out_features))
        if bias:
            self.bias = Parameter(torch.Tensor(self.out_features * 4))
        else:
            self.register_parameter('bias', None)
        self.init_criterion = init_criterion
        self.weight_init = weight_init
        self.seed = seed if seed is not None else np.random.randint(0, 1234)
        self.rng = RandomState(self.seed)
        self.reset_parameters()

    def reset_parameters(self):
        winit = {'quaternion': quaternion_init, 'unitary': unitary_init}[self
            .weight_init]
        if self.bias is not None:
            self.bias.data.fill_(0)
        affect_init(self.r_weight, self.i_weight, self.j_weight, self.
            k_weight, winit, self.rng, self.init_criterion)

    def __repr__(self):
        return self.__class__.__name__ + '(' + 'in_features=' + str(self.
            in_features) + ', out_features=' + str(self.out_features
            ) + ', bias=' + str(self.bias is not None
            ) + ', init_criterion=' + str(self.init_criterion
            ) + ', weight_init=' + str(self.weight_init) + ', seed=' + str(self
            .seed) + ')'

    def forward(self, input_0):
        arg1_1 = self.r_weight
        arg2_1 = self.i_weight
        arg3_1 = self.j_weight
        arg4_1 = self.k_weight
        arg5_1 = self.bias
        arg0_1 = input_0
        output = call([arg0_1, arg1_1, arg2_1, arg3_1, arg4_1, arg5_1])
        return output[0]


Measured latency_ms: 0.024387