ROCm/jax

Pull Requests Commits

[ROCm/Mosaic] Initial support for Mosaic in ROCm

Arech8 committed 22 days ago

fba368fa

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

Google-ML-Automation committed 22 days ago

8288263b

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

bchetioui committed 22 days ago

40c3359e

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

bchetioui committed 22 days ago

a5854e44

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

Google-ML-Automation committed 22 days ago

1066aa7b

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

superbobry committed 23 days ago

150a6d9d

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

superbobry committed 23 days ago

79fd127f

Regenerate python stubs for nanobind 2.11.0.

hawkinsp committed 23 days ago

6fb3f3da

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

Google-ML-Automation committed 23 days ago

3a93274b

Merge pull request #34462 from alekstheod:add_support_of_building_jax_under_umbrella_workspace

Google-ML-Automation committed 24 days ago

4a5ad80e

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

loislo committed 24 days ago

50b394dc

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

Google-ML-Automation committed 24 days ago

91871a92

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

bchetioui committed 24 days ago

01d76c43

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

danielsuo committed 24 days ago

46c02566

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

Google-ML-Automation committed 24 days ago

9c5d90dd

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

bchetioui committed 24 days ago

139f212f

Fix ValAccum to initialize Zero using .to_cotangent_aval()

yashk2810 committed 24 days ago

12462e05

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

bchetioui committed 24 days ago

457340ea

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

bchetioui committed 24 days ago

fac19bbf

mattjj committed 25 days ago

39cc010b

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

bchetioui committed 25 days ago

d751a3f7

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

yueshengys committed 25 days ago

b1359ede

[ci] add pyrefly pre-commit hook

Jake VanderPlas committed 25 days ago

382ed1b7

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

yashk2810 committed 25 days ago

b582e359

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

bchetioui committed 25 days ago

e786d22e

Clean up `LoadedExecutable::GetReadyFuture()`

junwhanahn committed 25 days ago

f0eddf93

[indexing] avoid printing full array repr in indexing-related error messages.

Jake VanderPlas committed 25 days ago

17b1bfe7

Fix test behavior.

yueshengys committed 25 days ago

54add9ec

Removed a few unused functions

superbobry committed 25 days ago

b58b083e

Merge pull request #34710 from jakevdp:ref-indexing

Google-ML-Automation committed 25 days ago

ab376494

Older