keras/src/backend/torch/numpy.py/count_nonzero
def count_nonzero(x, axis=None):
    x = convert_to_tensor(x)
    if axis == () or axis == []:
        # Torch handles the empty axis case differently from numpy.
        return cast(torch.ne(x, 0), "int32")
    return cast(torch.count_nonzero(x, dim=axis).T, "int32")