ROCm/jax

Pull Requests Commits

Add support for rocm release rpaths

alekstheod committed 1 day ago

a35f5781

Set rpath based on the type (link_only)

alekstheod committed 1 day ago

511e7dc2

[mosaic] Added a new tpu.fetch_and_add_sync op

superbobry committed 14 days ago

2f1645f1

Add some more information to deprecation messages in pxla modules

yashk2810 committed 14 days ago

0775e2b3

Reverts 4a5ad80ec771ec08bea2b28169d2f298aecbb46e

hawkinsp committed 14 days ago

5c4c7f05

Merge pull request #34682 from olupton:workaround-tcgen05-llvm-errors

Google-ML-Automation committed 14 days ago

8288263b

[Mosaic GPU] Clean up all the linter warnings in `dialect_lowering.py`.

bchetioui committed 14 days ago

40c3359e

[Pallas/Mosaic GPU] Remove unnecessary jaxlib version guard-related code.

bchetioui committed 14 days ago

a5854e44

Update XLA dependency to use revision http://github.com/openxla/xla/commit/80529bdfc1c3bc206d0809a474031e9c9008c9f0

Google-ML-Automation committed 15 days ago

1066aa7b

[mosaic] Added a macro which defines AbslStringify for all MLIR enums in the enclosing namespace

superbobry committed 15 days ago

150a6d9d

[mosaic] Avoid using the deprecated `output_memory_colors` in `CustomCallBackendConfig`

superbobry committed 15 days ago

79fd127f

Regenerate python stubs for nanobind 2.11.0.

hawkinsp committed 15 days ago

6fb3f3da

Update XLA dependency to use revision http://github.com/openxla/xla/commit/2a7188001603266f57a661a10beee3418412c312

Google-ML-Automation committed 16 days ago

3a93274b

Merge pull request #34462 from alekstheod:add_support_of_building_jax_under_umbrella_workspace

Google-ML-Automation committed 16 days ago

4a5ad80e

Adds `jax.lax.scaled_dot` for scaled dot products.

loislo committed 16 days ago

50b394dc

Update XLA dependency to use revision http://github.com/openxla/xla/commit/47fe49edce6c27225a0cb6675bf08d43550316a6

Google-ML-Automation committed 17 days ago

91871a92

[Mosaic GPU][NFC] Clean up the lowering of `slice_smem`.

bchetioui committed 17 days ago

01d76c43

Update jax.host_id() and jax.host_count() to jax.process_index() and jax.process_count(), respectively.

danielsuo committed 17 days ago

46c02566

Merge pull request #34742 from mattjj:remove-zero-dot-from-primal-value

Google-ML-Automation committed 17 days ago

9c5d90dd

[Mosaic GPU] Fix handling of sliced/transposed types in `transformed_smem_ref_type`.

bchetioui committed 17 days ago

139f212f

Fix ValAccum to initialize Zero using .to_cotangent_aval()

yashk2810 committed 17 days ago

12462e05

[Mosaic GPU][NFC] Linter + type checker fixes in layout inference code.

bchetioui committed 17 days ago

457340ea

[Mosaic GPU][NFC] Delete dead code and clean up linter/pytype warnings in `layout_inference.py`.

bchetioui committed 17 days ago

fac19bbf

mattjj committed 17 days ago

39cc010b

[Mosaic GPU] Fix layout inference for `tcgen05_mma` to be compatible with user-specified tile sizes.

bchetioui committed 17 days ago

d751a3f7

[Mosaic TPU] Update some lowering rules to avoid using `tpu.repeat` since the op is going to be deprecated.

yueshengys committed 17 days ago

b1359ede

[ci] add pyrefly pre-commit hook

Jake VanderPlas committed 17 days ago

382ed1b7

Add the infer_params cache back that was removed during the lu.cache -> weakref_lru_cache move.

yashk2810 committed 17 days ago

b582e359

[Mosaic GPU] Fix bug in layout inference of WGMMA.

bchetioui committed 17 days ago

e786d22e

Clean up `LoadedExecutable::GetReadyFuture()`

junwhanahn committed 17 days ago

f0eddf93

Older