How to run a tensorflow_decision_forests RandomForestModel with TPU?

It seems that when running on TPU the RandomForestModel class cannot be used for fitting because op 'SimpleMLNumericalFeature' is not implemented :

model_RF = tfdf.keras.RandomForestModel(task=tfdf.keras.Task.REGRESSION)
model_RF.fit(x=X_ent,y=Y_ent)

I gopt the error :

NotFoundError: Op type not registered 'SimpleMLNumericalFeature' in binary running on n-3f90423c-w-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. [Op:__inference_train_function_833]

So I would like to fit it without TPU ... It it possible to do that ?

I need TPU enable because I would like to use this model with other models that run under TPU..

Thanks !

Solution

I don't think it's possible to run a random forest model on a TPU simply because TPUs were specifically designed to only run neural network computations like matrix multiplication.