def METHOD_NAME(cls, ctx, op):
    a = ctx[op.inputs[0].key]
    (inp,), device_id, xp = as_same_device([a], device=op.device, ret_extra=True)
    with device(device_id):
        ctx[op.outputs[0].key] = xp.count_nonzero(inp, axis=op.axis)