ROCm/jax

Pull Requests Commits

Fix lowering code for ROCm RNN

Ruturaj4 committed 303 days ago

45d8d852

Change CODEOWNERS (#237)

charleshofer committed 326 days ago

Verified 9cac506b

Use bazel for PR tests (#216)

charleshofer committed 327 days ago

Verified dd3f34ca

Add CODEOWNERS file (#236)

charleshofer committed 330 days ago

Verified 60f51d21

Use a GitHub app for syncing rocm-main and upstream main (#224)

charleshofer committed 338 days ago

Verified 79c0cd46

Merge pull request #229 from ROCm/ci-upstream-sync-112_1

github-actions[bot] committed 340 days ago

Verified 81f8e65c

Merge branch 'rocm-main' into ci-upstream-sync-112_1

charleshofer committed 340 days ago

3745591d

Merge pull request #227 from ROCm/ci-upstream-sync-110_1

github-actions[bot] committed 340 days ago

Verified 073bc928

Change to make CI run

charleshofer committed 340 days ago

9133253c

[Mosaic GPU] Add support for collective MMA in the Blackwell matmul example

apaszke committed 340 days ago

74e86bab

[Mosaic GPU] Add support for non-multicast `.cta_group::2` async_copies

apaszke committed 340 days ago

21598d02

[pallas:triton] Lift dot_general restriction on minimal tile size for a.

Google-ML-Automation committed 340 days ago

2c165bff

Add "sequential_unrolled" vmap method for callbacks.

dfm committed 340 days ago

c502332e

[Mosaic GPU] Add support for 2-CTA MMA on Blackwell

apaszke committed 340 days ago

849ea268

[Mosaic GPU] Handle TMEM allocation in the compiler

apaszke committed 340 days ago

0209eee1

Update XLA dependency to use revision

Google-ML-Automation committed 340 days ago

5a223516

Merge pull request #26399 from gnecula:debug_info_jaxpr_6

Google-ML-Automation committed 340 days ago

49eccd6c

Disable XLA autotuning to speed up tests

apaszke committed 340 days ago

c171fc60

[better_errors] Continue adding debug info to Jaxprs (step 6)

gnecula committed 340 days ago

550d1aa1

Merge pull request #26450 from gnecula:debug_info_win2

Google-ML-Automation committed 340 days ago

e7acb20e

Fix debug_info_test.py regexp for Windows paths.

gnecula committed 341 days ago

872b8c3a

#jax Optimize `jax.numpy.take_along_axis` along the dimension satisfies

ZixuanJiang committed 341 days ago

4b1400db

Fix Python version test against Python 3.12.

hawkinsp committed 341 days ago

1e447c8a

Re-enable oss paged attn kernel

Google-ML-Automation committed 341 days ago

d3ed6ca0

[TPU[Mosaic] Fix missing sfences in smem DMAs

Google-ML-Automation committed 341 days ago

ffd3faad

Add `--config=cuda_libraries_from_stubs` in the end of all additional bazel options for CUDA wheels.

Google-ML-Automation committed 341 days ago

e64650e2

Merge pull request #26423 from gnecula:debug_info_jaxpr_7

Google-ML-Automation committed 341 days ago

b7d01228

Fix missing batch partitioning hook in GPU plugin.

dfm committed 341 days ago

ba8c3a92

[StableHLO] Allow composites to use dtype values in attributes

GleasonK committed 341 days ago

319f7f5a

Merge pull request #26426 from NKlug:nklug-linalg-solve-docu

Google-ML-Automation committed 341 days ago

6a638ac8

Older