def preprocessing_fn(inputs):
      """A helper function for validating analyzer outputs."""
      # Get tensors representing the outputs of the analyzers
      analyzer_outputs = analyzer_fn(inputs)

      # Check that keys of analyzer_outputs match expected_output.
      self.assertCountEqual(analyzer_outputs.keys(), expected_outputs.keys())

      # Get batch size from any input tensor.
      an_input = next(iter(inputs.values()))
      if isinstance(an_input, tf.RaggedTensor):
        batch_size = an_input.bounding_shape(axis=0)
      else:
        batch_size = tf.shape(input=an_input)[0]

      # Add a batch dimension and broadcast the analyzer outputs.
      result = {}
      for key, output_tensor in analyzer_outputs.items():
        # Get the expected shape, and set it.
        expected_output_shape = list(expected_outputs[key].shape)
        try:
          output_tensor.set_shape(expected_output_shape)
        except ValueError as e:
          raise ValueError(
              f'Error for key {key}, shapes are incompatible. Got '
              f'{output_tensor.shape}, expected {expected_output_shape}.'
          ) from e
        # Add a batch dimension
        output_tensor = tf.expand_dims(output_tensor, 0)
        # Broadcast along the batch dimension
        result[key] = tf.tile(
            output_tensor,
            multiples=[batch_size] + [1] * len(expected_output_shape))

      return result
A helper function for validating analyzer outputs.
tensorflow_transform/beam/tft_unit.py