[NNX] NNX migration prep (10/N): vocab tiling custom_vjp with output-head carve-out#3849

copybara-service[bot] merged 1 commit into

mainAI-Hypercomputer/maxtext:mainfrom

feat/nnx-vocab-tiling-custom-vjpAI-Hypercomputer/maxtext:feat/nnx-vocab-tiling-custom-vjpCopy head branch name to clipboard