# AOT ID: ['0_forward']
from ctypes import c_void_p, c_long, c_int
import torch
import math
import random
import os
import tempfile
from math import inf, nan
from torch._inductor.hooks import run_intermediate_hooks
from torch._inductor.utils import maybe_profile
from torch._inductor.codegen.memory_planning import _align as align
from torch import device, empty_strided
from torch._inductor.async_compile import AsyncCompile
from torch._inductor.select_algorithm import extern_kernels
from torch._inductor.codegen.multi_kernel import MultiKernelCall
import triton
import triton.language as tl
from torch._inductor.runtime.triton_heuristics import grid, split_scan_grid, grid_combo_kernels, start_graph, end_graph
from torch._C import _cuda_getCurrentRawStream as get_raw_stream

aten = torch.ops.aten
inductor_ops = torch.ops.inductor
_quantized = torch.ops._quantized
assert_size_stride = torch._C._dynamo.guards.assert_size_stride
empty_strided_cpu = torch._C._dynamo.guards._empty_strided_cpu
empty_strided_cuda = torch._C._dynamo.guards._empty_strided_cuda
empty_strided_xpu = torch._C._dynamo.guards._empty_strided_xpu
reinterpret_tensor = torch._C._dynamo.guards._reinterpret_tensor
alloc_from_pool = torch.ops.inductor._alloc_from_pool
async_compile = AsyncCompile()


# kernel path: runs/run_shard_7/inductor_cache/td/ctdj5kazgiki6gdaadhqtp2x7tq2ee5ey5hqqdcoqmp54jyhf74f.py
# Topologically Sorted Source Nodes: [out_2], Original ATen: [aten._log_softmax]
# Source node to ATen node mapping:
#   out_2 => amax, sub
# Graph fragment:
#   %amax : [num_users=1] = call_function[target=torch.ops.aten.amax.default](args = (%view_3, [1], True), kwargs = {})
#   %sub : [num_users=2] = call_function[target=torch.ops.aten.sub.Tensor](args = (%view_3, %amax), kwargs = {})
triton_poi_fused__log_softmax_0 = async_compile.triton('triton_poi_fused__log_softmax_0', '''
import triton
import triton.language as tl
from triton.compiler.compiler import AttrsDescriptor

from torch._inductor.runtime import triton_helpers, triton_heuristics
from torch._inductor.runtime.triton_helpers import libdevice, math as tl_math
from torch._inductor.runtime.hints import AutotuneHint, ReductionHint, TileHint, instance_descriptor, DeviceProperties

@triton_heuristics.pointwise(
    size_hints=[256], 
    filename=__file__,
    triton_meta={'signature': {0: '*fp32', 1: '*fp32', 2: 'i32'}, 'device': DeviceProperties(type='cuda', index=0, cc=80, major=8, regs_per_multiprocessor=65536, max_threads_per_multi_processor=2048, multi_processor_count=108), 'constants': {}, 'configs': [AttrsDescriptor(divisible_by_16=(0, 1, 2), equal_to_1=())]},
    inductor_meta={'autotune_hints': set(), 'kernel_name': 'triton_poi_fused__log_softmax_0', 'mutated_arg_names': [], 'no_x_dim': False, 'num_load': 5, 'num_reduction': 0, 'backend_hash': 'A9C866B4A14FD3277824029365D703C2427B2E685E54EC9B3EF4ADC8D1EEAC1D', 'are_deterministic_algorithms_enabled': False, 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': None, 'force_disable_caches': False, 'dynamic_scale_rblock': True, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 16, 'store_cubin': False},
    min_elem_per_thread=0
)
@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
    x3 = xindex
    x0 = xindex % 16
    x2 = (xindex // 64)
    tmp0 = tl.load(in_ptr0 + (x3), xmask)
    tmp1 = tl.load(in_ptr0 + (x0 + (64*x2)), xmask, eviction_policy='evict_last')
    tmp2 = tl.load(in_ptr0 + (16 + x0 + (64*x2)), xmask, eviction_policy='evict_last')
    tmp4 = tl.load(in_ptr0 + (32 + x0 + (64*x2)), xmask, eviction_policy='evict_last')
    tmp6 = tl.load(in_ptr0 + (48 + x0 + (64*x2)), 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 + (x3), tmp8, xmask)
''', device_str='cuda')


