class_docstr: Computes Mutual Information(MI) between each feature and the label.
  
  This statistic is the estimated Adjusted Mutual Information(AMI) between all
  features and the label. AMI prevents overestimation of MI for high entropy
  features. It is defined as MI(feature, labels) - MI(feature, shuffle(labels)).
  
  To use this statistic, use the `NonStreamingCustomStatsGenerator`. This
  generator can then be specified in the `stats_options` when calling
  `GenerateStatistics`.
  
  Example usage:
  ```
  generator = partitioned_stats_generator.NonStreamingCustomStatsGenerator(
    MutualInformation(...))
  ```
docstr: Initializes MutualInformation.
  
  Args:
    label_feature: The key used to identify labels in the ExampleBatch.
    schema: An optional schema describing the the dataset. Either a schema or
      a list of categorical and multivalent features must be provided.
    max_encoding_length: An int value to specify the maximum length of
      encoding to represent a feature value.
    seed: An int value to seed the RNG used in MI computation.
    multivalent_features: An optional set of features that are multivalent.
    categorical_features: An optional set of the features that are
      categorical.
    features_to_ignore: An optional set of features that should be ignored by
      the mutual information calculation.
    normalize_by_max: If True, AMI values are normalized to a range 0 to 1 by
      dividing by the maximum possible information AMI(Y, Y).
    allow_invalid_partitions: If True, generator tolerates input partitions
      that are invalid (e.g. size of partion is < the k for the KNN), where
      invalid partitions return no stats. The min_partitions_stat_presence arg
      to PartitionedStatisticsAnalyzer controls how many partitions may be
      invalid while still reporting the metric.
    custom_stats_key: A string that determines the key used in the custom
      statistic. This defaults to `_ADJUSTED_MUTUAL_INFORMATION_KEY`.
    column_partitions: If > 1, self.partitioner returns a PTransform that
      partitions input RecordBatches by column (feature), in addition to the
      normal row partitioning (by batch). The total number of effective
      partitions is column_partitions * row_partitions, where row_partitions
      is passed to self.partitioner.
  
  Raises:
    ValueError: If label_feature does not exist in the schema.
code: def __init__(
      self,
      label_feature: types.FeaturePath,
      schema: Optional[schema_pb2.Schema] = None,
      max_encoding_length: int = 512,
      seed: int = 12345,
      multivalent_features: Optional[Set[types.FeaturePath]] = None,
      categorical_features: Optional[Set[types.FeaturePath]] = None,
      features_to_ignore: Optional[Set[types.FeaturePath]] = None,
      normalize_by_max: bool = False,
      allow_invalid_partitions: bool = False,
      custom_stats_key: str = _ADJUSTED_MUTUAL_INFORMATION_KEY,
      column_partitions: int = 1,
  ):
      """Initializes MutualInformation.
  
      Args:
        label_feature: The key used to identify labels in the ExampleBatch.
        schema: An optional schema describing the the dataset. Either a schema or
          a list of categorical and multivalent features must be provided.
        max_encoding_length: An int value to specify the maximum length of
          encoding to represent a feature value.
        seed: An int value to seed the RNG used in MI computation.
        multivalent_features: An optional set of features that are multivalent.
        categorical_features: An optional set of the features that are
          categorical.
        features_to_ignore: An optional set of features that should be ignored by
          the mutual information calculation.
        normalize_by_max: If True, AMI values are normalized to a range 0 to 1 by
          dividing by the maximum possible information AMI(Y, Y).
        allow_invalid_partitions: If True, generator tolerates input partitions
          that are invalid (e.g. size of partion is < the k for the KNN), where
          invalid partitions return no stats. The min_partitions_stat_presence arg
          to PartitionedStatisticsAnalyzer controls how many partitions may be
          invalid while still reporting the metric.
        custom_stats_key: A string that determines the key used in the custom
          statistic. This defaults to `_ADJUSTED_MUTUAL_INFORMATION_KEY`.
        column_partitions: If > 1, self.partitioner returns a PTransform that
          partitions input RecordBatches by column (feature), in addition to the
          normal row partitioning (by batch). The total number of effective
          partitions is column_partitions * row_partitions, where row_partitions
          is passed to self.partitioner.
  
      Raises:
        ValueError: If label_feature does not exist in the schema.
      """
      self._label_feature = label_feature
      self._schema = schema
      self._normalize_by_max = normalize_by_max
      if multivalent_features is not None:
          self._multivalent_features = multivalent_features
      elif self._schema is not None:
          self._multivalent_features = (
              schema_util.get_multivalent_features(self._schema)
          )
      else:
          raise ValueError(
              "Either multivalent feature set or schema must be provided"
          )
      if categorical_features is not None:
          self._categorical_features = categorical_features
      elif self._schema is not None:
          self._categorical_features = (
              schema_util.get_categorical_features(self._schema)
          )
      else:
          raise ValueError(
              "Either categorical feature set or schema must be provided"
          )
      if schema:
          assert schema_util.get_feature(
              self._schema, self._label_feature
          )
      self._label_feature_is_categorical = (
          self._label_feature in self._categorical_features
      )
      self._max_encoding_length = max_encoding_length
      self._seed = seed
      self._features_to_ignore = features_to_ignore
      self._allow_invalid_partitions = allow_invalid_partitions
      self._custom_stats_key = custom_stats_key
      self._column_partitions = column_partitions
  
imports: import collections
  from typing import Any, Dict, Iterable, List, Optional, Set, Tuple
  from absl import logging
  import apache_beam
  import numpy
  import pandas
  import pyarrow
  from tensorflow_data_validation import types
  from tensorflow_data_validation.statistics.generators import partitioned_stats_generator
  from tensorflow_data_validation.utils import feature_partition_util
  from tensorflow_data_validation.utils import mutual_information_util
  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 1 to create a new MutualInformation object from inherited base classes: partitioned_stats_generator with example: obj = MutualInformation(label_feature, schema, max_encoding_length, seed, multivalent_features, categorical_features, features_to_ignore, normalize_by_max, allow_invalid_partitions, custom_stats_key, column_partitions)