Skip to content

[JAX] Allow input tensors to be sharded with TP in EP dispatch/combine ops - #3283

Draft
phu0ngng wants to merge 1 commit into
NVIDIA:mainfrom
phu0ngng:jax_etp
Draft

[JAX] Allow input tensors to be sharded with TP in EP dispatch/combine ops#3283
phu0ngng wants to merge 1 commit into
NVIDIA:mainfrom
phu0ngng:jax_etp

[JAX] Support tp-sharded hidden dim in EP dispatch/combine partition …

a0f930c
Select commit
Loading
Failed to load commit list.
DCO / DCO succeeded Jul 30, 2026 in 0s

DCO

All commits are signed off!

View more details on DCO