wrap_trax_keras.py

# Create a full Keras model using the layer you loaded from trax.
inputs = tf.keras.Input(shape=(None,), dtype='int32')
# Use default Keras syntax to link an input to the layer
hidden = keras_layer(inputs)
# Get the outputs from the trax-loaded layer
outputs = hidden
# Finally, wrap everything with a Keras Model
keras_model = tf.keras.Model(inputs=inputs, outputs=outputs)