How to export tensorflow model

[fsx950223]

Description

Could mode.pkl transfer to tensorflow model?
...

Environment information

OS: <your answer here>

$ pip freeze | grep tensor
# your output here

$ pip freeze | grep jax
# your output here

$ python -V
# your output here

For bugs: reproduction and error logs

# Steps to reproduce:
...

# Error logs:
...

[lukaszkaiser]

We're working on code to make this work, keeping this issue to track.

[lukaszkaiser]

This is done, you can make a Keras layer from any Trax layer like here:
https://github.com/google/trax/blob/master/trax/trax2keras.py

It will transfer weights too. So when you have a Trax model, do:
model.init_from_file(...model.pkl)
keras_model = TraxKerasLayer(model)

Now you can use it as any Keras layer:

keras_model.build(...)
keras_model.save(...)