# kernel path: runs/run_shard_7/inductor_cache/u5/cu5ua3ykbptipkew3i3zng4a7a4hy4f6xs547ovdooepce7uyfwz.py
# Topologically Sorted Source Nodes: [out_2], Original ATen: [aten._log_softmax]
# Source node to ATen node mapping:
#   out_2 => exp, log, sub_1, sum_1
# Graph fragment:
#   %exp : [num_users=1] = call_function[target=torch.ops.aten.exp.default](args = (%sub,), kwargs = {})
#   %sum_1 : [num_users=1] = call_function[target=torch.ops.aten.sum.dim_IntList](args = (%exp, [1], True), kwargs = {})
#   %log : [num_users=1] = call_function[target=torch.ops.aten.log.default](args = (%sum_1,), kwargs = {})
#   %sub_1 : [num_users=1] = call_function[target=torch.ops.aten.sub.Tensor](args = (%sub, %log), kwargs = {})
triton_poi_fused__log_softmax_1 = async_compile.triton('triton_poi_fused__log_softmax_1', '''
import triton
import triton.language as tl
from triton.compiler.compiler import AttrsDescriptor

from torch._inductor.runtime import triton_helpers, triton_heuristics
from torch._inductor.runtime.triton_helpers import libdevice, math as tl_math
from torch._inductor.runtime.hints import AutotuneHint, ReductionHint, TileHint, instance_descriptor, DeviceProperties

@triton_heuristics.pointwise(
    size_hints=[256], 
    filename=__file__,
    triton_meta={'signature': {0: '*fp32', 1: '*fp32', 2: 'i32'}, 'device': DeviceProperties(type='cuda', index=0, cc=80, major=8, regs_per_multiprocessor=65536, max_threads_per_multi_processor=2048, multi_processor_count=108), 'constants': {}, 'configs': [AttrsDescriptor(divisible_by_16=(0, 1, 2), equal_to_1=())]},
    inductor_meta={'autotune_hints': set(), 'kernel_name': 'triton_poi_fused__log_softmax_1', 'mutated_arg_names': [], 'no_x_dim': False, 'num_load': 5, 'num_reduction': 0, 'backend_hash': 'A9C866B4A14FD3277824029365D703C2427B2E685E54EC9B3EF4ADC8D1EEAC1D', 'are_deterministic_algorithms_enabled': False, 'assert_indirect_indexing': True, 'autotune_local_cache': True, 'autotune_pointwise': True, 'autotune_remote_cache': None, 'force_disable_caches': False, 'dynamic_scale_rblock': True, 'max_autotune': False, 'max_autotune_pointwise': False, 'min_split_scan_rblock': 256, 'spill_threshold': 16, 'store_cubin': False},
    min_elem_per_thread=0
)
@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
    x3 = xindex
    x0 = xindex % 16
    x2 = (xindex // 64)
    tmp0 = tl.load(in_ptr0 + (x3), xmask)
    tmp1 = tl.load(in_ptr0 + (x0 + (64*x2)), xmask, eviction_policy='evict_last')
    tmp3 = tl.load(in_ptr0 + (16 + x0 + (64*x2)), xmask, eviction_policy='evict_last')
    tmp6 = tl.load(in_ptr0 + (32 + x0 + (64*x2)), xmask, eviction_policy='evict_last')
    tmp9 = tl.load(in_ptr0 + (48 + x0 + (64*x2)), 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 + (x3), tmp13, xmask)
''', device_str='cuda')


async_compile.wait(globals())
del async_compile

