[mpmd] Add custom pipeline API for user-defined scheduling and merging#33333

copybara-service[bot] wants to merge 1 commit intomainjax-ml/jax:mainfrom

test_832290784jax-ml/jax:test_832290784Copy head branch name to clipboard