ROCm/jax

Pull Requests Commits

jaxlib_extension_version == 355 after 0.6.2 release. So remove the conditionals.

yashk2810 committed 1 year ago

1d077471

Add `gather` to `roofline`.

zacmustin committed 1 year ago

34d88cc5

Merge pull request #29528 from jax-ml:postrelease

Google-ML-Automation committed 1 year ago

8b88cc8e

Postrelease (0.6.2) changes

yashk2810 committed 1 year ago

feab6f4b

[JAX] Remove sleeping from colocated Python execution tests

hyeontaek committed 1 year ago

19f34a06

Add `cum{logsumexp, min, max, prod, sum}` to JAX roofline.

zacmustin committed 1 year ago

c944c656

Rollback https://github.com/jax-ml/jax/pull/29410 due to downstream pytype failures.

Jake VanderPlas committed 1 year ago

7dd13d70

Merge pull request #29135 from rosiezou:main

Google-ML-Automation committed 1 year ago

755bb676

Prepare for JAX release 0.6.2

yashk2810 committed 1 year ago

8f81490a

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

e4de90e6

[Pallas][Mosaic GPU] Enable collective MMA from TMEM.

justinjfu committed 1 year ago

02688e18

Merge pull request #29410 from DanisNone:nn-type

Google-ML-Automation committed 1 year ago

dc9ef614

Merge pull request #29516 from jakevdp:enable-x64-warning

Google-ML-Automation committed 1 year ago

353e7fac

Merge pull request #29504 from vfdev-5:tsan-ft-removed-fixed-suppression

Google-ML-Automation committed 1 year ago

3d37b0d7

[pallas] `AbstractMemoryRef` now implements all functional update methods via `update`

superbobry committed 1 year ago

c2cc9f9c

add psend and precv to jax/lax/parallel

rosiezou committed 1 year ago

784be1f9

[Pallas TPU] Add flag to enable using registers to keep track of slot info

Google-ML-Automation committed 1 year ago

0fd08213

jax.experimental.enable_x64: add warning to docstring

jakevdp committed 1 year ago

f22896ac

Add an API to overwrite the current execution_stream_id and respect it in XLA CPU dispatch.

Kuangyuan Chen committed 1 year ago

332aa354

Prefer binaries in NVIDIA `nvcc` wheel over system CUDA installation in Mosaic GPU implementation.

Google-ML-Automation committed 1 year ago

124c723a

Update Pallas debugging doc with TPU interpret mode + dynamic race detector.

jburnim committed 1 year ago

077f2b66

[Mosaic] Use BF16 ops for math::PowF on TPUv6+.

WindQAQ committed 1 year ago

91651f80

Pass through the `use_shardy_partitioner` with `jax.config.jax_use_shardy_partitioner`.

ZixuanJiang committed 1 year ago

93453f54

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

4be1402b

Add test for programmatic tracing with options.

sannidhyachauhan committed 1 year ago

acf99e5e

Removing Tensorflow references from the document.

sannidhyachauhan committed 1 year ago

0f5cdba0

Add custom-call ops to roofline.

zacmustin committed 1 year ago

2dfaceda

[JAX] Relax the return type of `colocated_python` decorator

hyeontaek committed 1 year ago

b6225124

[Mosaic:TPU][NFC] Delete unused variable

tlongeri committed 1 year ago

748b39f4

Removed fixed suppressions

vfdev-5 committed 1 year ago

Verified 97e580c8

Older