from typing import Any, List, Tuple

import torch

from vllm.distributed.parallel_state import TensorMetadata
from vllm_mlu.mlu_hijack_utils import MluHijackObject
from vllm_mlu.model_executor.custom_model.custom import CustomForCausalLM

def vllm__module_executor__models__custom_model__CustomForCausalLM__get_intermediate_tensor_metadata(
    self,
    batch_size: int,
    dtype: torch.dtype,
    device: torch.device,
) -> List[Tuple[str, Any]]:
    metadata_list: List[Tuple[str, Any]] = []
    size = torch.Size([batch_size, self.config.hidden_size])
    metadata_list.append(("hidden_states", TensorMetadata(device.type, dtype, size)))
    metadata_list.append(("residual", None))
    return metadata_list

MluHijackObject.apply_hijack(
    CustomForCausalLM,
    "get_intermediate_tensor_metadata",
    vllm__module_executor__models__custom_model__CustomForCausalLM__get_intermediate_tensor_metadata
)
