How to define a custom gate function in moe_layer

[lihuanihao]

Hello, thank you for your very useful codebase. I would like to know how to define a custom gate function in the moe_layer, like self._layer = moe_layer(gate_type={'type': 'custom','module': CustomGate, ... }?

[ghostplant]

You can upgrade tutel to include latest commit, then follow this example:

https://github.com/microsoft/Tutel/blob/main/tutel/examples/helloworld_custom_gate_expert.py#L53-L118

Example script: python3 -m tutel.examples.helloworld_custom_gate_expert