Skip to content
CI - Cloud TPU (presubmit)

Fix a bug where full and use_mesh outside jit did not work because the shard passed to make_array_from_callback was sharded on all devices instead of just 1 device. #5767

Sign in to view logs

Fix a bug where full and use_mesh outside jit did not work because the shard passed to make_array_from_callback was sharded on all devices instead of just 1 device.

Fix a bug where full and use_mesh outside jit did not work because the shard passed to make_array_from_callback was sharded on all devices instead of just 1 device. #5767

Triggered via pull request March 11, 2025 22:25
@copybara-servicecopybara-service[bot]
synchronize #27076
test_735894275
Status Success
Total duration 10m 1s
Artifacts –

cloud-tpu-ci-presubmit.yml

on: pull_request
Matrix: build-jax-artifacts
TPU test (jaxlib=head, v5e-8)  /  Pytest TPU (v5e-8, Python 3.10, libtpu=nightly)
6m 59s
TPU test (jaxlib=head, v5e-8) / Pytest TPU (v5e-8, Python 3.10, libtpu=nightly)
Fit to window
Zoom out
Zoom in