ROCm/jax

Pull Requests Commits

[ROCm]: Pallas updates for ROCm

Rahul Batra committed 2 years ago

5f87fd1e

Merge pull request #21828 from gnecula:exp_calling_convention

a-googler committed 2 years ago

a9edaeb3

Merge pull request #21846 from tilakrayal:patch-1

a-googler committed 2 years ago

3f4c2119

Port Cholesky Factorization to XLA's FFI

Paweł Paruzel committed 2 years ago

3d39b6e7

Fixing the naming conventions in signal.py

tilakrayal committed 2 years ago

Verified 3ef89a21

[export] Rename some API entry points

gnecula committed 2 years ago

7c3a4db3

Replace xla_computation in name_stack_test with `jit(f).lower(...).compiler_ir()`

yashk2810 committed 2 years ago

c5c7fa70

[Pallas] Add missing trace_stop ops for jaxprs that end without a non-scoped op.

justinjfu committed 2 years ago

e96b28c4

Revert: Improve tensorstore I/O efficiency

Google-ML-Automation committed 2 years ago

5462d2e3

Merge pull request #21837 from jakevdp:tree-docs

a-googler committed 2 years ago

cc22b6be

[Rollback] Bumped the minimum ml_dtypes version to 0.4.0

yashk2810 committed 2 years ago

b1f7627c

[Pallas] Allow keys as input to Pallas kernels.

justinjfu committed 2 years ago

4b81680b

Merge pull request #21832 from jakevdp:serialization-version-doc

a-googler committed 2 years ago

b7a8f9d5

Document jax.tree.* directly

jakevdp committed 2 years ago

d82b66f7

Enable runtime uptime telemetry for JAX on Cloud TPU.

Google-ML-Automation committed 2 years ago

dca542cc

Update XLA dependency to use revision

Google-ML-Automation committed 2 years ago

06fe7052

Skip LRUCacheTest if filelock is not installed

superbobry committed 2 years ago

69f437d2

Merge pull request #21823 from superbobry:pallas

a-googler committed 2 years ago

8b849975

Document jax.export serialization version numbers

jakevdp committed 2 years ago

6e837da3

[JAX] Disable qdwh_test in asan/msan/tsan configurations on TPU.

hawkinsp committed 2 years ago

339027d7

Enable jax's cloud-tpu configs when libtpu is present via through "pip install" or set by custom through the $TPU_LIBRARY_PATH env var

Google-ML-Automation committed 2 years ago

987a2f08

Merge pull request #21769 from gnecula:doc_export2

a-googler committed 2 years ago

544975f6

[export] Add documentation for jax.export

gnecula committed 2 years ago

105cc9a1

[PJRT:PLUGIN] Support both string and bytes as the input type of function name for register_custom_call_target in jax-cuda-plugin.

Jieying Luo committed 2 years ago

ad9f35ae

Copybara import of the project:

rahulbatra85 committed 2 years ago

4400ac45

Merge pull request #21799 from gnecula:pallas_cross

a-googler committed 2 years ago

73f67e22

[Mosaic GPU] Change FlashAttention implementation to support Grouped Query Attention.

bchetioui committed 2 years ago

25a47649

Removed BlockSpec.__init__

superbobry committed 2 years ago

c41e52a7

Integrate LLVM at llvm/llvm-project@c012e487b724

Google-ML-Automation committed 2 years ago

a0e5e0f4

[pallas] Add support for cross-platform lowering

gnecula committed 2 years ago

97db0e75

Older