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 math as tl_math
import logging
import random
import torch.nn as nn
from typing import List
import torch.onnx.operators
from functools import wraps
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__log_softmax_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')
    tmp2 = tl.load(in_ptr0 + (1 + 4 * x1), xmask, eviction_policy='evict_last')
    tmp4 = tl.load(in_ptr0 + (2 + 4 * x1), xmask, eviction_policy='evict_last')
    tmp6 = tl.load(in_ptr0 + (3 + 4 * x1), xmask, eviction_policy='evict_last')
    tmp3 = triton_helpers.maximum(tmp1, tmp2)
    tmp5 = triton_helpers.maximum(tmp3, tmp4)
    tmp7 = triton_helpers.maximum(tmp5, tmp6)
    tmp8 = tmp0 - tmp7
    tl.store(out_ptr0 + x2, tmp8, xmask)


@triton.jit
def triton_poi_fused__log_softmax_1(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 = tl_math.exp(tmp1)
    tmp4 = tl_math.exp(tmp3)
    tmp5 = tmp2 + tmp4
    tmp7 = tl_math.exp(tmp6)
    tmp8 = tmp5 + tmp7
    tmp10 = tl_math.exp(tmp9)
    tmp11 = tmp8 + tmp10
    tmp12 = tl_math.log(tmp11)
    tmp13 = tmp0 - tmp12
    tl.store(out_ptr0 + x2, tmp13, xmask)


def call(args):
    primals_1, primals_2 = args
    args.clear()
    assert_size_stride(primals_1, (4, 4), (4, 1))
    assert_size_stride(primals_2, (4, 4, 4, 4), (64, 16, 4, 1))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf0 = empty_strided_cuda((64, 4), (4, 1), torch.float32)
        extern_kernels.mm(reinterpret_tensor(primals_2, (64, 4), (4, 1), 0),
            reinterpret_tensor(primals_1, (4, 4), (1, 4), 0), out=buf0)
        del primals_1
        buf1 = empty_strided_cuda((4, 4, 4, 4), (64, 16, 4, 1), torch.float32)
        get_raw_stream(0)
        triton_poi_fused__log_softmax_0[grid(256)](buf0, buf1, 256, XBLOCK=
            256, num_warps=4, num_stages=1)
        buf2 = reinterpret_tensor(buf0, (4, 4, 4, 4), (64, 16, 4, 1), 0)
        del buf0
        triton_poi_fused__log_softmax_1[grid(256)](buf1, buf2, 256, XBLOCK=
            256, num_warps=4, num_stages=1)
        del buf1
    return buf2, reinterpret_tensor(primals_2, (64, 4), (4, 1), 0), buf2


def singleton(cls):
    """
    Singleton decorator

    Args:
        cls: singleton class

    Returns:
        - an instance of a singleton class
    """
    instances = {}

    @wraps(cls)
    def getinstance(*args, **kwargs):
        if cls not in instances:
            instances[cls] = cls(*args, **kwargs)
        return instances[cls]
    return getinstance


def get_invalid_class_mask(classes: 'int', invalid_classes: 'List'):
    """
    Create mask for invalid classes

    Args:
        classes: number of labels
        invalid_classes: invalid class list

    Returns:
        - mask for invalid class
            :math:`(1, L)` where L is the number of classes
    """
    invalid_class_mask = torch.zeros(classes).bool()
    if invalid_classes:
        for idx in invalid_classes:
            invalid_class_mask[idx] = True
    invalid_class_mask = invalid_class_mask.unsqueeze(dim=0)
    env = Environment()
    if env.device.startswith('cuda'):
        invalid_class_mask = invalid_class_mask
    return invalid_class_mask


@singleton
class Environment:
    """
    Environment is a running environment class.

    Args:
        profiling_window: profiling window size
        configs: configs for running tasks
        debug: running with debug information
        no_warning: do not output warning informations
        seed: initial seed for random and torch
        device: running device
        fp16: running with fp16
        no_progress_bar: do not show progress bar
        pb_interval: show progress bar with an interval
    """

    def __init__(self, configs=None, profiling_window: 'int'=0, debug:
        'bool'=False, no_warning: 'bool'=False, seed: 'int'=0, device:
        'str'=None, fp16: 'bool'=False, no_progress_bar: 'bool'=False,
        pb_interval: 'int'=1, custom_libs: 'str'=None):
        self.profiling_window = profiling_window
        self.configs = configs
        self.debug = debug
        self.no_warning = no_warning
        self.seed = seed
        self.fp16 = fp16
        self.no_progress_bar = no_progress_bar
        self.pb_interval = pb_interval
        self.distributed_world = 1
        self.rank = 0
        self.local_rank = 0
        if device is None:
            if torch.cuda.is_available():
                self.device = 'cuda'
            else:
                self.device = 'cpu'
        else:
            self.device = device
        if self.device == 'cuda':
            self._init_cuda()
        self._init_log()
        self._init_seed()
        self._import_custom_lib(custom_libs)

    def _init_log(self):
        FORMAT = (
            f"%(asctime)s ｜ %(levelname)s | %(name)s |{f' RANK {self.rank} | ' if not self.is_master() else ' '}%(message)s"
            )
        logging.basicConfig(format=FORMAT, datefmt='%Y-%m-%d,%H:%M:%S',
            level=logging.INFO)
        if not self.is_master():
            logging.disable(logging.INFO)

    def _import_custom_lib(self, path):
        """
        Import library manually

        Args:
            path: external libraries split with `,`
        """
        if path:
            path = path.strip('\n')
            for line in path.split(','):
                logger.info(f'import module from {line}')
                line = line.replace('/', '.')
                importlib.import_module(line)

    def _init_cuda(self):
        """
        Initialize cuda device

        We assume that the user will not run ParaGen on more than one workers with only 1 GPU
        used on each worker.
        """
        if torch.cuda.device_count() > 1:
            hvd.init()
            torch.cuda.set_device(hvd.local_rank())
            self.rank = hvd.rank()
            self.local_rank = hvd.local_rank()
            self.distributed_world = hvd.size()
        torch.cuda.empty_cache()

    def _init_seed(self):
        """
        Initialize global seed
        """
        random.seed(self.seed)
        import torch
        torch.manual_seed(self.seed)
        if self.device == 'cuda':
            torch.manual_seed(self.seed)

    def is_master(self):
        """
        check the current process is the master process
        """
        return self.rank == 0


class LinearClassifierNew(nn.Module):
    """
    Classifier with only on a linear projection.

    Args:
        d_model: feature dimensionality
        labels: number of classes
        invalid_classes (List): class that is not allowed to produce
    """

    def __init__(self, d_model, labels, invalid_classes: 'List'=None):
        super().__init__()
        self._linear = nn.Linear(d_model, labels, bias=False)
        self._invalid_class_mask = get_invalid_class_mask(labels,
            invalid_classes) if invalid_classes else None

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


Measured latency_ms: 0.0339