//  [python]
def testScaleToZScorePerKey(self, key_vocabulary_filename):
    # TODO(b/131852830) Add elementwise tests.
    def preprocessing_fn(inputs):

      def scale_to_z_score_per_key(tensor, key, var_name=''):
        if key_vocabulary_filename is None:
          filename = None
        else:
          filename = key_vocabulary_filename + var_name
        z_score = tft.scale_to_z_score_per_key(
            t.cast(tensor, tf.float32), key=key, elementwise=False,
            key_vocabulary_filename=filename)
        self.assertEqual(z_score.dtype, tf.float32)
        return z_score

      return {
          'x_scaled': scale_to_z_score_per_key(inputs['x'], inputs['key'], 'x'),
          'y_scaled': scale_to_z_score_per_key(inputs['y'], inputs['key'], 'y'),
          's_scaled': scale_to_z_score_per_key(inputs['s'], inputs['key'], 's'),
      }

    np_dtype = np.float32
    input_data = [
        {
            'x': np.array([-4], dtype=np_dtype),
            'y': np.array([0], dtype=np_dtype),
            's': 3,
            'key': 'a',
        },
        {
            'x': np.array([10], dtype=np_dtype),
            'y': np.array([0], dtype=np_dtype),
            's': -3,
            'key': 'a',
        },
        {
            'x': np.array([1], dtype=np_dtype),
            'y': np.array([0], dtype=np_dtype),
            's': 3,
            'key': 'b',
        },
        {
            'x': np.array([2], dtype=np_dtype),
            'y': np.array([0], dtype=np_dtype),
            's': 3,
            'key': 'a',
        },
        {
            'x': np.array([4], dtype=np_dtype),
            'y': np.array([0], dtype=np_dtype),
            's': -3,
            'key': 'a',
        },
        {
            'x': np.array([-1], dtype=np_dtype),
            'y': np.array([0], dtype=np_dtype),
            's': -3,
            'key': 'b',
        },
        {
            'x': np.array([np.nan], dtype=np_dtype),
            'y': np.array([np.nan], dtype=np_dtype),
            's': np.nan,
            'key': 'b',
        },
    ]
    # 'a':
    # Mean(x) = 3, Mean(y) = 0
    # Var(x) = (-7^2 + -1^2 + 7^2 + 1^2) / 4 = 25, Var(y) = 0
    # StdDev(x) = 5, StdDev(y) = 0
    # 'b':
    # Mean(x) = 0, Mean(y) = 0
    # Var(x) = 1, Var(y) = 0
    # StdDev(x) = 1, StdDev(y) = 0
    expected_data = [
        {
            'x_scaled': [-1.4],  # [(-4 - 3) / 5, (2 - 3) / 5]
            'y_scaled': [0.],
            's_scaled': 1.,
        },
        {
            'x_scaled': [1.4],  # [(10 - 3) / 5, (4 - 3) / 5]
            'y_scaled': [0.],
            's_scaled': -1.,
        },
        {
            'x_scaled': [1.],  # [(1 - 0) / 1, (-1 - 0) / 1]
            'y_scaled': [0.],
            's_scaled': 1.,
        },
        {
            'x_scaled': [-.2],  # [(-4 - 3) / 5, (2 - 3) / 5]
            'y_scaled': [0.],
            's_scaled': 1.,
        },
        {
            'x_scaled': [.2],  # [(10 - 3) / 5, (4 - 3) / 5]
            'y_scaled': [0.],
            's_scaled': -1.,
        },
        {
            'x_scaled': [-1.],  # [(1 - 0) / 1, (-1 - 0) / 1]
            'y_scaled': [0.],
            's_scaled': -1.,
        },
        {
            'x_scaled': [np.nan],
            'y_scaled': [np.nan],
            's_scaled': np.nan,
        },
    ]

    input_metadata = tft.DatasetMetadata.from_feature_spec({
        'x':
            tf.io.FixedLenFeature([1],
                                  tft_unit.canonical_numeric_dtype(tf.float32)),
        'y':
            tf.io.FixedLenFeature([1],
                                  tft_unit.canonical_numeric_dtype(tf.float32)),
        's':
            tf.io.FixedLenFeature([],
                                  tft_unit.canonical_numeric_dtype(tf.float32)),
        'key':
            tf.io.FixedLenFeature([], tf.string),
    })
    expected_metadata = tft.DatasetMetadata.from_feature_spec({
        'x_scaled': tf.io.FixedLenFeature([1], tf.float32),
        'y_scaled': tf.io.FixedLenFeature([1], tf.float32),
        's_scaled': tf.io.FixedLenFeature([], tf.float32),
    })
    self.assertAnalyzeAndTransformResults(input_data, input_metadata,
                                          preprocessing_fn, expected_data,
                                          expected_metadata)