ROCm/jax

Pull Requests Commits

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

yashk2810 committed 213 days ago

1d077471

Add `gather` to `roofline`.

zacmustin committed 213 days ago

34d88cc5

Merge pull request #29528 from jax-ml:postrelease

Google-ML-Automation committed 213 days ago

8b88cc8e

Postrelease (0.6.2) changes

yashk2810 committed 213 days ago

feab6f4b

[JAX] Remove sleeping from colocated Python execution tests

hyeontaek committed 213 days ago

19f34a06

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

zacmustin committed 213 days ago

c944c656

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

Jake VanderPlas committed 213 days ago

7dd13d70

Merge pull request #29135 from rosiezou:main

Google-ML-Automation committed 214 days ago

755bb676

Prepare for JAX release 0.6.2

yashk2810 committed 214 days ago

8f81490a

Update XLA dependency to use revision

Google-ML-Automation committed 214 days ago

e4de90e6

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

justinjfu committed 214 days ago

02688e18

Merge pull request #29410 from DanisNone:nn-type

Google-ML-Automation committed 214 days ago

dc9ef614

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

Google-ML-Automation committed 214 days ago

353e7fac

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

Google-ML-Automation committed 214 days ago

3d37b0d7

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

superbobry committed 214 days ago

c2cc9f9c

add psend and precv to jax/lax/parallel

rosiezou committed 214 days ago

784be1f9

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

Google-ML-Automation committed 214 days ago

0fd08213

jax.experimental.enable_x64: add warning to docstring

jakevdp committed 214 days ago

f22896ac

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

cky9301 committed 214 days ago

332aa354

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

Google-ML-Automation committed 214 days ago

124c723a

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

jburnim committed 214 days ago

077f2b66

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

WindQAQ committed 214 days ago

91651f80

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

ZixuanJiang committed 214 days ago

93453f54

Update XLA dependency to use revision

Google-ML-Automation committed 214 days ago

4be1402b

Add test for programmatic tracing with options.

sannidhyachauhan committed 214 days ago

acf99e5e

Removing Tensorflow references from the document.

sannidhyachauhan committed 214 days ago

0f5cdba0

Add custom-call ops to roofline.

zacmustin committed 214 days ago

2dfaceda

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

hyeontaek committed 214 days ago

b6225124

[Mosaic:TPU][NFC] Delete unused variable

tlongeri committed 214 days ago

748b39f4

Removed fixed suppressions

vfdev-5 committed 214 days ago

Verified 97e580c8

Older