ROCm/jax

Pull Requests Commits

[Mosaic:TPU] Byte-granularity dynamic gathers

tlongeri committed 346 days ago

66795a13

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

yashk2810 committed 346 days ago

361cf02d

Add `gather` to `roofline`.

zacmustin committed 346 days ago

4fb6f1bc

Postrelease (0.6.2) changes

yashk2810 committed 346 days ago

98e91122

Prepare for JAX release 0.6.2

yashk2810 committed 346 days ago

2a1f66b9

[JAX] Remove sleeping from colocated Python execution tests

hyeontaek committed 346 days ago

f3e5c89f

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

zacmustin committed 346 days ago

8adb773c

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

Jake VanderPlas committed 346 days ago

95f5eacf

add psend and precv to jax/lax/parallel

rosiezou committed 346 days ago

1f46aa4d

Update XLA dependency to use revision

Google-ML-Automation committed 346 days ago

6803fbfd

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

justinjfu committed 346 days ago

3b5c51f3

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

DanisNone committed 346 days ago

d17dc18f

jax.experimental.enable_x64: add warning to docstring

jakevdp committed 346 days ago

54e34127

Removed fixed suppressions

vfdev-5 committed 346 days ago

cfbfd287

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

superbobry committed 346 days ago

dbb7794d

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

Google-ML-Automation committed 346 days ago

6a7c4d3a

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

Kuangyuan Chen committed 346 days ago

a3b8c508

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

Google-ML-Automation committed 346 days ago

2d2c58b0

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

jburnim committed 346 days ago

b2d89da8

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

WindQAQ committed 346 days ago

52ae9e22

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

ZixuanJiang committed 346 days ago

a0aa266a

Update XLA dependency to use revision

Google-ML-Automation committed 346 days ago

4c3f40f0

Add test for programmatic tracing with options.

sannidhyachauhan committed 346 days ago

412d80fb

Removing Tensorflow references from the document.

sannidhyachauhan committed 346 days ago

bff97363

Add custom-call ops to roofline.

zacmustin committed 346 days ago

d71e5fcd

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

hyeontaek committed 346 days ago

bb0291e4

[Mosaic:TPU][NFC] Delete unused variable

tlongeri committed 346 days ago

8c14f6cc

Fix a missing bounds check in traceback code.

hawkinsp committed 346 days ago

2fa8e357

[doc] add missing axis_types documentation

jakevdp committed 346 days ago

005acd2f

Remove legacy CPU custom calls.

dfm committed 346 days ago

6140f530

Newer Older