import inspect
import torch


def apply_chunking_to_forward(forward_fn, chunk_size: int, chunk_dim: int, *input_tensors):
    """ 将运算分割为多个chunk计算来节省显存 当chunk=0，不分块执行
    参数:
        forward_fn (:obj:`Callable[..., torch.Tensor]`):
        chunk_size : `num_chunks = len(input_tensors[0]) / chunk_size`.
        chunk_dim (:obj:`int`): chunk的维度
        input_tensors :tuple(tensors)
    """
    assert len(input_tensors) > 0, f"{input_tensors} has to be a tuple/list of tensors"
    tensor_shape = input_tensors[0].shape[chunk_dim]
    # 每个tensor chunk的维度需相等
    assert all(input_tensor.shape[chunk_dim] == tensor_shape for input_tensor in input_tensors), "All input tenors have to be of the same shape"
    # inspect.signature exist since python 3.5 and is a python method -> no problem with backward compatibility
    num_args_in_forward_chunk_fn = len(inspect.signature(forward_fn).parameters)
    if num_args_in_forward_chunk_fn != len(input_tensors):
        raise ValueError(f"forward_chunk_fn expects {num_args_in_forward_chunk_fn} arguments, but only {len(input_tensors)} input tensors are given")

    if chunk_size > 0:
        if input_tensors[0].shape[chunk_dim] % chunk_size != 0:
            raise ValueError(f"The dimension to be chunked {input_tensors[0].shape[chunk_dim]} has to be a multiple of the chunk size {chunk_size}")
        num_chunks = input_tensors[0].shape[chunk_dim] // chunk_size
        # 将每个tensor按chunk_dim分块
        input_tensors_chunks = tuple(input_tensor.chunk(num_chunks, dim=chunk_dim) for input_tensor in input_tensors)
        # 执行forward到每个分块
        output_chunks = tuple(forward_fn(*input_tensors_chunk) for input_tensors_chunk in zip(*input_tensors_chunks))
        # 合并结果
        return torch.cat(output_chunks, dim=chunk_dim)

    return forward_fn(*input_tensors)