def make_input_output_pools(
    *, pool_size, eval_pt_func, input_filter_func, output_filter_func
):
    """
    Make input and output pools for benchmarking. The rationale is avoiding retrieving the same input from the device cache for fair perf assessment.
    Parameters
    ----------
    pool_size : int
        The size of the pool.
    eval_pt_func : callable
        A callable that returns a dict of inputs and outputs.
    input_filter_func : callable
        A callable that takes a key and a value and returns True if the key-value pair from the `eval_pt_func` result should be included in the input pool.
    output_filter_func : callable
        A callable that takes a key and a value and returns True if the key-value pair from the `eval_pt_func` result should be included in the output pool.
    Returns
    -------
    inputs_pool : List[Dict[str, torch.Tensor]]
        A list of inputs to pass into Model.RunWithTensors.
    outputs_pool : List[Dict[str, torch.Tensor]]
        A list of outputs to pass into Model.RunWithTensors.
    """
    return zip(
        *[
            [
                {k: v for k, v in d.items() if input_filter_func(k, v)},
                {
                    k: torch.empty_like(v)
                    for k, v in d.items()
                    if output_filter_func(k, v)
                },
            ]
            for d in [eval_pt_func() for _ in range(pool_size)]
        ]
    )

    Make input and output pools for benchmarking. The rationale is avoiding retrieving the same input from the device cache for fair perf assessment.
    Parameters
    ----------
    pool_size : int
        The size of the pool.
    eval_pt_func : callable
        A callable that returns a dict of inputs and outputs.
    input_filter_func : callable
        A callable that takes a key and a value and returns True if the key-value pair from the `eval_pt_func` result should be included in the input pool.
    output_filter_func : callable
        A callable that takes a key and a value and returns True if the key-value pair from the `eval_pt_func` result should be included in the output pool.
    Returns
    -------
    inputs_pool : List[Dict[str, torch.Tensor]]
        A list of inputs to pass into Model.RunWithTensors.
    outputs_pool : List[Dict[str, torch.Tensor]]
        A list of outputs to pass into Model.RunWithTensors.
    
python/aitemplate/testing/benchmark_ait.py