def call(args):
    primals_1, primals_2, primals_3, primals_4, primals_5 = 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))
    assert_size_stride(primals_4, (4, 4), (4, 1))
    assert_size_stride(primals_5, (4, ), (1, ))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf0 = empty_strided_cuda((64, 4), (4, 1), torch.float32)
        # Topologically Sorted Source Nodes: [out], Original ATen: [aten.addmm]
        extern_kernels.addmm(primals_2, reinterpret_tensor(primals_3, (64, 4), (4, 1), 0), reinterpret_tensor(primals_1, (4, 4), (1, 4), 0), alpha=1, beta=1, out=buf0)
        del primals_1
        del primals_2
        buf1 = empty_strided_cuda((64, 4), (4, 1), torch.float32)
        # Topologically Sorted Source Nodes: [out_1], Original ATen: [aten.addmm]
        extern_kernels.addmm(primals_5, buf0, reinterpret_tensor(primals_4, (4, 4), (1, 4), 0), alpha=1, beta=1, out=buf1)
        del primals_5
        buf2 = empty_strided_cuda((4, 4, 4, 4), (64, 16, 4, 1), torch.float32)
        # Topologically Sorted Source Nodes: [out_2], Original ATen: [aten._log_softmax]
        stream0 = get_raw_stream(0)
        triton_poi_fused__log_softmax_0.run(buf1, buf2, 256, grid=grid(256), stream=stream0)
        buf3 = reinterpret_tensor(buf1, (4, 4, 4, 4), (64, 16, 4, 1), 0); del buf1  # reuse
        # Topologically Sorted Source Nodes: [out_2], Original ATen: [aten._log_softmax]
        triton_poi_fused__log_softmax_1.run(buf2, buf3, 256, grid=grid(256), stream=stream0)
        del buf2
    return (buf3, reinterpret_tensor(primals_3, (64, 4), (4, 1), 0), buf0, buf3, primals_4, )


def benchmark_compiled_module(times=10, repeat=10):
    from torch._dynamo.testing import rand_strided
    from torch._inductor.utils import print_performance
    primals_1 = rand_strided((4, 4), (4, 1), device='cuda:0', dtype=torch.float32)
    primals_2 = rand_strided((4, ), (1, ), device='cuda:0', dtype=torch.float32)
    primals_3 = rand_strided((4, 4, 4, 4), (64, 16, 4, 1), device='cuda:0', dtype=torch.float32)
    primals_4 = rand_strided((4, 4), (4, 1), device='cuda:0', dtype=torch.float32)
    primals_5 = rand_strided((4, ), (1, ), device='cuda:0', dtype=torch.float32)
    fn = lambda: call([primals_1, primals_2, primals_3, primals_4, primals_5])
    return print_performance(fn, times=times, repeat=repeat)


if __name__ == "__main__":
    from torch._inductor.wrapper_benchmark import compiled_module_main
    compiled_module_main('None', benchmark_compiled_module)

from torch.nn import Module
import random
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torch.utils.data import TensorDataset


class MyNetwork(Module):

    def __init__(self, size_input, size_hidden, size_output):
        """Create simple network"""
        super().__init__()
        self.layer_1 = nn.Linear(size_input, size_hidden)
        self.layer_2 = nn.Linear(size_hidden, size_output)
        self.softmax = nn.LogSoftmax(dim=1)

    def forward(self, X):
        """Forward through network"""
        out = self.layer_1(X)
        out = self.layer_2(out)
        out = self.softmax(out)
        return out


