[JAX] Better error message when Q, K, V are sharded differently#2440

Open

jberchtold-nvidia wants to merge 4 commits intoNVIDIA:mainfrom

jberchtold-nvidia:jberchtold/jax-attn-assert-qkv-sharding

Commits on Dec 2, 2025

Commits on Dec 5, 2025

Merge branch 'main' into jberchtold/jax-attn-assert-qkv-sharding
jberchtold-nvidia
authored