ROCm/jax

Pull Requests Commits

Updated the Triton lowering of jnp.transpose

superbobry committed 1 year ago

60d03285

Moved atomic_* operations out of the Triton compatibility layer

superbobry committed 1 year ago

b132c176

Inlined tl.core._to_tensor

superbobry committed 1 year ago

fee7e0ca

[ROCm]: More changes for jax-triton

Rahul Batra committed 1 year ago

504c3f02

_triton_ext no longer links in MLIR C APIs

superbobry committed 1 year ago

cd45df3e

Add MLIR_CAPI_EXPORTED to Triton dialect C bindings

superbobry committed 1 year ago

5c5ce1cc

jaxlib/triton: ensure python 3.9 compatibility

jakevdp committed 1 year ago

d11e2b75

[ROCm]: Add jax-triton support

Rahul Batra committed 1 year ago

435bd8a1

[ROCm]: Add support to continue on fail

Rahul Batra committed 1 year ago

f6d31628

Disable linalg_test::EighTinyNorm

Rahul Batra committed 2 years ago

8a11b400

Fixed a typo in min/max Triton lowering rules

superbobry committed 2 years ago

9e94e6ef

Merge pull request #19667 from jakevdp:array-empty-repr

a-googler committed 2 years ago

b53f7572

Merge pull request #19634 from jakevdp:key-reuse-scan

a-googler committed 2 years ago

c1c0c1cf

Improve repr for empty jax.Array

jakevdp committed 2 years ago

d9cbd7bd

Merge pull request #19629 from jakevdp:key-reuse-pjit

a-googler committed 2 years ago

69a9f7f4

Merge pull request #19599 from ROCm:rocm-add-triton_command_buffer

a-googler committed 2 years ago

206398a2

Merge pull request #19665 from jakevdp:disable-jit-doc

a-googler committed 2 years ago

be994512

document that under disable_jit, individual primitives are still compiled

jakevdp committed 2 years ago

82611eb8

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

Rahul Batra committed 2 years ago

f01c27f6

Update XLA dependency to use revision

a-googler committed 2 years ago

e224c3d6

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

Older