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.triton_helpers import math as tl_math
from torch.nn import Module
import math
import numpy as np
from torch.nn import *
from torch.optim import AdamW
from typing import Union
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_pow_view_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
    x0 = xindex
    tmp0 = tl.load(in_ptr0 + x0, xmask)
    tmp1 = tmp0 * tmp0
    tl.store(out_ptr0 + x0, tmp1, xmask)


@triton.jit
def triton_poi_fused_pow_1(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 100
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex
    tmp0 = tl.load(in_ptr0 + x0, xmask)
    tmp1 = tmp0 * tmp0
    tl.store(out_ptr0 + x0, tmp1, xmask)


@triton.jit
def triton_poi_fused_add_mul_pow_sub_2(in_out_ptr0, in_ptr0, in_ptr1,
    in_ptr2, xnumel, XBLOCK: tl.constexpr):
    xnumel = 1600
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x1 = xindex // 25
    x2 = xindex
    tmp0 = tl.load(in_ptr0 + 0)
    tmp1 = tl.broadcast_to(tmp0, [XBLOCK])
    tmp2 = tl.load(in_ptr1 + x1, xmask, eviction_policy='evict_last')
    tmp4 = tl.load(in_ptr2 + x2, xmask)
    tmp9 = tl.load(in_out_ptr0 + x2, xmask)
    tmp3 = tmp1 + tmp2
    tmp5 = tmp4 * tmp4
    tmp6 = 0.5
    tmp7 = tmp5 * tmp6
    tmp8 = tmp3 + tmp7
    tmp10 = tmp8 - tmp9
    tl.store(in_out_ptr0 + x2, tmp10, xmask)


@triton.jit
def triton_poi_fused_abs_sgn_3(in_ptr0, out_ptr0, out_ptr1, xnumel, XBLOCK:
    tl.constexpr):
    xnumel = 32
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex % 8
    x1 = xindex // 8
    x2 = xindex
    tmp0 = tl.load(in_ptr0 + (8 + x0 + 32 * x1), xmask)
    tmp1 = tl_math.abs(tmp0)
    tmp2 = tl.full([1], 0, tl.int32)
    tmp3 = tmp2 < tmp0
    tmp4 = tmp3.to(tl.int8)
    tmp5 = tmp0 < tmp2
    tmp6 = tmp5.to(tl.int8)
    tmp7 = tmp4 - tmp6
    tmp8 = tmp7.to(tmp0.dtype)
    tl.store(out_ptr0 + x2, tmp1, xmask)
    tl.store(out_ptr1 + x2, tmp8, xmask)


def call(args):
    primals_1, primals_2, primals_3, primals_4, primals_5, primals_6 = args
    args.clear()
    assert_size_stride(primals_1, (4, 25), (25, 1))
    assert_size_stride(primals_2, (4, 4, 4, 4), (64, 16, 4, 1))
    assert_size_stride(primals_3, (4, 1), (1, 1))
    assert_size_stride(primals_4, (1,), (1,))
    assert_size_stride(primals_5, (2, 25), (25, 1))
    assert_size_stride(primals_6, (2,), (1,))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf0 = empty_strided_cuda((64, 25), (25, 1), torch.float32)
        extern_kernels.mm(reinterpret_tensor(primals_2, (64, 4), (4, 1), 0),
            primals_1, out=buf0)
        buf1 = empty_strided_cuda((64, 4), (4, 1), torch.float32)
        get_raw_stream(0)
        triton_poi_fused_pow_view_0[grid(256)](primals_2, buf1, 256, XBLOCK
            =128, num_warps=4, num_stages=1)
        buf2 = empty_strided_cuda((4, 25), (25, 1), torch.float32)
        triton_poi_fused_pow_1[grid(100)](primals_1, buf2, 100, XBLOCK=128,
            num_warps=4, num_stages=1)
        buf3 = empty_strided_cuda((64, 25), (25, 1), torch.float32)
        extern_kernels.mm(buf1, buf2, out=buf3)
        del buf2
        buf4 = empty_strided_cuda((64, 1), (1, 1), torch.float32)
        extern_kernels.mm(reinterpret_tensor(primals_2, (64, 4), (4, 1), 0),
            primals_3, out=buf4)
        del primals_3
        buf5 = reinterpret_tensor(buf3, (4, 4, 4, 25), (400, 100, 25, 1), 0)
        del buf3
        triton_poi_fused_add_mul_pow_sub_2[grid(1600)](buf5, primals_4,
            buf4, buf0, 1600, XBLOCK=256, num_warps=4, num_stages=1)
        del buf4
        del primals_4
        buf6 = empty_strided_cuda((64, 2), (2, 1), torch.float32)
        extern_kernels.addmm(primals_6, reinterpret_tensor(buf5, (64, 25),
            (25, 1), 0), reinterpret_tensor(primals_5, (25, 2), (1, 25), 0),
            alpha=1, beta=1, out=buf6)
        del primals_6
        buf7 = empty_strided_cuda((4, 4, 2), (8, 2, 1), torch.float32)
        buf8 = empty_strided_cuda((4, 4, 2), (8, 2, 1), torch.float32)
        triton_poi_fused_abs_sgn_3[grid(32)](buf6, buf7, buf8, 32, XBLOCK=
            32, num_warps=1, num_stages=1)
    return reinterpret_tensor(buf6, (4, 4, 2), (32, 2, 1), 0
        ), buf7, primals_1, buf0, reinterpret_tensor(buf5, (64, 25), (25, 1), 0
        ), buf8, primals_5, reinterpret_tensor(primals_2, (4, 64), (1, 4), 0
        ), reinterpret_tensor(buf1, (4, 64), (1, 4), 0)


class FactorizationMachineNew(Module):
    """
    [Factorization Machine Recommendation Model]
    Learns latent space features to characterize similarity of dataset features
    to compute a recommendation as a function of dataset features. Dataset
    features can be mixed / hybrid such that you can combine information
    on both the recommended object and the recommendation target to generate
    an informed similarity or recommendation / ranking metric.
    """

    def __init__(self, data_dim, hidden_dim=25, seed=None) ->None:
        """
        Instantiate class attributes for FM. Constructs a feature similarity matrix
        F of shape (x_features, hidden_dim) to learn implicit representations of
        all trainable features in the data for recommendation or ranking.
        :param data_dim <int>:      Number of features to learn from in the dataset.
        :param hidden_dim <int>:    Dimension of the latent space of features.
        :param seed <int>:          Random seed fixture for reproducibility.
        """
        super().__init__()
        self.input_dim = data_dim
        self.hidden_dim = hidden_dim
        self.torch_gen = None
        if seed is not None:
            self.torch_gen = torch.manual_seed(seed)
        """ Matrix Factorization """
        self.F = Parameter(torch.empty((self.input_dim, self.hidden_dim)),
            requires_grad=True)
        init.xavier_uniform_(self.F)
        """ Linear Regression """
        self.V = Parameter(torch.empty((self.input_dim, 1)), requires_grad=True
            )
        init.xavier_uniform_(self.V)
        self.bias = Parameter(torch.zeros(1), requires_grad=True)
        """ Gaussian Regression """
        self.gaussian_dist = Linear(in_features=self.hidden_dim, out_features=2
            )

    def fit(self, X: 'Union[torch.Tensor, np.ndarray]', Y:
        'Union[torch.Tensor, np.ndarray]', mask:
        'Union[torch.Tensor, np.ndarray]'=None, cycles=100, lr=0.002,
        batch_frac=0.01, regularize=0.01, patience=3, verbose=False):
        """
        Train the Factorization Machine.
        :param X <torch.Tensor>:        Input training data features of shape (N, X).
        :param Y <torch.Tensor>:        Target training data class / score vector of shape (N, 1).
        :param mask <torch.Tensor>:     Feature observability mask for X of shape (N, X).
        :param cycles <int>:            Number of gradient descent cycles.
        :param lr <float>:              Learning rate. Re-calibrated to order of values in matrix M.
        :param batch_frac <float>:      Fraction of the dataset to set as the batch size.
        :param regularize <float>:      Weight decay lambda for regularization in AdamW.
        :param patience <int>:          Number of cycles of convergence before termination.
        :param verbose <bool>:          Output training progress information.
        """
        if any([len(X.shape) != 2, len(Y.shape) != 2, mask is not None and 
            mask.shape != X.shape, X.shape[1] != self.input_dim, Y.shape[1] !=
            1, cycles <= 0, lr <= 0, batch_frac <= 0, regularize < 0]):
            None
            return
        N = X.shape[0]
        if not torch.is_tensor(X):
            X = torch.Tensor(X)
        if not torch.is_tensor(Y):
            Y = torch.Tensor(Y)
        mask_tensor = torch.ones(X.shape)
        if mask is not None:
            mask_tensor = torch.where(torch.Tensor(mask) != 0, 1, 0)
        optimizer = AdamW(self.parameters(), lr=lr, weight_decay=regularize)
        model_opt = dict(self.state_dict())
        loss_opt = float('inf')
        timer = 0
        for i in range(cycles):
            for _ in range(math.ceil(1 / batch_frac)):
                rand_idx = torch.randint(N, size=(math.ceil(batch_frac * N)
                    ,), generator=self.torch_gen)
                X_batch = X[rand_idx]
                Y_batch = Y[rand_idx]
                mask_batch = mask_tensor[rand_idx]
                self.zero_grad()
                Y_mu, Y_sigma = self(X_batch * mask_batch)
                loss = GaussianNLLLoss()(Y_mu, Y_batch, Y_sigma)
                loss.sum().backward()
                optimizer.step()
            if i % math.ceil(cycles / 5) == 0 and verbose:
                None
            if loss.sum().item() < loss_opt:
                model_opt = dict(self.state_dict())
                loss_opt = loss.sum().item()
                timer = 0
            else:
                timer += 1
                if timer > patience:
                    self.load_state_dict(model_opt)
                    break

    def forward(self, input_0):
        primals_1 = self.F
        primals_3 = self.V
        primals_4 = self.bias
        primals_5 = self.gaussian_dist.weight
        primals_6 = self.gaussian_dist.bias
        primals_2 = input_0
        output = call([primals_1, primals_2, primals_3, primals_4,
            primals_5, primals_6])
        return output[0], output[1]


Measured latency_ms: 0.088331