[JAX] Migrate EP collective-stream annotation to compute_on2 - #3304

[JAX] Migrate EP collective-stream annotation to compute_on2#3304

phu0ngng wants to merge 1 commit into

NVIDIA:mainNVIDIA/TransformerEngine:mainfrom

phu0ngng:jax_compute_on_versionphu0ngng/TransformerEngine:jax_compute_on_versionCopy head branch name to clipboard