Handle when JAX saves the manual axes on the `CallOp` of the `shmap_body`.#92

Open

copybara-service[bot] wants to merge 1 commit intomainfrom test_672543437

+49-41

Commits on Sep 9, 2024

Handle when JAX saves the manual axes on the CallOp of the shmap_body.

bartchr808
authored and
copybara-github
committed