Moving 'activation_vocab' sharding rule with single axis to the top

Without this change, the automatic sharding resolution fails to pick up the correct sharding for activation_vocab when tensor_sequence is used.

Author: abhinavgoel95

src/MaxText/configs/base.yml

@@ -378,9 +378,9 @@ logical_axis_rules: [
                       ['activation_kv_batch', ['data', 'fsdp', 'fsdp_transpose', 'expert']],
                       ['activation_kv_batch_no_exp', ['data', 'fsdp', 'fsdp_transpose']],
                       ['activation_kv_head_dim', ['tensor', 'tensor_transpose', 'tensor_sequence']],
+                      ['activation_vocab', 'tensor_sequence'],
                       ['activation_vocab', ['tensor', 'tensor_transpose', 'sequence', 'tensor_sequence']],
                       ['activation_vocab', ['tensor', 'tensor_transpose']],
-                      ['activation_vocab', 'tensor_sequence'],
                       ['activation_vocab', ['sequence','context']],
                       ['activation_stage', 'stage'],
                       ['activation_exp', ['expert']],