`torch.count_nonzero` fails with `axis` tuple and `keepdims=True`

[mdhaber]

from array_api_compat import torch as xp
x = xp.asarray([1])
xp.count_nonzero(x, axis=(0,), keepdims=True)
# TypeError: unsqueeze(): argument 'dim' (position 1) must be int, not tuple