def METHOD_NAME(rank):
    if tp_env.mode == '2d':
        check_2d_tensor_parallel_rank(rank)
    elif tp_env == '2.5d':
        check_2p5d_tensor_parallel_rank(rank)
    elif tp_env == '3d':
        check_3d_tensor_parallel_rank(rank)