Skip to content
ci

Adding JAX config change jax_pmap_shmap_merge=False to the top of the optimizer.py file. The new pmap implementation currently seems to slow down projects using kfac_jax. #1836

Sign in to view logs

Adding JAX config change jax_pmap_shmap_merge=False to the top of the optimizer.py file. The new pmap implementation currently seems to slow down projects using kfac_jax.

Adding JAX config change jax_pmap_shmap_merge=False to the top of the optimizer.py file. The new pmap implementation currently seems to slow down projects using kfac_jax. #1836

Triggered via pull request December 5, 2025 21:02
@copybara-servicecopybara-service[bot]
synchronize #367
test_840716303
Status Success
Total duration 5m 18s
Artifacts –

ci.yml

on: pull_request
Matrix: build-and-test
Fit to window
Zoom out
Zoom in