keras/src/ops/numpy.py/CountNonzero/__init__
class CountNonzero:    def __init__(self, axis=None):
        super().__init__()
        if isinstance(axis, int):
            self.axis = (axis,)
        else:
            self.axis = axis