ROCm/jax

Pull Requests Commits

[ROCm] Use compiler-rt builtins for all clang-based ROCm builds

mminutoli committed 127 days ago

1fdea476

Merge pull request #37564 from Saumay:add-scipy-special-boxcox

Google-ML-Automation committed 127 days ago

5799b549

Don't squeeze singleton dimension if their spec is not None in jnp.matmul.

yashk2810 committed 128 days ago

085710a8

Add a jax.extend.xla.register_hlo_module_transformation API.

hawkinsp committed 128 days ago

3a5e9c27

Roll back PR #37872; it led to timeouts in downstream projects.

Jake VanderPlas committed 128 days ago

92ae3541

Consume memory space + core type in isGather instead of MemRefs.

nvgrw committed 128 days ago

f3a724d8

[JAX][CPU] Skip running `testMathSpecialFloatValues` on ARM for `float16` `tanh`.

bchetioui committed 128 days ago

80df7b53

[pallas:mgpu] Deprecated using `pl.pallas_call` for Mosaic GPU kernels

superbobry committed 128 days ago

395f039c

[pallas:sc] Removed a few unnecessary test skips

superbobry committed 128 days ago

709095b9

[Mosaic GPU] Support slicing for non defaults TMEM layouts with 128 rows.

allanrenucci committed 128 days ago

80e49d9a

[pallas] Removed tests for cross-platform exports

superbobry committed 128 days ago

e59492a5

[pallas:sc] Do not skip SC tests on Cloud TPUs

superbobry committed 128 days ago

ae32a398

[pallas:mosaic] Removed a few unnecessary skips in tpu_pallas_distributed_test

superbobry committed 128 days ago

55aa2cf9

[pallas:sc] Do not skip test_gather_add on 7x

superbobry committed 128 days ago

a4caecad

Add TPU v8i to enum

Google-ML-Automation committed 128 days ago

a7d70b2a

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

Google-ML-Automation committed 128 days ago

74198ebe

Merge pull request #38037 from olupton:sm120-skip

Google-ML-Automation committed 128 days ago

05e445b7

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

Google-ML-Automation committed 129 days ago

58d4bf81

Add jax.scipy.special.boxcox and boxcox1p

Saumay committed 129 days ago

8e5dc0eb

Merge pull request #38080 from mattjj:tri-solve-rules

Google-ML-Automation committed 129 days ago

8880324e

restore triangular_solve batching rule, improve sharding rule

mattjj committed 129 days ago

507e331e

Support while_p in GPU Interpret Mode.

jburnim committed 130 days ago

5c8d2bed

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

Google-ML-Automation committed 130 days ago

7ec563a0

Merge pull request #38075 from jakevdp:geometric-fix

Google-ML-Automation committed 130 days ago

edfffefa

Fix export and typed prng key. While lowering call_exported, if the out_avals were a typed prng key, it was not being converted to physical avals before lowering.

yashk2810 committed 130 days ago

a12e5d56

Delete not_mapped and replace it with None

yashk2810 committed 130 days ago

66e0fd37

Merge pull request #37674 from TheSkyentist:feat/dawsn

Google-ML-Automation committed 130 days ago

9e9184de

Merge pull request #38065 from magaonka-amd:fix/rocm-skip-complex-gemm-ci-tests

Google-ML-Automation committed 130 days ago

2d7a47df

[pallas] Mark additional higher-order primitives as uncacheable.

hawkinsp committed 130 days ago

b8a49c58

Turn verifyScatter / verifyGather helpers into free functions.

nvgrw committed 130 days ago

512586c3

Older