ROCm/jax

Pull Requests Commits

[Mosaic GPU] Rework the CUDA_ROOT detection once again

apaszke committed 227 days ago

51fea783

[mosaic] Added a `k` prefix to `TPU_MemorySpace` members

superbobry committed 227 days ago

6348ed9f

[Mosaic:TPU] Byte-granularity dynamic gathers

tlongeri committed 227 days ago

66795a13

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

yashk2810 committed 227 days ago

361cf02d

Add `gather` to `roofline`.

zacmustin committed 227 days ago

4fb6f1bc

Postrelease (0.6.2) changes

yashk2810 committed 227 days ago

98e91122

Prepare for JAX release 0.6.2

yashk2810 committed 227 days ago

2a1f66b9

[JAX] Remove sleeping from colocated Python execution tests

hyeontaek committed 227 days ago

f3e5c89f

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

zacmustin committed 227 days ago

8adb773c

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

Jake VanderPlas committed 227 days ago

95f5eacf

add psend and precv to jax/lax/parallel

rosiezou committed 227 days ago

1f46aa4d

Update XLA dependency to use revision

Google-ML-Automation committed 227 days ago

6803fbfd

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

justinjfu committed 227 days ago

3b5c51f3

add jax.nn module type hints (__init__.pyi)

DanisNone committed 227 days ago

d17dc18f

jax.experimental.enable_x64: add warning to docstring

jakevdp committed 227 days ago

54e34127

Removed fixed suppressions

vfdev-5 committed 227 days ago

cfbfd287

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

superbobry committed 227 days ago

dbb7794d

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

Google-ML-Automation committed 227 days ago

6a7c4d3a

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

cky9301 committed 227 days ago

a3b8c508

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

Google-ML-Automation committed 227 days ago

2d2c58b0

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

jburnim committed 227 days ago

b2d89da8

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

WindQAQ committed 227 days ago

52ae9e22

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

ZixuanJiang committed 227 days ago

a0aa266a

Update XLA dependency to use revision

Google-ML-Automation committed 227 days ago

4c3f40f0

Add test for programmatic tracing with options.

sannidhyachauhan committed 227 days ago

412d80fb

Removing Tensorflow references from the document.

sannidhyachauhan committed 227 days ago

bff97363

Add custom-call ops to roofline.

zacmustin committed 227 days ago

d71e5fcd

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

hyeontaek committed 227 days ago

bb0291e4

[Mosaic:TPU][NFC] Delete unused variable

tlongeri committed 227 days ago

8c14f6cc

Fix a missing bounds check in traceback code.

hawkinsp committed 227 days ago

2fa8e357

Newer Older