ROCm/jax

Pull Requests Commits

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

43f2da86

[pallas:mosaic_gpu] Added WG lowering rules for TMA primitives and run_scoped_p

superbobry committed 1 year ago

7eadc64b

[Pallas:MGPU] Consistently use i32 as the grid index type in emit_pipeline

apaszke committed 1 year ago

80848ad8

Remove `-y` and extra `--pre` arguments

nitins17 committed 1 year ago

cc830748

Pass in "JAX_" variables to the docker container

nitins17 committed 1 year ago

e4d63dff

[sharding_in_types] Error out when using `auto_axes` or `explicit_axes` API when there is no context mesh.

yashk2810 committed 1 year ago

b707f0bd

Sharp bits: add note on subnormal flush-to-zero

jakevdp committed 1 year ago

f5ca46f5

Merge pull request #26715 from Rifur13:normalize

Google-ML-Automation committed 1 year ago

41faf51a

Add a param to not normalize the attention weights

Rifur13 committed 1 year ago

8f1dd02e

Add sharding mismatch to explain_tracing_cache_miss

yashk2810 committed 1 year ago

6f8bab3c

Merge pull request #26713 from hawkinsp:docs4

Google-ML-Automation committed 1 year ago

d2867333

Change documentation to recommend libtpu from pypi instead of GCS.

hawkinsp committed 1 year ago

85add667

Merge pull request #26708 from hawkinsp:postrelease

Google-ML-Automation committed 1 year ago

11b45a6f

Update version numbers after 0.5.1 release.

hawkinsp committed 1 year ago

c8c4cfa0

Merge branch 'release/0.5.1' into main

hawkinsp committed 1 year ago

54f70724

Share the logic of allowing propagation to input/output between JAX and australis

yashk2810 committed 1 year ago

8739232f

Remove extra "--system" argument

nitins17 committed 1 year ago

cadb5311

Merge pull request #26690 from h-vetinari:gcc

Google-ML-Automation committed 1 year ago

d0cb30b7

Make mesh and *_spec parameters optional.

matthiaskramm committed 1 year ago

79e1e1fc

Fix shard_map debug_nan leakage of manual out_avals in the impl rules of jit i.e. impl rule of jit saw a manual out_aval which is not expected. This is a band-aid for now with a TODO to do a proper fix

yashk2810 committed 1 year ago

6d8be966

Use `uv` to install Python packages

nitins17 committed 1 year ago

4b4f2f9c

Update JVP rule for lax.linalg.lu to use vmap instead of broadcasted_iotas.

dfm committed 1 year ago

62530d59

Fix rank promotion error in JVP of batched eigh.

dfm committed 1 year ago

6bd99207

Plumb layout through the creation of PjRtArrays.

emilyfertig committed 1 year ago

9d421c91

Update lax.linalg.svd primitive to use registration helper functions.

dfm committed 1 year ago

ae656e15

Prepare for JAX release 0.5.1

yashk2810 committed 1 year ago

07440f4a

Merge pull request #25053 from JanLuca:gesvd

Google-ML-Automation committed 1 year ago

c74f497e

Mark error_check_test as thread-unsafe.

hawkinsp committed 1 year ago

aeff675e

Merge pull request #26569 from gnecula:debug_info_arg_names

Google-ML-Automation committed 1 year ago

c17ea805

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

b0cfcb8e

Newer Older