class VariableDataLoader(object):
    """Load data from variable length inputs

        Attributes
        ----------
        lengths : dict()
            Dictionary of input-length -> input samples

        index : boolean, default=False
            If True, also returns original index

        batch_size : int, default=1
            Size of each batch to output

        shuffle : boolean, default=True
            If True, shuffle the data randomly, each yielded batch contains
            only input items of the same length
        """

    def __init__(self, X, y, index=False, batch_size=1, shuffle=True):
        """Load data from variable length inputs

            Parameters
            ----------
            X : iterable of shape=(n_samples,)
                Input sequences
                Each item in iterable should be a sequence (of variable length)

            y : iterable of shape=(n_samples,)
                Labels corresponding to X

            index : boolean, default=False
                If True, also returns original index

            batch_size : int, default=1
                Size of each batch to output

            shuffle : boolean, default=True
                If True, shuffle the data randomly, each yielded batch contains
                only input items of the same length
            """
        self.lengths = dict()
        for i, (X_, y_) in enumerate(zip(X, y)):
            X_length, y_length, i_length = self.lengths.get(len(X_), (list(
                ), list(), list()))
            X_length.append(X_)
            y_length.append(y_)
            i_length.append(i)
            self.lengths[len(X_)] = X_length, y_length, i_length
        for k, v in self.lengths.items():
            self.lengths[k] = torch.as_tensor(v[0]), torch.as_tensor(v[1]
                ), torch.as_tensor(v[2])
        self.index = index
        self.batch_size = batch_size
        self.shuffle = shuffle
        self.reset()
        self.keys = set(self.data.keys())

    def reset(self):
        """Reset the VariableDataLoader"""
        self.done = set()
        self.data = {k: iter(DataLoader(TensorDataset(v[0], v[1], v[2]),
            batch_size=self.batch_size, shuffle=self.shuffle)) for k, v in
            self.lengths.items()}

    def __iter__(self):
        """Returns iterable of VariableDataLoader"""
        self.reset()
        return self

    def __next__(self):
        """Get next item of VariableDataLoader"""
        if self.done == self.keys:
            self.reset()
            raise StopIteration
        if self.shuffle:
            key = random.choice(list(self.keys - self.done))
        else:
            key = sorted(self.keys - self.done)[0]
        try:
            X_, y_, i = next(self.data.get(key))
            if self.index:
                item = X_, y_, i
            else:
                item = X_, y_
        except StopIteration:
            self.done.add(key)
            item = next(self)
        return item


class Module(nn.Module):
    """Extention of nn.Module that adds fit and predict methods
        Can be used for automatic training.

        Attributes
        ----------
        progress : Progress()
            Used to track progress of fit and predict methods
    """

    def __init__(self, *args, **kwargs):
        """Only calls super method nn.Module with given arguments."""
        super().__init__(*args, **kwargs)

    def fit(self, X, y, epochs=10, batch_size=32, learning_rate=0.01,
        criterion=nn.NLLLoss(), optimizer=optim.SGD, variable=False,
        verbose=True, **kwargs):
        """Train the module with given parameters

            Parameters
            ----------
            X : torch.Tensor
                Tensor to train with

            y : torch.Tensor
                Target tensor

            epochs : int, default=10
                Number of epochs to train with

            batch_size : int, default=32
                Default batch size to use for training

            learning_rate : float, default=0.01
                Learning rate to use for optimizer

            criterion : nn.Loss, default=nn.NLLLoss()
                Loss function to use

            optimizer : optim.Optimizer, default=optim.SGD
                Optimizer to use for training

            variable : boolean, default=False
                If True, accept inputs of variable length

            verbose : boolean, default=True
                If True, prints training progress

            Returns
            -------
            result : self
                Returns self
            """
        optimizer = optimizer(params=self.parameters(), lr=learning_rate)
        if variable:
            'cuda' if torch.cuda.is_available() else 'cpu'
            data = VariableDataLoader(X, y, batch_size=batch_size, shuffle=True
                )
        else:
            data = DataLoader(TensorDataset(X, y), batch_size=batch_size,
                shuffle=True)
        for epoch in range(1, epochs + 1):
            try:
                for X_, y_ in tqdm.tqdm(data, desc=
                    '[Epoch {:{width}}/{:{width}}]'.format(epoch, epochs,
                    width=len(str(epochs)))):
                    optimizer.zero_grad()
                    X_ = X_.clone().detach()
                    y_pred = self(X_)
                    loss = criterion(y_pred, y_)
                    loss.backward()
                    optimizer.step()
            except KeyboardInterrupt:
                None
                break
        return self

    def predict(self, X, batch_size=32, variable=False, verbose=True, **kwargs
        ):
        """Makes prediction based on input data X.
            Default implementation just uses the module forward(X) method,
            often the predict method will be overwritten to fit the specific
            needs of the module.

            Parameters
            ----------
            X : torch.Tensor
                Tensor from which to make prediction

            batch_size : int, default=32
                Batch size in which to predict items in X

            variable : boolean, default=False
                If True, accept inputs of variable length

            verbose : boolean, default=True
                If True, print progress of prediction

            Returns
            -------
            result : torch.Tensor
                Resulting prediction
            """
        with torch.no_grad():
            result = list()
            indices = torch.arange(len(X))
            if variable:
                indices = list()
                data = VariableDataLoader(X, torch.zeros(len(X)), index=
                    True, batch_size=batch_size, shuffle=False)
                for X_, y_, i in tqdm.tqdm(data, desc='Predicting'):
                    result.append(self(X_))
                    indices.append(i)
                indices = torch.cat(indices)
            else:
                for batch in tqdm.tqdm(range(0, X.shape[0], batch_size),
                    desc='Predicting'):
                    X_ = X[batch:batch + batch_size]
                    result.append(self(X_))
            return torch.cat(result)[indices]

    def fit_predict(self, X, y, epochs=10, batch_size=32, learning_rate=
        0.01, criterion=nn.NLLLoss, optimizer=optim.SGD, variable=False,
        verbose=True, **kwargs):
        """Train the module with given parameters

            Parameters
            ----------
            X : torch.Tensor
                Tensor to train with

            y : torch.Tensor
                Target tensor

            epochs : int, default=10
                Number of epochs to train with

            batch_size : int, default=32
                Default batch size to use for training

            learning_rate : float, default=0.01
                Learning rate to use for optimizer

            criterion : nn.Loss, default=nn.NLLLoss
                Loss function to use

            optimizer : optim.Optimizer, default=optim.SGD
                Optimizer to use for training

            variable : boolean, default=False
                If True, accept inputs of variable length

            verbose : boolean, default=True
                If True, prints training progress

            Returns
            -------
            result : torch.Tensor
                Resulting prediction
            """
        return self.fit(X, y, epochs, batch_size, learning_rate, criterion,
            optimizer, variable, verbose, **kwargs).predict(X, batch_size,
            variable, verbose, **kwargs)


