def METHOD_NAME(args) -> Tuple[list, list, List[bool], TreeSpec]:
    flat_args, spec = tree_flatten(args)
    grad_args = []
    other_args = []
    grad_flags = []
    for arg in flat_args:
        flag = _is_grad_tensor(arg)
        grad_flags.append(flag)
        if flag:
            grad_args.append(arg)
        else:
            other_args.append(arg)
    assert len(grad_args) > 0
    return grad_args, other_args, grad_flags, spec