Questions regarding trax and it's roadmap

[jimexist]

Description

This is a quick note for some of my questions while using trax for the first time so the sections below are not applicable.

I came from the deeplearning.ai's cource and it's of high quality, packed with SOTA models, and surprisingly using trax for the coding projects. So far I have enjoyed the library, esp. it's adaptation to numpy and jax, so ideally it would have a widened audience beyond just NLP, deep learning, to include e.g. general community of scientific computing, or at least this is my understanding of it.

However I still have a few questions / suggestions in my mind regarding the future of this project and it matters to me before I can commit more time into using or contributing to it, so I'd appreciate if someone on the Google brain team could give some hint:

Production usage

New SOTA models are being written in this trax library (performer, i recon?) but to really put it into usage we'd really like more documentation, examples, guidance on how to deploy them - huggingface did a great job in providing a dedicated pipeline abstraction and a detailed index for all pretrained models and how people can contribute to that. spaCy 3.0 also would be coming out with a pipeline API for transformers along with pretrained models, which seems nice.

It might be too much to ask but specifically for huggingface models, after PyTorch being the first citizen there, they have since added support for TensorFlow 2.0 APIs, so I guess maybe they (or trax) can provide a Trax API docs? Just some thoughts here, but in general developers would not want to do things again here and there - if there's a unified model to do these we'd be keen.

non-NLP models

For NLP (attention based models really) Trax really shines in its simplicity. But in other areas e.g. SOTA CV models there's still a shortage. If attention based models are not the limiting factors maybe we'd need more examples, documentations, etc. in here. Take albumentations for example, it's where the community of PyTorch starts to leapfrog TensorFlow (but it's adding TensorFlow 2.0 API now) - so I guess that can be one area for improvement

Error messages, UX, visualizations

So far while coding #1199 I found several places where error messages should've been improved: e.g. missing a batch dimension gives unhelpful error messages, there's no place that mentions train_batches_stream must be added with a 3rd weighted loss value (and if you miss that in data pipeline you'd get cryptic error messages). Fixing this really requires a fresh pair of eyes I guess :-)

From a visualization point of view there's little docs on how to use with TensorBoard, Wandb, etc. but they'll really be cool things to add and of little risk.

Model zoo and inter-operatability

I guess the one biggest missing piece is the lack of a model zoo or at least a place for people to find pre-trained models (even better for them to share and contribute). There's also little mentioning of the gin config system.

I guess all the above are sometimes too much to ask for a newly started projects so i would emphasize that above all I am enjoying this library and these are just my hopes and wishes for it being more popular :-)

---

Environment information

OS: <your answer here>

$ pip freeze | grep trax
# your output 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]

Thanks for your questions! In general, the gitter chat may be a better place for questions like this since they often benefit from a back-and-forth discussion and these issues are reserved for code. But here a few quick answers (happy to continue on chat):

• Production usage: we suggest to deploy using the TF backend and SavedModel, which works with JavaScript and Servo and other TF packages. Here's a colab showing how to get a SavedModel from Trax: https://trax-ml.readthedocs.io/en/latest/notebooks/tf_numpy_and_keras.html

• Non-NLP models: yes' it'd be great to have more of them! We're just waiting for some vision enthusiasts to send PRs 👍

• We do need better error messages - yes! Feel free to open PRs, we're also considering a bigger design change (to remove delayed init like pyTorch) that could improve a lot of them in one go.

• TensorBoard is fully supported, we probably just need to document it better (please send PRs!)

• Model ZOO and colab examples are a big missing part. We were not sure how to structure it, but just opened and examples/ directory today, with the idea that we'll start gathering colabs and putting them there, with pre-trained models on GCP

Hope that helps and looking forward to talking more!

[INF800]

Hey @lukaszkaiser, I am interested to send some PRs for vision models. Can I know where can I follow up ?