ROCm/jax

Pull Requests Commits

Add missing subprocess and re imports

phambinhfin committed 46 days ago

82484b48

Add JAX version check for is_finite skip on ROCm

phambinhfin committed 55 days ago

37268302

Use specific exception types in get_rocm_shared_memory_limit

phambinhfin committed 55 days ago

d2ca35a7

Compress shared memory check logic in test_dot

phambinhfin committed 55 days ago

2125ee20

Fix and enable Pallas ops tests on ROCm

phambinhfin committed 55 days ago

09cd8c8b

Merge pull request #35309 from superbobry:pyrefly

Google-ML-Automation committed 55 days ago

4d3395f0

[pyrefly] Fixed or suppressed type errors in Pallas backends

superbobry committed 55 days ago

9ce7295f

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

Google-ML-Automation committed 56 days ago

8dac162f

Merge pull request #35306 from mattjj:remat-jaxpr

Google-ML-Automation committed 56 days ago

dea31b6c

[hijax][remat3] add basic remat_jaxpr

mattjj committed 56 days ago

b17e41dd

Merge pull request #35304 from mattjj:remat-transform

Google-ML-Automation committed 56 days ago

53d6b350

[hijax][remat3] start implementing remat on hijax

mattjj committed 56 days ago

5f43a720

Merge pull request #35302 from mattjj:hijax-dce-fix

Google-ML-Automation committed 56 days ago

0cff9fe3

[hijax] fix followup on #35301

mattjj committed 56 days ago

acad3005

Merge pull request #35301 from mattjj:better-hijax-dce-interface

Google-ML-Automation committed 56 days ago

7f80f430

[hijax] improved dce interface

mattjj committed 56 days ago

b1981cbb

Canonicalize indexer.slice_sharding. Fixes https://github.com/jax-ml/jax/issues/35273

yashk2810 committed 56 days ago

7fbbe019

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

Google-ML-Automation committed 57 days ago

4f22f9d2

Merge pull request #35285 from mattjj:is-vjp-bit

Google-ML-Automation committed 57 days ago

4181fe49

[hijax] add is_vjp bit to LinearizeTrace

mattjj committed 57 days ago

79f3519e

Propagate input sharding to output in full_like if it's fully replicated. Then even if input/output shapes don't match, it should be fine. Fixes https://github.com/jax-ml/jax/issues/35273

yashk2810 committed 57 days ago

3b6fd04e

Fix pallas fuser usage of call_hi_primitive_p

yashk2810 committed 57 days ago

1c781867

Remove dead code

yashk2810 committed 57 days ago

f2f1d3ab

Fix some jax tests when custom_vjp3 defaults to True. There are still 6 test targets failing.

yashk2810 committed 57 days ago

1371c2ba

Remove primitive_transpose registrations for primitives that have a transpose_fancy rule defined

yashk2810 committed 57 days ago

6481bb28

Update rules_cc version to 0.2.9

akuegel committed 57 days ago

8b3bc0b4

[Mosaic:GPU] Use 2 GPUs for collective metadata test

PatriosTheGreat committed 57 days ago

1ff638f9

[Mosaic GPU] Add support for tcgen05 MMAs that are both sparse and block-scaled

apaszke committed 58 days ago

4553a01e

Re-enable CI for B200 MGPU test + increase shard_count

apaszke committed 58 days ago

cf870cbe

[pallas:mosaic] Renamed the `KernelType` enum to `CoreType`

superbobry committed 58 days ago

1d0504b1

Older