# AOT ID: ['0_inference']
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_6/inductor_cache/t4/ct4abpxorr7ywpffawu2dyuam53y6icpnxs5d4cvirqopww6wnuh.py
# Topologically Sorted Source Nodes: [concatinated_batch], Original ATen: [aten.cat]
# Source node to ATen node mapping:
#   concatinated_batch => cat
# Graph fragment:
#   %cat : [num_users=1] = call_function[target=torch.ops.aten.cat.default](args = ([%arg0_1, %arg1_1], 2), kwargs = {})
triton_poi_fused_cat_0 = async_compile.triton('triton_poi_fused_cat_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=[512], 
    filename=__file__,
    triton_meta={'signature': {0: '*fp32', 1: '*fp32', 2: '*fp32', 3: '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, 3), equal_to_1=())]},
    inductor_meta={'autotune_hints': set(), 'kernel_name': 'triton_poi_fused_cat_0', 'mutated_arg_names': [], 'no_x_dim': False, 'num_load': 2, '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_cat_0(in_ptr0, in_ptr1, out_ptr0, xnumel, XBLOCK : tl.constexpr):
    xnumel = 512
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x1 = (xindex // 4) % 8
    x0 = xindex % 4
    x2 = (xindex // 32)
    x3 = xindex
    tmp0 = x1
    tmp1 = tl.full([1], 0, tl.int64)
    tmp2 = tmp0 >= tmp1
    tmp3 = tl.full([1], 4, tl.int64)
    tmp4 = tmp0 < tmp3
    tmp5 = tl.load(in_ptr0 + (x0 + (4*x1) + (16*x2)), tmp4 & xmask, other=0.0)
    tmp6 = tmp0 >= tmp3
    tmp7 = tl.full([1], 8, tl.int64)
    tmp8 = tmp0 < tmp7
    tmp9 = tl.load(in_ptr1 + (x0 + (4*((-4) + x1)) + (16*x2)), tmp6 & xmask, other=0.0)
    tmp10 = tl.where(tmp4, tmp5, tmp9)
    tl.store(out_ptr0 + (x3), tmp10, xmask)
''', device_str='cuda')


async_compile.wait(globals())
del async_compile

def call(args):
    arg0_1, arg1_1 = args
    args.clear()
    assert_size_stride(arg0_1, (4, 4, 4, 4), (64, 16, 4, 1))
    assert_size_stride(arg1_1, (4, 4, 4, 4), (64, 16, 4, 1))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf0 = empty_strided_cuda((4, 4, 8, 4), (128, 32, 4, 1), torch.float32)
        # Topologically Sorted Source Nodes: [concatinated_batch], Original ATen: [aten.cat]
        stream0 = get_raw_stream(0)
        triton_poi_fused_cat_0.run(arg0_1, arg1_1, buf0, 512, grid=grid(512), stream=stream0)
        del arg0_1
        del arg1_1
    return (buf0, )


def benchmark_compiled_module(times=10, repeat=10):
    from torch._dynamo.testing import rand_strided
    from torch._inductor.utils import print_performance
    arg0_1 = rand_strided((4, 4, 4, 4), (64, 16, 4, 1), device='cuda:0', dtype=torch.float32)
    arg1_1 = rand_strided((4, 4, 4, 4), (64, 16, 4, 1), device='cuda:0', dtype=torch.float32)
    fn = lambda: call([arg0_1, arg1_1])
    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)

import torch
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 math
import numpy as np
from abc import ABC
from abc import abstractmethod
from abc import abstractproperty
from torch import nn
from enum import Enum
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, out_ptr0, xnumel, XBLOCK: tl.
    constexpr):
    xnumel = 512
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x1 = xindex // 4 % 8
    x0 = xindex % 4
    x2 = xindex // 32
    x3 = xindex
    tmp0 = x1
    tl.full([1], 0, tl.int64)
    tmp3 = tl.full([1], 4, tl.int64)
    tmp4 = tmp0 < tmp3
    tmp5 = tl.load(in_ptr0 + (x0 + 4 * x1 + 16 * x2), tmp4 & xmask, other=0.0)
    tmp6 = tmp0 >= tmp3
    tl.full([1], 8, tl.int64)
    tmp9 = tl.load(in_ptr1 + (x0 + 4 * (-4 + x1) + 16 * x2), tmp6 & xmask,
        other=0.0)
    tmp10 = tl.where(tmp4, tmp5, tmp9)
    tl.store(out_ptr0 + x3, tmp10, xmask)


def call(args):
    arg0_1, arg1_1 = args
    args.clear()
    assert_size_stride(arg0_1, (4, 4, 4, 4), (64, 16, 4, 1))
    assert_size_stride(arg1_1, (4, 4, 4, 4), (64, 16, 4, 1))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf0 = empty_strided_cuda((4, 4, 8, 4), (128, 32, 4, 1), torch.float32)
        get_raw_stream(0)
        triton_poi_fused_cat_0[grid(512)](arg0_1, arg1_1, buf0, 512, XBLOCK
            =256, num_warps=4, num_stages=1)
        del arg0_1
        del arg1_1
    return buf0,


def tensor_to_numpy(tensor):
    return tensor.detach().cpu().numpy()


class MLPParamHandler(ABC):

    def __init__(self) ->None:
        """Interface for parameter handler. For Algorithms that require data on past model parameters this module handels this data such that is it is in the right format according to the specific algorithm.
        """
        super().__init__()

    @abstractmethod
    def get_policy_replay_data(self, model: 'torch.nn.Module'):
        """Function to extract replay data from a module in a format that it can be saved to the replay buffer.

        Args:
            model (torch.nn.Module): Module to extract data from.
        """
        ...

    @abstractmethod
    def get_policy_critic_data(self, model: 'torch.nn.Module'):
        """Function to extract data from a policy that the critic requires to evaluate it.

        Args:
            model (torch.nn.Module): Module to extract data from.
        """
        ...

    @abstractmethod
    def format_replay_buffer_data(self, **kwargs):
        """Function to format data from the replay buffer such that it can be used as input to the critic afterwards.
        """
        ...

    @abstractproperty
    def replay_data_keys(self):
        """Keys used to save data to the replay buffer.
        """
        ...

    @abstractproperty
    def replay_data_info(self):
        """Description of the data needed to initialize the replay buffer.
        """
        ...


class FlatParamHandler(MLPParamHandler):

    def __init__(self, example_policy: 'torch.nn.Module') ->None:
        """Parameter handler that simply takes all parameters flattens them and saves them in one vector.

        Args:
            example_policy (torch.nn.Module): Example policy network to acquire shape of data.
        """
        super().__init__()
        self._replay_data_keys = [constants.DATA_PARAMETERS]
        self._replay_data_info = {self._replay_data_keys[0]: {'shape': self
            .get_policy_critic_data(example_policy).shape}}

    def get_policy_replay_data(self, model: 'torch.nn.Module'):
        return {self._replay_data_keys[0]: tensor_to_numpy(torch.nn.utils.
            parameters_to_vector(model.parameters())).reshape(1, -1)}

    def get_policy_critic_data(self, model: 'torch.nn.Module'):
        return torch.nn.utils.parameters_to_vector(model.parameters()).reshape(
            1, -1)

    def format_replay_buffer_data(self, **kwargs):
        return kwargs[self._replay_data_keys[0]]

    @property
    def replay_data_keys(self):
        return self._replay_data_keys

    @property
    def replay_data_info(self):
        return self._replay_data_info


class NamedParamHandler(MLPParamHandler):

    def __init__(self, example_policy: 'torch.nn.Module') ->None:
        """Parameter handler that saves parameters in a dictionary shape such that the parameters are saved in a similar format of how they are used in the actual module. Useful if the Parameters are later reused similarly to how the are used within the module they are extracted from.

        Args:
            example_policy (torch.nn.Module): Example policy network to acquire structure of module and according dictionary.
        """
        super().__init__()
        actor_parameter_dict = self.get_policy_critic_data(example_policy)
        self._replay_data_keys = actor_parameter_dict.keys()
        self._replay_data_info = {key: {'shape': actor_parameter_dict[key].
            shape[1:]} for key in self._replay_data_keys}

    def get_policy_replay_data(self, model: 'torch.nn.Module'):
        batched_param_dict = self.get_policy_critic_data(model)
        return {key: tensor_to_numpy(value) for key, value in
            batched_param_dict.items()}

    def get_policy_critic_data(self, model: 'torch.nn.Module'):
        param_dict = dict(model.named_parameters())
        return {key: torch.unsqueeze(tensor, dim=0) for key, tensor in
            param_dict.items()}

    def format_replay_buffer_data(self, **kwargs):
        return {key: kwargs[key] for key in self._replay_data_keys}

    @property
    def replay_data_info(self):
        return self._replay_data_info

    @property
    def replay_data_keys(self):
        return self._replay_data_keys


class StateActionHandler(MLPParamHandler):

    def __init__(self, num_state_action_pairs: 'int', episode_length: 'int',
        rollout_handler: 'RolloutHandler') ->None:
        """Parameter handler that does not actually use parameters but rather state action pairs as representation of policies.

        Args:
            num_state_action_pairs (int): Number of state action pairs used as a representation for a policy.
            episode_length (int): Maximal time steps of a episode used for representation purposes.
            rollout_handler (RolloutHandler): Rollout handler used to execute rollouts when needed.
        """
        super().__init__()
        self.rollout_handler = rollout_handler
        self.num_state_action_pairs = num_state_action_pairs
        self.episode_length = episode_length
        self._replay_data_keys = [constants.DATA_OBSERVATIONS, constants.
            DATA_ACTIONS]
        self._replay_data_info = {constants.DATA_OBSERVATIONS: {'shape': (
            episode_length, *rollout_handler.environment_handler.
            exploration_environment.observation_space.shape)}, constants.
            DATA_ACTIONS: {'shape': (episode_length, *rollout_handler.
            environment_handler.exploration_environment.action_space.shape)}}

    def get_policy_replay_data(self, model: 'torch.nn.Module'):
        return {}

    def get_policy_critic_data(self, model: 'torch.nn.Module'):
        rollout_data = self.rollout_handler.update_rollout(policy=model,
            extraction_keys=[constants.DATA_OBSERVATIONS, constants.
            DATA_ACTIONS])
        sampled_states, sampeled_actions = self.format_replay_buffer_data(**
            rollout_data)
        return sampled_states, sampeled_actions

    def format_replay_buffer_data(self, **kwargs):
        states = kwargs[constants.DATA_OBSERVATIONS]
        actions = kwargs[constants.DATA_ACTIONS]
        sampled_states, sampeled_actions = self._sample_state_action_paris(
            states, actions)
        return sampled_states, sampeled_actions

    def _sample_state_action_paris(self, states, actions):
        """To make sure the number of state actions paris is always the same, this function sub samples the desired amount from a state action batch. This also acts as a kind of data augmentation as the representation of a single policy will consist of different state action pairs if called multiple times.

        Args:
            states (np.ndarray): Batch of states to sub sample from.
            actions (np.ndarray): Batch of actions (according to states) to sum sample from.

        Returns:
            sampled_states (np.ndarray): Sub sampled states
            sampeled_actions (np.ndarray): Sub sampled actions
        """
        sample_id = np.random.choice(range(self.episode_length), size=self.
            num_state_action_pairs, replace=False)
        sampled_states = states[:, sample_id]
        sampeled_actions = actions[:, sample_id]
        return sampled_states, sampeled_actions

    @property
    def replay_data_info(self):
        return self._replay_data_info

    @property
    def replay_data_keys(self):
        return self._replay_data_keys


class ParameterFormat(Enum):
    FlatParameters = FlatParamHandler
    NamedParameters = NamedParamHandler
    StateAction = StateActionHandler


class MLPEmbeddingNetwork(nn.Module):

    def __init__(self):
        super(MLPEmbeddingNetwork, self).__init__()

    @abstractproperty
    def embedding_size(self) ->int:
        ...

    @abstractproperty
    def input_type(self) ->ParameterFormat:
        ...


class StateActionEmbeddingNew(MLPEmbeddingNetwork):

    def __init__(self, num_state_action_pairs: 'int', observation_space:
        'gym.Space', action_space: 'gym.Space'):
        super(StateActionEmbeddingNew, self).__init__()
        self.observation_shape = observation_space.shape
        self.action_shape = action_space.shape
        self.num_state_action_pairs = num_state_action_pairs
        self._embedding_size = num_state_action_pairs * (math.prod(self.
            observation_shape) + math.prod(self.action_shape))

    @property
    def embedding_size(self) ->int:
        return self._embedding_size

    @property
    def input_type(self) ->ParameterFormat:
        return ParameterFormat.StateAction

    def forward(self, input_0, input_1):
        arg0_1 = input_0
        arg1_1 = input_1
        output = call([arg0_1, arg1_1])
        return output[0]
