help with categorical encoding

[g00phy]

Trying to port a kaggle housing regressor from Tensorflow to Jax but have hit a snag.

The dataset has a few categorical features(like "MSZoning", "Street","Alley") which need to hot encoded. Tensorflow has a few preprocessors like TextVectorization, StringLookup, or IntegerLookup to handle categorical features of various types.

I could build a text category encoder as shown below which would adapt on a particular categorical feature.

And finally plug the encoder in the model as shown below.

Is there any provision of StringLookup/textvectorizer in jax/flax which could adapt on multiple categories?