Hi all, I’m trying to incorporate a GBT model using TFDF into an existing TFX project. Thus far, we are able to successfully train a GBT using TFDF inside a trainer component, but we are having issues creating an Evaluator stage for this model.
I believe our issue lies with the serving functions we’re creating for this model. We are following the same technique used previously for normal TF/Keras NNs. Here is the error that occurs during the evaluation stage:
RuntimeError: tensorflow.python.framework.errors_impl.InvalidArgumentError: slice index 0 of dimension 1 out of bounds.
            [[{{node StatefulPartitionedCall/gradient_boosted_trees_model/StatefulPartitionedCall/strided_slice_3}}]] [Op:__inference_signature_wrapper_1194552]

Here is the code for creating the serving functions & signatures:
def _get_serve_tf_examples_fn(model, tf_transform_output, transform=False):
    """Returns a function that parses a serialized tf.Example and applies TFT."""

    model.tft_layer = tf_transform_output.transform_features_layer()

    @tf.function
    def serve_tf_examples_fn(serialized_tf_examples):
        """Returns the output to be used in the serving signature."""
        raw_feature_spec = tf_transform_output.raw_feature_spec().copy()
        feature_spec = { k: v for k, v in raw_feature_spec.items() if k in FEATURES }
        label_spec = raw_feature_spec.get(_LABEL_KEY) or {}

        if transform:
            parsed_features_with_label = tf.io.parse_example(
                serialized_tf_examples, { **feature_spec, **label_spec }
            )

            return model.tft_layer(parsed_features_with_label)

        parsed_features = tf.io.parse_example(serialized_tf_examples, feature_spec)
        transformed_features = model.tft_layer(parsed_features)

        return model(transformed_features)

    return serve_tf_examples_fn

...

# Code for creating signatures & saving model
 _serving_default = _get_serve_tf_examples_fn(model, tf_transform_output) \
          .get_concrete_function(
              tf.TensorSpec(shape=[None], dtype=tf.string, name='examples')
          )

_transform = _get_serve_tf_examples_fn(model, tf_transform_output, transform=True) \
          .get_concrete_function(
              tf.TensorSpec(shape=[None], dtype=tf.string, name='examples')
          )

signatures = { 'serving_default': _serving_default, 'transform': _transform }

model.save(fn_args.serving_model_dir, save_format='tf', signatures=signatures)

Has anyone successfully used TFDF with TFX?