ROCm/jax

Pull Requests Commits

Trigger the continuous Bazel tests on changes to the requirements lock files.

hawkinsp committed 5 days ago

aac8a1c3

[trace_to_jaxpr] replace trace_to_jaxpr_dynamic in state_test

gnecula committed 5 days ago

e63f9460

[Pallas/Mosaic GPU] Temporarily increase `shard_count` on `mosaic_gpu_test`.

bchetioui committed 5 days ago

9b684827

Reverts b550aadebd4efc1865925abe451eedeb84d65725

bchetioui committed 5 days ago

2304bb45

[Pallas][GPU kernel interpreter] Remove explicit conversions to Jax arrays in host callbacks.

Google-ML-Automation committed 5 days ago

aed2ca2a

[trace_to_jaxpr] replace trace_to_jaxpr_dynamic in Pallas

gnecula committed 5 days ago

7015be36

Also disable multiprocess tests on TPU v4 in CI.

hawkinsp committed 5 days ago

ec81f98c

[Mosaic TPU] More general support for reshape on old TPUs.

yueshengys committed 5 days ago

8792e141

Update XLA dependency to use revision http://github.com/openxla/xla/commit/1808d03c51382336c4635f605e84cc1531b2199f

Google-ML-Automation committed 5 days ago

2528c5df

Merge pull request #38486 from Cjkkkk:flex_attn

Google-ML-Automation committed 5 days ago

b550aade

[doc] Add some hints for debugging and fixing slow tracing and compilation times

gnecula committed 5 days ago

e2f2e9e2

Merge pull request #38415 from mattjj:leave-no-trace

Google-ML-Automation committed 5 days ago

03652b8b

Make broadcast_in_dim a standard_primitive

yashk2810 committed 5 days ago

0404721b

Upgrade gRPC in XLA and TensorFlow to v1.81.0

Google-ML-Automation committed 6 days ago

613eaa50

Disable multiprocess tests on TPU v5e.

hawkinsp committed 6 days ago

4ab7f97c

Change CI builds not to set JAXCI_CLONE_MAIN_XLA.

hawkinsp committed 6 days ago

79216ce7

[Mosaic GPU] Halve shard count for `tests/mosaic/gpu_test.py`.

bchetioui committed 6 days ago

b7dc7114

Fix test skip in gpu_pallas_distributed_test.

hawkinsp committed 6 days ago

b84048dd

Merge pull request #38518 from zhenying-liu:add_sync

Google-ML-Automation committed 6 days ago

a29032ec

Reverts 30db8844d7e106c64fada35b7f57e689ce3496bc

Cristian Garcia committed 6 days ago

4e50bebc

Synchronize PGLE profiler tests before stopping traces

zhenying-liu committed 6 days ago

d7ff05e0

Don't exclude "pinned" libtpu from the nightly builds.

hawkinsp committed 6 days ago

cd306c62

Error out if the arr is sharded over the axes being rolled over in `jnp.roll`

yashk2810 committed 6 days ago

098e1332

Pull the mesh from context if the aval's mesh is empty when creating an empty array

yashk2810 committed 6 days ago

acfff26b

Fix portserver initialization failure in TPU postsubmit CI job.

hawkinsp committed 6 days ago

4bdef2a9

[Mosaic GPU] Enforce len(tiling) <= len(shape) for `TileTransform.transform_shape`.

golechwierowicz committed 6 days ago

a59cbecb

[Mosaic GPU] Optimize test runtime for `test_smem_gmem_registers_load_store`.

bchetioui committed 6 days ago

5b948f63

Reduce verbosity of profiler test cases when they fail.

hawkinsp committed 6 days ago

0249bbb5

[Mosaic GPU] Remove `num_col_tiles=3` from `test_store_tiled` parametrization.

bchetioui committed 6 days ago

f16bce9d

[Mosaic GPU] Reduce the number of column tiles in `WGMMALayoutTest.test_store_tiled_with_tiling_rank`.

bchetioui committed 6 days ago

e04dbc29

Newer Older