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
import re
import warnings
from torch import nn
from collections import OrderedDict
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_add_view_0(in_out_ptr0, in_ptr0, xnumel, XBLOCK: tl.
    constexpr):
    xnumel = 256
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x4 = xindex
    x0 = xindex % 4
    tmp0 = tl.load(in_out_ptr0 + x4, xmask)
    tmp1 = tl.load(in_ptr0 + x0, xmask, eviction_policy='evict_last')
    tmp2 = tmp0 + tmp1
    tl.store(in_out_ptr0 + x4, tmp2, xmask)


def call(args):
    primals_1, primals_2, primals_3 = args
    args.clear()
    assert_size_stride(primals_1, (4, 4), (4, 1))
    assert_size_stride(primals_2, (4,), (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, 4), (4, 1), torch.float32)
        extern_kernels.mm(reinterpret_tensor(primals_3, (64, 4), (4, 1), 0),
            reinterpret_tensor(primals_1, (4, 4), (1, 4), 0), out=buf0)
        del primals_1
        buf1 = reinterpret_tensor(buf0, (4, 4, 4, 4), (64, 16, 4, 1), 0)
        del buf0
        buf2 = buf1
        del buf1
        get_raw_stream(0)
        triton_poi_fused_add_view_0[grid(256)](buf2, primals_2, 256, XBLOCK
            =256, num_warps=4, num_stages=1)
        del primals_2
    return buf2, reinterpret_tensor(primals_3, (64, 4), (4, 1), 0)


class MetaModule(nn.Module):
    """
    Base class for PyTorch meta-learning modules. These modules accept an
    additional argument `params` in their `forward` method.

    Notes
    -----
    Objects inherited from `MetaModule` are fully compatible with PyTorch
    modules from `torch.nn.Module`. The argument `params` is a dictionary of
    tensors, with full support of the computation graph (for differentiation).
    """

    def __init__(self):
        super(MetaModule, self).__init__()
        self._children_modules_parameters_cache = dict()

    def meta_named_parameters(self, prefix='', recurse=True):
        gen = self._named_members(lambda module: module._parameters.items() if
            isinstance(module, MetaModule) else [], prefix=prefix, recurse=
            recurse)
        for elem in gen:
            yield elem

    def meta_parameters(self, recurse=True):
        for name, param in self.meta_named_parameters(recurse=recurse):
            yield param

    def get_subdict(self, params, key=None):
        if params is None:
            return None
        all_names = tuple(params.keys())
        if (key, all_names) not in self._children_modules_parameters_cache:
            if key is None:
                self._children_modules_parameters_cache[key, all_names
                    ] = all_names
            else:
                key_escape = re.escape(key)
                key_re = re.compile('^{0}\\.(.+)'.format(key_escape))
                self._children_modules_parameters_cache[key, all_names] = [
                    key_re.sub('\\1', k) for k in all_names if key_re.match
                    (k) is not None]
        names = self._children_modules_parameters_cache[key, all_names]
        if not names:
            warnings.warn(
                'Module `{0}` has no parameter corresponding to the submodule named `{1}` in the dictionary `params` provided as an argument to `forward()`. Using the default parameters for this submodule. The list of the parameters in `params`: [{2}].'
                .format(self.__class__.__name__, key, ', '.join(all_names)),
                stacklevel=2)
            return None
        return OrderedDict([(name, params[f'{key}.{name}']) for name in names])


class BatchLinearNew(nn.Linear, MetaModule):
    """A linear meta-layer that can deal with batched weight matrices and biases, as for instance output by a
    hypernetwork."""
    __doc__ = nn.Linear.__doc__

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


Measured latency_ms: 0.022198