class_docstr: Tests for NonStreamingCustomStatsGenerator.
class_docstr_tok: ["Tests","for","NonStreamingCustomStatsGenerator","."]
code: def setUp(self):
      # Integration tests involving Beam and AMI are challenging to write
      # because Beam PCollections are unordered while the results of adjusted MI
      # depend on the order of the data for small datasets. This test case tests
      # MI with one label which will give a value of 0 regardless of
      # the ordering of elements in the PCollection. The purpose of this test is
      # to ensure that the Mutual Information pipeline is able to handle a
      # variety of input types. Unit tests ensuring correctness of the MI value
      # itself are included in mutual_information_test.
  
      # fa is categorical, fb is numeric, fc is multivalent categorical, fd is
      # multivalent numeric
  
      self.record_batches = [
          pa.RecordBatch.from_arrays(
              [
                  pa.array([["1"]]),
                  pa.array([[1.1]]),
                  pa.array([["1", "1", "1"]]),
                  pa.array([[1.0, 1.2, 0.8]]),
                  pa.array([["label"]]),
              ],
              ["fa", "fb", "fc", "fd", "label_key"],
          ),
          pa.RecordBatch.from_arrays(
              [
                  pa.array([["0"]]),
                  pa.array([[0.3]]),
                  pa.array([["0", "1"]]),
                  pa.array([[0.1, 0.0]]),
                  pa.array([["label"]]),
              ],
              ["fa", "fb", "fc", "fd", "label_key"],
          ),
          pa.RecordBatch.from_arrays(
              [
                  pa.array([["1"]]),
                  pa.array([[np.NaN]], type=pa.list_(pa.float64())),
                  pa.array([["0", "0"]]),
                  pa.array([[0.0, 0.2]]),
                  pa.array([["label"]]),
              ],
              ["fa", "fb", "fc", "fd", "label_key"],
          ),
          pa.RecordBatch.from_arrays(
              [
                  pa.array([None]),
                  pa.array([None]),
                  pa.array([["1", "0", "0", "1"]]),
                  pa.array([None]),
                  pa.array([["label"]]),
              ],
              ["fa", "fb", "fc", "fd", "label_key"],
          ),
          pa.RecordBatch.from_arrays(
              [
                  pa.array([["1"]]),
                  pa.array([[1.0]]),
                  pa.array([["1", "1"]]),
                  pa.array([[1.0]]),
                  pa.array([["label"]]),
              ],
              ["fa", "fb", "fc", "fd", "label_key"],
          ),
          pa.RecordBatch.from_arrays(
              [
                  pa.array([["0"]]),
                  pa.array([[0.3]]),
                  pa.array([["0"]]),
                  pa.array([[0.0, 0.2]]),
                  pa.array([["label"]]),
              ],
              ["fa", "fb", "fc", "fd", "label_key"],
          ),
          pa.RecordBatch.from_arrays(
              [
                  pa.array([None]),
                  pa.array([None]),
                  pa.array([["1", "0", "0", "1"]]),
                  pa.array([None]),
                  pa.array([["label"]]),
              ],
              ["fa", "fb", "fc", "fd", "label_key"],
          ),
          pa.RecordBatch.from_arrays(
              [
                  pa.array([["1"]]),
                  pa.array([[1.0]]),
                  pa.array([["1", "1"]]),
                  pa.array([[1.0]]),
                  pa.array([["label"]]),
              ],
              ["fa", "fb", "fc", "fd", "label_key"],
          ),
          pa.RecordBatch.from_arrays(
              [
                  pa.array([["0"]]),
                  pa.array([[0.3]]),
                  pa.array([["0"]]),
                  pa.array([[0.0, 0.2]]),
                  pa.array([["label"]]),
              ],
              ["fa", "fb", "fc", "fd", "label_key"],
          ),
          pa.RecordBatch.from_arrays(
              [
                  pa.array([None]),
                  pa.array([None]),
                  pa.array([["1", "0", "0", "1"]]),
                  pa.array([None]),
                  pa.array([["label"]]),
              ],
              ["fa", "fb", "fc", "fd", "label_key"],
          ),
          pa.RecordBatch.from_arrays(
              [
                  pa.array([["1"]]),
                  pa.array([[1.0]]),
                  pa.array([["1", "1"]]),
                  pa.array([[1.0]]),
                  pa.array([["label"]]),
              ],
              ["fa", "fb", "fc", "fd", "label_key"],
          ),
          pa.RecordBatch.from_arrays(
              [
                  pa.array([["0"]]),
                  pa.array([[0.3]]),
                  pa.array([["0"]]),
                  pa.array([[0.0, 0.2]]),
                  pa.array([["label"]]),
              ],
              ["fa", "fb", "fc", "fd", "label_key"],
          ),
      ]
  
      self.schema = text_format.Parse(
          """
          feature {
            name: "fa"
            type: BYTES
            value_count {
              min: 0
              max: 1
            }
          }
          feature {
            name: "fb"
            type: FLOAT
            value_count {
              min: 0
              max: 1
            }
          }
          feature {
            name: "fc"
            type: BYTES
            value_count {
              min: 0
              max: 4
            }
          }
          feature {
            name: "fd"
            type: FLOAT
            value_count {
              min: 0
              max: 3
            }
          }
          feature {
            name: "label_key"
            type: BYTES
            value_count {
              min: 1
              max: 1
            }
          }""",
          schema_pb2.Schema(),
      )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from absl.testing import absltest
  from absl.testing import parameterized
  import apache_beam
  from apache_beam.testing import util
  import numpy
  import pyarrow
  from tensorflow_data_validation import types
  from tensorflow_data_validation.statistics.generators import mutual_information
  from tensorflow_data_validation.statistics.generators import partitioned_stats_generator
  from tensorflow_data_validation.utils import test_util
  from tfx_bsl.arrow import table_util
  from google.protobuf import text_format
  from tensorflow_metadata.proto.v0 import schema_pb2
  from tensorflow_metadata.proto.v0 import statistics_pb2
label_desc: Use image node_id 1 for calling the NonStreamingCustomStatsGeneratorTest obj's underlying member method code with example usage: obj.setUp() without return types