def get_inputs():
    return [torch.rand([4, 4, 4, 4])]


def get_init_inputs():
    return [[], {'size_input': 4, 'size_hidden': 4, 'size_output': 4}]

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
from torch.nn import Module
import random
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torch.utils.data import TensorDataset
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
    x3 = xindex
    x0 = xindex % 16
    x2 = xindex // 64
    tmp0 = tl.load(in_ptr0 + x3, xmask)
    tmp1 = tl.load(in_ptr0 + (x0 + 64 * x2), xmask, eviction_policy=
        'evict_last')
    tmp2 = tl.load(in_ptr0 + (16 + x0 + 64 * x2), xmask, eviction_policy=
        'evict_last')
    tmp4 = tl.load(in_ptr0 + (32 + x0 + 64 * x2), xmask, eviction_policy=
        'evict_last')
    tmp6 = tl.load(in_ptr0 + (48 + x0 + 64 * x2), 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 + x3, 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
    x3 = xindex
    x0 = xindex % 16
    x2 = xindex // 64
    tmp0 = tl.load(in_ptr0 + x3, xmask)
    tmp1 = tl.load(in_ptr0 + (x0 + 64 * x2), xmask, eviction_policy=
        'evict_last')
    tmp3 = tl.load(in_ptr0 + (16 + x0 + 64 * x2), xmask, eviction_policy=
        'evict_last')
    tmp6 = tl.load(in_ptr0 + (32 + x0 + 64 * x2), xmask, eviction_policy=
        'evict_last')
    tmp9 = tl.load(in_ptr0 + (48 + x0 + 64 * x2), 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 + x3, tmp13, xmask)


def call(args):
    primals_1, primals_2, primals_3, primals_4, primals_5 = 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))
    assert_size_stride(primals_4, (4, 4), (4, 1))
    assert_size_stride(primals_5, (4,), (1,))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf0 = empty_strided_cuda((64, 4), (4, 1), torch.float32)
        extern_kernels.addmm(primals_2, reinterpret_tensor(primals_3, (64, 
            4), (4, 1), 0), reinterpret_tensor(primals_1, (4, 4), (1, 4), 0
            ), alpha=1, beta=1, out=buf0)
        del primals_1
        del primals_2
        buf1 = empty_strided_cuda((64, 4), (4, 1), torch.float32)
        extern_kernels.addmm(primals_5, buf0, reinterpret_tensor(primals_4,
            (4, 4), (1, 4), 0), alpha=1, beta=1, out=buf1)
        del primals_5
        buf2 = 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)](buf1, buf2, 256, XBLOCK=
            128, num_warps=4, num_stages=1)
        buf3 = reinterpret_tensor(buf1, (4, 4, 4, 4), (64, 16, 4, 1), 0)
        del buf1
        triton_poi_fused__log_softmax_1[grid(256)](buf2, buf3, 256, XBLOCK=
            128, num_warps=4, num_stages=1)
        del buf2
    return buf3, reinterpret_tensor(primals_3, (64, 4), (4, 1), 0
        ), buf0, buf3, primals_4


