code: def _sklearn_calculate_mi_wrapper(
      feature: np.ndarray,
      label: np.ndarray,
      discrete_features: Union[bool, Sequence[bool]],
      seed: int,
      copy: bool,
      is_label_categorical: bool,
  ) -> Optional[np.ndarray]:
      """Wraps sklearn calculate mi with some additional validation.
  
      Args:
        feature: The features.
        label: The labels.
        discrete_features: If bool, then determines whether to consider all
          features discrete or continuous. If array, then it should be either a
          boolean mask with shape (n_features,) or array with indices of discrete
          features.
        seed: Determines random number generation for adding small noise to
          continuous variables in order to remove repeated values. Pass an int for
          reproducible results across multiple function calls.
        copy: Whether to make a copy of the given data. If set to False, the
          initial data will be overwritten.
        is_label_categorical: True if the label is a categorical feature.
  
      Returns:
        A numpy array of mutual information of each feature. Will return None if MI
        cannot be calculated.
      """
      if is_label_categorical:
          calc_mi_fn = mutual_info_classif
      else:
          # Skip if sample size is smaller than number of required neighbors plus
          # itself.
          if len(feature) <= _KNN_N_NEIGHBORS:
              return None
          calc_mi_fn = mutual_info_regression
  
      return calc_mi_fn(
          feature,
          label,
          discrete_features=discrete_features,
          n_neighbors=_KNN_N_NEIGHBORS,
          copy=copy,
          random_state=seed,
      )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import sys
  from typing import Dict, List, Optional, Sequence, Set, Text, Union
  import numpy
  import pandas
  import pyarrow
  from tensorflow_data_validation import types
  from tensorflow_data_validation.arrow import arrow_util
  from tensorflow_data_validation.statistics.generators import partitioned_stats_generator
  from tensorflow_data_validation.utils import schema_util
  from tensorflow_data_validation.utils import stats_util
  from tfx_bsl.arrow import array_util
  from tensorflow_metadata.proto.v0 import schema_pb2
  from tensorflow_metadata.proto.v0 import statistics_pb2
label_desc: Use image node_id 2 for calling a global function with example usage: _sklearn_calculate_mi_wrapper(feature, label, discrete_features, seed, copy, is_label_categorical) and returns: calc_mi_fn, None