code: def testBucketizationElementwise(
      self,
      test_inputs,
      expected_boundaries,
      do_shuffle,
      epsilon,
      should_apply,
      is_manual_boundaries,
      input_dtype,
  ):
      test_inputs = list(test_inputs)
  
      # Shuffle the input to add randomness to input generated with
      # simple range().
      if do_shuffle:
          random.shuffle(test_inputs)
  
      def preprocessing_fn(inputs):
          x = tf.cast(inputs["x"], input_dtype)
  
          num_buckets = len(expected_boundaries) + 1
          if should_apply:
              if is_manual_boundaries:
                  bucket_boundaries = [
                      expected_boundaries,
                      [2 * b for b in expected_boundaries],
                  ]
              else:
                  bucket_boundaries = tft.quantiles(
                      x,
                      num_buckets,
                      epsilon,
                      reduce_instance_dims=False,
                  )
                  bucket_boundaries = tf.unstack(
                      bucket_boundaries, axis=0
                  )
  
              result = []
              for i, boundaries in enumerate(bucket_boundaries):
                  boundaries = tf.cast(boundaries, tf.float32)
                  result.append(
                      tft.apply_buckets(
                          x[:, i], tf.expand_dims(boundaries, axis=0)
                      )
                  )
              result = tf.stack(result, axis=1)
  
          else:
              result = tft.bucketize(
                  x,
                  num_buckets=num_buckets,
                  epsilon=epsilon,
                  elementwise=True,
              )
          return {"q_b": result}
  
      input_data = [{"x": [x, 2 * x]} for x in test_inputs]
  
      input_metadata = tft.DatasetMetadata.from_feature_spec(
          {
              "x": tf.io.FixedLenFeature(
                  [2], tft_unit.canonical_numeric_dtype(input_dtype)
              )
          }
      )
  
      # Sort the input based on value, index is used to create expected_data.
      sorted_list = sorted(enumerate(test_inputs), key=lambda p: p[1])
  
      # Expected data has the same size as input, one bucket per input value.
      expected_data = [[None, None]] * len(test_inputs)
      bucket = 0
  
      for index, x in sorted_list:
          # Increment the bucket number when crossing the boundary
          if (
              bucket < len(expected_boundaries)
              and x >= expected_boundaries[bucket]
          ):
              bucket += 1
          expected_data[index] = {"q_b": [bucket, bucket]}
  
      expected_metadata = tft.DatasetMetadata.from_feature_spec(
          {
              "q_b": tf.io.FixedLenFeature([2], tf.int64),
          },
          None,
      )
  
      @contextlib.contextmanager
      def no_assert():
          yield None
  
      assertion = no_assert()
      if input_dtype == tf.float16:
          assertion = self.assertRaisesRegexp(
              TypeError,
              ".*DataType float16 not in list of allowed values.*",
          )
  
      with assertion:
          self.assertAnalyzeAndTransformResults(
              input_data,
              input_metadata,
              preprocessing_fn,
              expected_data,
              expected_metadata,
              desired_batch_size=1000,
          )
  
imports: import contextlib
  import random
  import numpy
  import tensorflow
  import tensorflow_transform
  from tensorflow_transform import analyzers
  from tensorflow_transform.beam import impl
  from tensorflow_transform.beam import tft_unit
  from tensorflow_metadata.proto.v0 import schema_pb2
label_desc: Use image node_id 4 for calling the BucketizeIntegrationTest obj's underlying member method code with example usage: obj.testBucketizationElementwise(test_inputs, expected_boundaries, do_shuffle, epsilon, should_apply, is_manual_boundaries, input_dtype) and returns: dict