class MyNetworkNew(Module):

    def __init__(self, size_input, size_hidden, size_output):
        """Create simple network"""
        super().__init__()
        self.layer_1 = nn.Linear(size_input, size_hidden)
        self.layer_2 = nn.Linear(size_hidden, size_output)
        self.softmax = nn.LogSoftmax(dim=1)

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


class VariableDataLoader(object):
    """Load data from variable length inputs

        Attributes
        ----------
        lengths : dict()
            Dictionary of input-length -> input samples

        index : boolean, default=False
            If True, also returns original index

        batch_size : int, default=1
            Size of each batch to output

        shuffle : boolean, default=True
            If True, shuffle the data randomly, each yielded batch contains
            only input items of the same length
        """

    def __init__(self, X, y, index=False, batch_size=1, shuffle=True):
        """Load data from variable length inputs

            Parameters
            ----------
            X : iterable of shape=(n_samples,)
                Input sequences
                Each item in iterable should be a sequence (of variable length)

            y : iterable of shape=(n_samples,)
                Labels corresponding to X

            index : boolean, default=False
                If True, also returns original index

            batch_size : int, default=1
                Size of each batch to output

            shuffle : boolean, default=True
                If True, shuffle the data randomly, each yielded batch contains
                only input items of the same length
            """
        self.lengths = dict()
        for i, (X_, y_) in enumerate(zip(X, y)):
            X_length, y_length, i_length = self.lengths.get(len(X_), (list(
                ), list(), list()))
            X_length.append(X_)
            y_length.append(y_)
            i_length.append(i)
            self.lengths[len(X_)] = X_length, y_length, i_length
        for k, v in self.lengths.items():
            self.lengths[k] = torch.as_tensor(v[0]), torch.as_tensor(v[1]
                ), torch.as_tensor(v[2])
        self.index = index
        self.batch_size = batch_size
        self.shuffle = shuffle
        self.reset()
        self.keys = set(self.data.keys())

    def reset(self):
        """Reset the VariableDataLoader"""
        self.done = set()
        self.data = {k: iter(DataLoader(TensorDataset(v[0], v[1], v[2]),
            batch_size=self.batch_size, shuffle=self.shuffle)) for k, v in
            self.lengths.items()}

    def __iter__(self):
        """Returns iterable of VariableDataLoader"""
        self.reset()
        return self

    def __next__(self):
        """Get next item of VariableDataLoader"""
        if self.done == self.keys:
            self.reset()
            raise StopIteration
        if self.shuffle:
            key = random.choice(list(self.keys - self.done))
        else:
            key = sorted(self.keys - self.done)[0]
        try:
            X_, y_, i = next(self.data.get(key))
            if self.index:
                item = X_, y_, i
            else:
                item = X_, y_
        except StopIteration:
            self.done.add(key)
            item = next(self)
        return item


