Fix misuses of jax._src.lib.mlir.#538

copybara-service[bot] merged 1 commit intomainjax-ml/jax-tpu-embedding:mainfrom

test_826263298jax-ml/jax-tpu-embedding:test_826263298Copy head branch name to clipboard