Add SetCanonicalShardingParams to TransformerLayer
Add SetCanonicalShardingParams to TransformerLayer

PiperOrigin-RevId: 350229390

tensorflow/lingvo,tensorflow/lingvo,tensorflow/lingvo,tensorflow/lingvo