ROCm/jax

Pull Requests Commits

Skip test_tridiagonal_solve_grad on ROCm due to rocSparse issue

AratiGanesh committed 9 hours ago

fa7d45fb

Reverts e121d21893d46a613685363f1321dfe336ffbd2b

Google-ML-Automation committed 10 hours ago

6654cd9c

Allow internal mutable array effect to be DCEd

sharadmv committed 10 hours ago

54603e2e

[pmap] Update migrate pmap docs to include sharded reduction issues.

danielsuo committed 10 hours ago

9bf9b863

[pmap] Cache jit wrappers in pmap's shard_map path to reduce dispatch overhead

danielsuo committed 11 hours ago

280f6b4e

Pass unused recoverable option to coordination service.

mwhittaker committed 12 hours ago

83e8cf58

[Pallas/Mosaic GPU] Temporary disable tests due to the lack of cross-device barrier before kernel execution with collective metadata.

Google-ML-Automation committed 14 hours ago

8f021d1d

[JAX] Add a configuration option jax_compilation_cache_check_contents.

hawkinsp committed 16 hours ago

dd967b05

Temporary disable TMA with collective metadata.

Google-ML-Automation committed 18 hours ago

07026a66

Merge pull request #34908 from ezhulenev:no-ffi-api-0

Google-ML-Automation committed 21 hours ago

9e8222c0

Update XLA dependency to use revision http://github.com/openxla/xla/commit/721dcddd667ae0867be72346d0906185702edc83

Google-ML-Automation committed 22 hours ago

f2952d9c

Make IndivisibleError more understandable by reporting the correct dim_size.

Google-ML-Automation committed 1 day ago

1014bf3d

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

Google-ML-Automation committed 1 day ago

f9a90ff2

Reverts c868c0748d69fd750f125d74843aacfeb2ee84d3

Google-ML-Automation committed 2 days ago

e121d218

Reverts a7a4b6c7e1dceaf9ad49937170997dfa3a135c13

yashk2810 committed 2 days ago

644d5e47

Add APIs to support Mosaic Kernel specialization

Google-ML-Automation committed 2 days ago

c868c074

[jax:ffi] Do not use internal XLA ffi_api.h header

ezhulenev committed 2 days ago

74783118

Update XLA dependency to use revision http://github.com/openxla/xla/commit/9e665ebe7b6fad3bef7147bc79f65440d19bfc58

Google-ML-Automation committed 2 days ago

bb0ce4e8

Move the canonicalization code so that axis=None is handled correctly in explicit mode

yashk2810 committed 3 days ago

51d6376a

In explicit mode, error out by default if the specs of arguments do not match in_specs passed to shard_map. In other words, in_specs in explicit_mode acts as an assert to make sure (`jax.typeof(arg).sharding.spec == in_specs`)

yashk2810 committed 3 days ago

a7a4b6c7

[Mosaic] Add target core type parameter to tpu.enqueue_dma

vsytch committed 3 days ago

9f654f56

Refactor profiler metadata collection to use a registry.

Google-ML-Automation committed 3 days ago

ba8ebca5

[Pallas] Clean up state Transforms

sharadmv committed 3 days ago

f6101063

Make IndivisibleError inherit from ValueError so it can be caught by user code

marksandler2 committed 3 days ago

f9a823b0

[Mosaic:TPU] Fold tpu.erase_layout into memref.load and memref.store

tlongeri committed 3 days ago

7d6fd6da

Increase shard count for gpu_pallas_interpret_test.

belitskiy committed 3 days ago

0f1e9c54

[logging] Avoid attaching the logging handler to jax for "NOTSET".

rdyro committed 3 days ago

63cd5913

Fix a Lowering issue for Pallas Kernel with Dynamic Shapes

Google-ML-Automation committed 3 days ago

4f82c7c3

[pallas:sc] `pltpu.sync_copy` now accepts `add=`

superbobry committed 3 days ago

21efe9a6

[pallas:sc] Added a primitive for atomic synchronous fetch-and-add in SMEM

superbobry committed 3 days ago

89fc1a93

Older