[JAX] Add warning if using BSHD and max_segments_per_seq > 1#2796

jberchtold-nvidia merged 7 commits intoNVIDIA:mainNVIDIA/TransformerEngine:mainfrom

jberchtold-nvidia:jberchtold/te-max-segments-per-seq-warning-bshdjberchtold-nvidia/TransformerEngine:jberchtold/te-max-segments-per-seq-warning-bshdCopy head branch name to clipboard