transform: Fail to load transform when `scale_to_z_score_per_key` is used

Hey team,

I have a couple of weighted sparse features in the form of:

{
    'feat_keys': tf.io.VarLenFeature(tf.string),
    'feat_weights': tf.io.VarLenFeature(tf.float32),
}

feat_keys and feat_weights are exactly the same dimension. In my case, the feat_weights is the value of multiple decayed counts and the feat_keys indicates the decay. Since each feat_weights entry correspond to a count distinct from the other values, I’m trying to use scale_to_z_score_per_key to normalize them. My pre-processing function looks like this:

def preprocessing_fn(feats):
    feats = dict(feats)

    for f in FEATS_TO_BUILD_VOCAB:
        tft.vocabulary(feats[f], vocab_filename=f)

    for k, v in CONTS_FEATS_TO_NORMALIZE_BY_KEY.items():
        feats[v + '__norm'] = tft.scale_to_z_score_per_key(
            x=feats[v],
            key=feats[k],
        )

    return feats

The transform is produced as expected. However, when I try to use it, I get this error:

InvalidArgumentError: Op type not registered
'scale_to_z_score_per_key_mean_and_var_per_key_assert_equal_1_Assert_AssertGuard_true_251'
in binary running on notebook-de-gpu-0. Make sure the Op and Kernel are registered in the
binary running in this process. Note that if you are loading a saved graph which used ops from
tf.contrib, accessing (e.g.) `tf.contrib.resampler` should be done before importing the graph, as
contrib ops are lazily registered when the module is first accessed. while building NodeDef
'model_1/transform_features_layer_1/transform/transform/scale_to_z_score_per_key/mean_and_var_per_key/assert_equal_1/Assert/AssertGuard/then/_0'
[Op:__inference_distributed_function_5652]

It looks that some parts of the transform are not available in the binary somehow.

I thought I was making some basic mistake in how I’m using TF Transform, but if I comment out the tft.scale_to_z_score_per_key piece, all works as expected, so this looks like a bug. Any ideas of what might be happening and how I can debug this?

I’m using TF 2.1 (eager mode) and TFT 0.21.2.