ROCm/jax

Pull Requests Commits

[JAX] Add a test verifying the behavior of module-level state accessed by colocated Python

hyeontaek committed 1 year ago

f43d2b68

Merge pull request #26339 from jakevdp:mean-doc

Google-ML-Automation committed 1 year ago

10363663

Merge pull request #26259 from Qazalbash:scipy-expon

Google-ML-Automation committed 1 year ago

1eda5e2e

Merge pull request #26313 from gnecula:debug_info_vjp

Google-ML-Automation committed 1 year ago

c46b0215

doc: add note about f16 casting in jnp.mean

jakevdp committed 1 year ago

9b402ecd

Merge branch 'main' into scipy-expon

Qazalbash committed 1 year ago

Verified 7fc605f7

Compute buffer aliasing on a per buffer basis.

pschuh committed 1 year ago

da0827b7

Refactor JAX wheel build rules to control the wheel filename and maintain reproducible wheel content and filename results.

Google-ML-Automation committed 1 year ago

d424f5b5

[pallas_mgpu] Fix emit_pipeline_with_wgmma test.

petebu committed 1 year ago

9f53dfae

[Mosaic GPU] Write a higher-level tcgen05.mma helper reusing WGMMA implementation

apaszke committed 1 year ago

6d6c8c2e

Return arrays from `ArrayImpl._check_and_rearrange`.

emilyfertig committed 1 year ago

4ae7fcf3

[better_errors] Add debug info to the Jaxprs formed for AD

gnecula committed 1 year ago

abcaec70

Roll back multinomial change from https://github.com/jax-ml/jax/pull/25688

Jake VanderPlas committed 1 year ago

e4dac395

Replace Python 3.12 with Python 3.13 and add Python 3.10 to the matrix

nitins17 committed 1 year ago

bc055569

Run Bazel CPU/CUDA RBE jobs on pushes to main/release branches

nitins17 committed 1 year ago

b5b913ac

#sdy Fix the format for the external link to jax-shardy-guide colab.

bixia1 committed 1 year ago

ae74d3e5

[Mosaic GPU] Add helpers for dealing with TMEM references + implement optimized loads

apaszke committed 1 year ago

f4dab0cf

[pallas:mgpu] Cast all indices to i32 during lowering.

cperivol committed 1 year ago

eeace3ce

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

3a1c63c5

[Mosaic GPU] Refactor the Blackwell matmul example and make it runnable

apaszke committed 1 year ago

b79ab01e

[Mosaic TPU] Add optimized casts for bf16->s4 in TPUv6

apaszke committed 1 year ago

e7a4f893

[Mosaic GPU] Infer whether A/B are row- or column-major from strides

apaszke committed 1 year ago

1fbc4a15

Merge pull request #26159 from andportnoy:aportnoy/mosaic-gpu-blackwell-simple-matmul

Google-ML-Automation committed 1 year ago

d6be2351

Skip broken tests at HEAD

yashk2810 committed 1 year ago

c07b6b52

Merge pull request #26324 from mattjj:skip-tests-with-extra-requirements

Google-ML-Automation committed 1 year ago

9aa98130

Merge pull request #26325 from hawkinsp:tpu2

Google-ML-Automation committed 1 year ago

781172c2

Remove libtpu-nightly dependency from jax[tpu].

hawkinsp committed 1 year ago

b1a2c27a

skip tests with extra requirements

mattjj committed 1 year ago

1ae02bc0

Set the mesh as manual during partial_eval_custom in shard_map so that `_add_reshapes` happens under the correct mesh.

yashk2810 committed 1 year ago

307006e1

[Pallas TPU] Add helpers for writing collectives

sharadmv committed 1 year ago

02f45313

Newer Older