Trax Versions Save

TODO(afrozm): Write this up.

TODO(afrozm): Write this up.

Some bug fixes.

TODO(afrozm): Describe

Lots of documentation, bugs squashed and misc changes.

Features

• Load checkpoint for fine-tuning in https://github.com/google/trax/commit/0349c3f2c32ed83d0ba54a0f6652766ea5467cca by @henrykmichalewski !
• Added generic input pipeline for GLUE tasks in https://github.com/google/trax/commit/2f490e83a9be1b802dccf866c52c8908645e10d6 by @henrykmichalewski - thanks a lot!
• Some nascent support for bfloat16s!

Models

• Performer's Favor and CausalFavor - https://github.com/google/trax/commit/77db199392ff967e49156ccb24c916b270b47eca thanks @lukaszkaiser !
• Funnel-Transformer in https://github.com/google/trax/pull/1156 thanks a lot @mvxxx !
• BERT for Trax in https://github.com/google/trax/pull/1254 , https://github.com/google/trax/pull/1223, etc by @piotrekp1 - thanks a lot!
• Residual Exchange Network by @kkanska in https://github.com/google/trax/commit/3a8f402b722a0d9ae934ed399e6532dba8401f9e ! Thanks!

PRs Merged

• Fixing broken example links in https://github.com/google/trax/pull/1263 thanks @amtagrwl !
• Added WideResnet, Deconv etc Example Notebook in https://github.com/google/trax/pull/1259 , https://github.com/google/trax/pull/1202, https://github.com/google/trax/pull/1232 thanks a lot! @SauravMaheshkar !
• Remove implicit object from the base class in https://github.com/google/trax/pull/1228 thanks @HarshCasper !
• Fashion MNIST example in https://github.com/google/trax/pull/1199 thanks @Jimexist !
• Fix PretrainedBERT init in https://github.com/google/trax/pull/1135 by @hepaajan !
• Typo in the TransformerDecoder input parameters description in https://github.com/google/trax/pull/1100 by @kujaomega !

TODO(afrozm): Write something.

PRs: Thanks @NathanHowell for fixing a bunch of typos in #962 Thanks @DarrenZhang01 for contributing to the TF-Numpy extensions code in #954

ReversibleSerialTrainer for memory efficient, layer by layer training.

Miscellaneous other issues.

Minor Fixes:

• Rename gumbel_sample to logsoftmax_sample in a2497cbd8477a11f2b96e87a4d53ce46b845ffa8
• A fix to storing checkpoints in Loop in 88b033c804a4da9021505f2ce8fbb5c7d6500574

Framework:

• Data pipeline combinators in 17d44710bee74e8bb6e3b34baa4114b0c78160af
• Multi-host training in the Loop api 524321e0f77afa58a9fba470e5bccb19ae6bdb92
• Auto-regressive sampling in fe0fa7886ee33bfac831a969a82b351ccb02941c
• T2T Tokenizers 6306231ace01ab55c5df0c401b1f3cf4c7b6a32f
• Early work on multi-task training in 423d664b470612ebbc0e5041e8288eb1a47c30ef thanks to @koz4k

Tasks / Gin files:

• IMDB dataset in 08bdb50db41a089b6e12c61a8a29981225bf3566
• Machine Translation en-pl in 5fb8aa8c5cb86dabb2338938c745996d5d87d996