Miscellaneous fixes.

• tl.Embedding now has the same signature as pytorch/tf
• train.lr_schedule (function object) -> train.lr_schedule_fn (function)
• Report loss back to training.Loop

Trax now has docs on - https://trax-ml.readthedocs.io/en/latest/trax.html thanks to @j2i2 !

Many usability changes, especially in trax.supervised.training.TrainTask/EvalTask/Loop, docs, comments etc.

• flat saved model/checkpoint representation
• lr schedule simplified, now they just take step number.
• configs are now in supervised/configs and rl/configs.
• RL obsolete code cleanup.

Also rapid development of the tf-numpy codebase !

Merged PRs:

• #459 by @w4-sjcho - adding names to layers, aiding debuggability thanks a lot!
• #256 and #300 by @stephenjfox and @satyarohith refining the README.md language, thanks a lot folks!
• #313 #312 #436 #396 from @pkol with lots of bugfixes, thanks a lot @pkol !
• #409 by @pkol -- a special shoutout to this PR, this fixes a long standing issue that prevented termination of the process by tracking the zombie threads -- thanks a lot for this @pkol specially !
• #386 another shoutout to @pkol for an amazing speedup in the RL code -- thanks a lot again !
• #344 a psum bugfix with tf backend from @fsx950223 - thanks a lot !
• #335 a bugfix from @friesel - thanks a lot Phillip !
• #315 better exception handling by @cool-RR - thanks a lot !

Reformer:

• BERT initialization and finetuning by Nikita!
• Many changes including ReformerLM on C4 dataset.

RL:

• New 'light' RL code in the Trax supervised style, check it out!
• AWR in the old code working with MuJoCo tasks.

And many more changes the Trax framework !

Reformer

• Reversible Transformer model for machine translation and other encoder-decoder tasks
• Add code for beam search, sampling, and greedy decoding (see trax.models.beam_search.Search)
• Memory-efficient attention classes have been re-written to use even less memory and to support faster decoding (see the new SelfAttention, LSHSelfAttention and EncDecAttention classes)

RL

• Implemented the Advantage-Weighted Regression algorithm, a simple off-policy reinforcement learning algorithm.
• Extracted out a PolicyBasedTrainer, so ppo_trainer.PPO and awr_trainer.AwrTrainer now both inherit from it.
• Refactoring of the serialization code in the RL part, thanks to @koz4k !

Framework

• A lot of code cleanup and refactoring of the core abstractions by Jonni, thanks Jonni!

TF Numpy

• More ops added by @wangpengmit !

• More bugfixes related to scan/remat.

• Minor bugfix to GRU and scan.

New Models

• Reformer Implementation - https://arxiv.org/abs/2001.04451 Thanks Nikita Kitaev, @lukaszkaiser and @levskaya !

Colabs

• Reformer Colabs https://github.com/google/trax/commit/9dbf8636a34cffbd421bedb2a1e3d7fe006346c0
• Transformer Colab https://github.com/google/trax/commit/d2c5b84b8db53888013c56980ccbafba077ee8f2

Framework Changes

• Ongoing cleanups and API simplifications.
• Optimization by @jekbradbury - thanks James!

PRs

• Consistent logging absl.logging and setup.py fixes thanks to @lkhphuc in #198
• Code cleanups by @cclauss in #196
• Code cleanup by @pkol in #134 - thanks!
• Bug fix by @pzielinski-nyc in #151 - thanks!