main
adds jais2 model support (#42684)
🚨 fix + tests dense & MoE TP all reduce (decoder only) (#43722)