ROCm/jax

Pull Requests Commits

[ROCm]: Add ROCm command buffer support for triton kernel

Rahul Batra committed 2 years ago

f01c27f6

Merge pull request #19662 from skye:version

a-googler committed 2 years ago

789ef4b0

jax 0.4.24 release take 2: update libtpu version

skye committed 2 years ago

4cca5fa1

Removed unnecessary calls to jax_triton.get_triton_type

superbobry committed 2 years ago

5c5e8f03

[shape_poly] Add micro-benchmarks for symbolic shape manipulation.

gnecula committed 2 years ago

313caba4

Remove an excessive log statement from the create_device_mesh function.

a-googler committed 2 years ago

77d41289

Inlined all builder.create_* calls

superbobry committed 2 years ago

0596c804

compile_ttir_to_ptx_inplace no longer requires an IR context

superbobry committed 2 years ago

c0c6a450

Migrated tl.semantic.cast to lower directly to MLIR

superbobry committed 2 years ago

c77d45d5

Update XLA dependency to use revision

a-googler committed 2 years ago

f2613387

Update XLA dependency to use revision

a-googler committed 2 years ago

2a5af098

Typo: MT19937 is 2.5 kilobytes, not kilobits.

a-googler committed 2 years ago

2cfbd351

Inlined some Triton-specific abstractions

superbobry committed 2 years ago

06d32808

Update XLA dependency to use revision

a-googler committed 2 years ago

a44c0e0a

Add `--use_clang` and `--clang_path` options to `build.py`

ddunl committed 2 years ago

6928465b

[Pallas] Add interpret mode support for dynamic grid

sharadmv committed 2 years ago

a7a6b40b

Only register profiler when the plugin is dynamically loaded.

Jieying Luo committed 2 years ago

c4b62660

Removed wrap_with_builder and tensor.to from the Triton compatibility layer

superbobry committed 2 years ago

cda40ece

Migrated dot to lower directly to Triton IR

superbobry committed 2 years ago

28eff4f9

Migrated store/load to lower directly to Triton IR

superbobry committed 2 years ago

5867a05c

[jax_triton] Only use side stream to do autotuning when doing graph capture

anlunx committed 2 years ago

16636f9c

Added a custom lowering rule for pow which special-cases weak dtypes

superbobry committed 2 years ago

e1ea936f

Fixed a few ir.Value type-tensor dtype mismatches in the Pallas lowering code on GPU

superbobry committed 2 years ago

d9f42c56

[export] Set default native serialization version to 9.

gnecula committed 2 years ago

fdf227e7

Update XLA dependency to use revision

a-googler committed 2 years ago

f7cbd1c4

Disable JAX compilation cache for XLA:CPU

ezhulenev committed 2 years ago

28ef77df

[Pallas/TPU] Allow 1-sized batch dim in vmap of dynamic grid

sharadmv committed 2 years ago

a41385c8

[Pallas/TPU] Add vmap support for dynamic grid

sharadmv committed 2 years ago

d76705da

Merge pull request #19622 from jakevdp:rng-uniform-doc

a-googler committed 2 years ago

b6cda052

Merge pull request #19621 from jakevdp:cholesky

a-googler committed 2 years ago

bf9e3e30

Older