I have been searching for this question for hours. Obviously, in Huggingface Transformers, if we don't do anything about the model, we can directly load it into jax/flax/optax. However, what if I want to train a TensorFlow model utilizing its TPU properties, see a graph network, and then use jax/flax/optax to do something like diffusion generation like this example? It would be extremely weird for google to develop two frameworks and not allow them to work together. There is a Google Library that seems to do this. However, it is directly from tf.hub. I am not sure whether people can train their own model to their form. Or even contribute to tf.hub.
Asked
Active
Viewed 40 times
