code: def count_nonzero(X, axis: Optional[int] = None, sample_weight=None):
      """A variant of X.getnnz() with extension to weighting on axis 0
  
      Useful in efficiently calculating multilabel metrics.
  
      Parameters
      ----------
      X : CSR sparse matrix of shape (n_samples, n_labels)
          Input data.
  
      axis : None, 0 or 1
          The axis on which the data is aggregated.
  
      sample_weight : array-like of shape (n_samples,), default=None
          Weight for each row of X.
      """
      if axis == -1:
          axis = 1
      elif axis == -2:
          axis = 0
      if axis is not None and axis not in (0, 1):
          raise ValueError(f"Unsupported axis: {axis}")
  
      X = mt.asarray(X)
      if sample_weight is not None:
          sample_weight = mt.asarray(sample_weight)
  
      op = LearnCountNonzero(axis=axis)
      return op(X, sample_weight=sample_weight)
  
file_path: mars/mars/learn/utils/sparsefuncs.py
imports: from typing import Optional
  import numpy
  from ...None import opcodes, tensor
  from ...core import OutputType, recursive_tile
  from ...serialization.serializables import Int16Field, ReferenceField
  from ...utils import has_unknown_shape
  from ..operands import LearnOperand, LearnOperandMixin
label_desc: Use image node_id 1 for calling a global function with example usage: count_nonzero(X, axis, sample_weight) and returns: op