Minor fix after jax changes.#203

copybara-service[bot] wants to merge 1 commit intomaingoogle-research/vmoe:mainfrom

test_773571781google-research/vmoe:test_773571781Copy head branch name to clipboard