class Module(nn.Module):
    """Extention of nn.Module that adds fit and predict methods
        Can be used for automatic training.

        Attributes
        ----------
        progress : Progress()
            Used to track progress of fit and predict methods
    """

    def __init__(self, *args, **kwargs):
        """Only calls super method nn.Module with given arguments."""
        super().__init__(*args, **kwargs)

    def fit(self, X, y, epochs=10, batch_size=32, learning_rate=0.01,
        criterion=nn.NLLLoss(), optimizer=optim.SGD, variable=False,
        verbose=True, **kwargs):
        """Train the module with given parameters

            Parameters
            ----------
            X : torch.Tensor
                Tensor to train with

            y : torch.Tensor
                Target tensor

            epochs : int, default=10
                Number of epochs to train with

            batch_size : int, default=32
                Default batch size to use for training

            learning_rate : float, default=0.01
                Learning rate to use for optimizer

            criterion : nn.Loss, default=nn.NLLLoss()
                Loss function to use

            optimizer : optim.Optimizer, default=optim.SGD
                Optimizer to use for training

            variable : boolean, default=False
                If True, accept inputs of variable length

            verbose : boolean, default=True
                If True, prints training progress

            Returns
            -------
            result : self
                Returns self
            """
        optimizer = optimizer(params=self.parameters(), lr=learning_rate)
        if variable:
            'cuda' if torch.cuda.is_available() else 'cpu'
            data = VariableDataLoader(X, y, batch_size=batch_size, shuffle=True
                )
        else:
            data = DataLoader(TensorDataset(X, y), batch_size=batch_size,
                shuffle=True)
        for epoch in range(1, epochs + 1):
            try:
                for X_, y_ in tqdm.tqdm(data, desc=
                    '[Epoch {:{width}}/{:{width}}]'.format(epoch, epochs,
                    width=len(str(epochs)))):
                    optimizer.zero_grad()
                    X_ = X_.clone().detach()
                    y_pred = self(X_)
                    loss = criterion(y_pred, y_)
                    loss.backward()
                    optimizer.step()
            except KeyboardInterrupt:
                None
                break
        return self

    def predict(self, X, batch_size=32, variable=False, verbose=True, **kwargs
        ):
        """Makes prediction based on input data X.
            Default implementation just uses the module forward(X) method,
            often the predict method will be overwritten to fit the specific
            needs of the module.

            Parameters
            ----------
            X : torch.Tensor
                Tensor from which to make prediction

            batch_size : int, default=32
                Batch size in which to predict items in X

            variable : boolean, default=False
                If True, accept inputs of variable length

            verbose : boolean, default=True
                If True, print progress of prediction

            Returns
            -------
            result : torch.Tensor
                Resulting prediction
            """
        with torch.no_grad():
            result = list()
            indices = torch.arange(len(X))
            if variable:
                indices = list()
                data = VariableDataLoader(X, torch.zeros(len(X)), index=
                    True, batch_size=batch_size, shuffle=False)
                for X_, y_, i in tqdm.tqdm(data, desc='Predicting'):
                    result.append(self(X_))
                    indices.append(i)
                indices = torch.cat(indices)
            else:
                for batch in tqdm.tqdm(range(0, X.shape[0], batch_size),
                    desc='Predicting'):
                    X_ = X[batch:batch + batch_size]
                    result.append(self(X_))
            return torch.cat(result)[indices]

    def fit_predict(self, X, y, epochs=10, batch_size=32, learning_rate=
        0.01, criterion=nn.NLLLoss, optimizer=optim.SGD, variable=False,
        verbose=True, **kwargs):
        """Train the module with given parameters

            Parameters
            ----------
            X : torch.Tensor
                Tensor to train with

            y : torch.Tensor
                Target tensor

            epochs : int, default=10
                Number of epochs to train with

            batch_size : int, default=32
                Default batch size to use for training

            learning_rate : float, default=0.01
                Learning rate to use for optimizer

            criterion : nn.Loss, default=nn.NLLLoss
                Loss function to use

            optimizer : optim.Optimizer, default=optim.SGD
                Optimizer to use for training

            variable : boolean, default=False
                If True, accept inputs of variable length

            verbose : boolean, default=True
                If True, prints training progress

            Returns
            -------
            result : torch.Tensor
                Resulting prediction
            """
        return self.fit(X, y, epochs, batch_size, learning_rate, criterion,
            optimizer, variable, verbose, **kwargs).predict(X, batch_size,
            variable, verbose, **kwargs)
