Maintained Documentation

[felipeboffnunes]

Description

Trax is a library for deep learning that focuses on sequence models and reinforcement learning. It combines performance with code clarity and maintained documentation and tests. ...

Sorry to bother, I'll be brief. I don't think the "maintained documentation" part of the statement is true (yet?). I like the work and I respect every project that goes deep down on neural network implementation, but I feel there is a critical lack of documentation for this project.

I was giving a look at Flax read the docs and, although the projects have different motives, I believe there should be something alike for Trax.

Again, sorry to bother. Wish all the luck and success for the project.

[jalammar]

Maybe this is something outside contributors can help with. To me, this also brings to mind the discussions recently on Twitter. While some of these discussions are on API design and library scope, It has been making me think of ways the community can support with the docs for Trax.

One of my favorite ways to think of this is how "What nobody tells you about documentation" breaks it down into:

• Tutorials
• How-To Guides
• Explanation
• References

I suspect that at this point in time, references are the key required artifact. @lukaszkaiser indicated the team is currently working on it. If the team is inclined, it could be worth their time to tag a few issues "good first issue" or "contributions welcome" for newcomers to start helping out.

Current Trax Docs Inventory Besides README.md and the code API, the other documentation artifacts are currently: Trax Quick Intro Colab Short, sweet, and to the point. I really like how to-the-point this colab was.

Reformer Colabs: These are How-to guides for three tasks currently:

• Image Generation
• Machine Translation
• Text Generation

*Edit:** layers/intro notebook.

What Docs Should be Produced Next? Curious to hear from the contributors and community what they feel is needed next.

Personally, it's important to know which aim of the library takes precedence: A - The examples/models -- e.g. people coming in to use the Reformer, or continue where Transformer2Transformer left off, or for a transformer implementation that works on CPUs/GPUs/TPUs. Or; B - to learn a new deep learning library from the ground up.

[cgarciae]

Same question. I was kind of expecting an API Reference given the statement.

[agrawalparth08]

@jalammar Just the other day I was trying to visualize the model layers in a better way (like keras model summary ) and I found some moderate code help in intro.ipynb in layers which is incomplete in the models part. Tried starting with this small cleaner function, may need more recursive way to pretty print/create a graphical tree like layout of the model which would help in debugging models as well.

[lukaszkaiser]

It took us a while to get familiar with readthedocs, but we do have the API docs there at last! https://trax-ml.readthedocs.io/en/latest/

We're now working on getting the colabs (quick intro and layer intro) there too.

It was hard to ask the community for help as we're making it work both on readthedocs and doing some internal work (e.g., so the colabs are unit-tested at each commit). But it is coming to an end, and from here we should be able to easily accept PRs and we'll definitely ask for help!

Some first examples where help would be welcome:

• API docs for convolution and pooling layers
• A colab with image models (say WIdeResNet on CIFAR or even MLP on MNIST or both) showing the results
• A Transformer colab in trax/models (explaining the model first, later we can add translation training?)

A meta-question: where would be a good place to put such open TODOs? Where would you expect them?

[jalammar]

That's wonderful, @lukaszkaiser! Congrats for shipping!

Since reading your comment, I have already started working on a Transformer colab for trax/model/transformer. I'll be sharing a draft in the next couple of days. I think I have a gentle narrative nailed down for it. If there's any additional guidance on what you and the team want to see in the colab, please let me know. I'm considering for the flow to start with a trained TransformerLM, build some confidence in the reader by establishing the model, the blocks, and the Trainer, then move to the Transformer model. I've always found it beneficial to readers if they go over a trained model first, then start learning about the training process once they're comfortable with the model.

For the other two (API docs for Conv/pooling and image models), maybe create an issue for each of the two (for visibility) and tag with "documentation", "good first issue" and "contributions welcome" (this repo has a "help wanted" tag, Tensorflow and JAX have "contributions welcome", either can work)?

[lukaszkaiser]

I created #897 for starters, please take it if you wish :). We also updated the main page README and the trax intro colab -- it has a direct Transformer decoding part with a pre-trained En-De model. So there is one trained model, I'm happy to add a LM too if you think having a decoder-only model would be helpful.

Will create more issues (and tags/labels) soon, thanks for the suggestions!!