Support embedding & Test auto-sharding on the whole BERT model & Refine auto-sharding interface

[merrymercy]

• Monkey patch the nn.Embed in flax to use one-hot + matmul instead of gather/scatter
• Test auto-sharing solver on the whole BERT model (copied from huggingface). The result (transformer + embedding) has exactly the same partition strategy and communication cost as Megatron-LM's solution. I will do some benchmark in the next PR.
• Refine the auto-sharding interface to better fit Combining Manual Pipeline Parallelism & Automatic SPMD Parallelism #46

Dependency:
need to update flax>=0.3.4