Adding switch_max routing that, when overflow occurs, will zero out the lowest probability tokens according to the router instead of the last tokens in the group.
Adding switch_max routing that, when overflow occurs, will zero out the lowest probability tokens according to the router instead of the last tokens in the group.

PiperOrigin-RevId: 365980105

tensorflow/mesh,tensorflow/mesh