def METHOD_NAME(cls, ctx, op):
    (x,), device_id, xp = as_same_device(
        [ctx[c.key] for c in op.inputs], op.device, ret_extra=True
    )
    axis = cls.get_arg_axis(op.axis, op.inputs[0].ndim)
    keepdims = op.keepdims
    with device(device_id):
        nz = xp.count_nonzero(x, axis=axis)
        if keepdims:
            slcs = [slice(None)] * op.inputs[0].ndim
            for ax in op.axis:
                slcs[ax] = np.newaxis
            nz = xp.asarray(nz)[tuple(slcs)]
        ctx[op.outputs[0].key] = nz