def METHOD_NAME(cls, ctx, op: "LearnCountNonzero"):
    axis = op.axis
    X = ctx[op.inputs[0].key]
    sample_weight = (
        ctx[op.sample_weight.key] if op.sample_weight is not None else None
    )
    # We rely here on the fact that np.diff(Y.indptr) for a CSR
    # will return the number of nonzero entries in each row.
    # A bincount over Y.indices will return the number of nonzeros
    # in each column. See ``csr_matrix.getnnz`` in scipy >= 0.14.
    if axis is None:
        if sample_weight is None:
            res = X.nnz
        else:
            res = np.dot(np.diff(X.indptr), sample_weight)
    elif axis == 1:
        out = np.diff(X.indptr)
        if sample_weight is None:
            # astype here is for consistency with axis=0 dtype
            res = out.astype("intp")
        else:
            res = out * sample_weight
    else:
        if sample_weight is None:
            res = np.bincount(X.indices, minlength=X.shape[1])
        else:
            weights = np.repeat(sample_weight, np.diff(X.indptr))
            res = np.bincount(X.indices, minlength=X.shape[1], weights=weights)
    if np.isscalar(res):
        res = np.array([res])
    out_shape = op.outputs[0].shape
    if any(np.isnan(s) for s in out_shape):
        new_shape = list(out_shape)
        for i, s in enumerate(out_shape):
            if np.isnan(s):
                new_shape[i] = -1
        out_shape = tuple(new_shape)
    ctx[op.outputs[0].key] = res.reshape(out_shape)