Tensorflow Mesh needs documentation. Will this be provided anytime soon?

[shyamalschandra]

I read the paper, Switch Transformers, as carefully as possible. However, none of these parameters were glossarized and well-defined in the code and paper. For example, you have the following uncommented lines with variables:

https://github.com/tensorflow/mesh/blob/5ce96838da567061515b71ca5b767e21bdca5768/mesh_tensorflow/transformer/moe.py#L41-L87

https://github.com/tensorflow/mesh/blob/5ce96838da567061515b71ca5b767e21bdca5768/mesh_tensorflow/transformer/moe.py#L133-L159

However, I see "some" documentation that is not inline with the code in the following section:

https://github.com/tensorflow/mesh/blob/5ce96838da567061515b71ca5b767e21bdca5768/mesh_tensorflow/transformer/moe.py#L504-L609

How can I understand the paper if the explanation for Switch Transformers in MoE is unclear and too abstract for most people unless they have access to the authors of the paper?

[shyamalschandra]

Anyone out there with some credence? Any help could help?