keras/src/backend/numpy/numpy.py/count_nonzero
def count_nonzero(x, axis=None):
    axis = standardize_axis_for_numpy(axis)
    # np.count_nonzero will return python int when axis=None, so we need
    # to convert_to_tensor
    return convert_to_tensor(np.count_nonzero(x, axis=axis)).astype("int32")