def METHOD_NAME(flat_args):
    for a in flat_args:
        if not (isinstance(a, (torch.Tensor, str, int, float, bool)) or a is None):
            raise AssertionError(
                f"Only Tensors or scalars are supported as pytree flattened inputs, got: {a}"
            )