class_docstr: Custom partitioner DoFn for MutualInformation.
docstr: Performs row-wise random key assignment and column-wise slicing.
  
  Each input RecordBatch is mapped to up to self._column_partitions output
  RecordBatch, each of which contains a subset of columns. Only the label
  column is duplicated across RecordBatches, so this is nearly a partitioning
  of columns. If self._column_partitions == 1, the output RecordBatch is
  unmodified.
  
  The total partition key space is _row_partitions * _column_partitions.
  
  Args:
    element: An input sliced record batch.
  
  Yields:
    A sequence of partitioned RecordBatches.
code: def process(
      self, element: types.SlicedRecordBatch
  ) -> Iterable[Tuple[Tuple[types.SliceKey, int], pa.RecordBatch]]:
      """Performs row-wise random key assignment and column-wise slicing.
  
      Each input RecordBatch is mapped to up to self._column_partitions output
      RecordBatch, each of which contains a subset of columns. Only the label
      column is duplicated across RecordBatches, so this is nearly a partitioning
      of columns. If self._column_partitions == 1, the output RecordBatch is
      unmodified.
  
      The total partition key space is _row_partitions * _column_partitions.
  
      Args:
        element: An input sliced record batch.
  
      Yields:
        A sequence of partitioned RecordBatches.
  
      """
  
      row_partition = self._rng.integers(
          0, self._row_partitions, dtype=int
      )
      if self._partitioner is None:
          slice_key, record_batch = element
          yield (slice_key, row_partition), record_batch
      else:
          for (
              slice_key,
              column_partition,
          ), record_batch in feature_partition_util.generate_feature_partitions(
              element, self._partitioner, self._label_column
          ):
              partition = (
                  row_partition * self._column_partitions
                  + column_partition
              )
              yield (slice_key, partition), record_batch
  
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 3 for calling the _PartitionFn obj's underlying member method code with example usage: obj.process(element) without return types