ROCm/jax

Pull Requests Commits

[ROCm]: Add jax-triton support for ROCm

Rahul Batra committed 2 years ago

b4b97cd8

Update XLA dependency to use revision

a-googler committed 2 years ago

eba73a7f

Testing triton integration 2023-09-14

a-googler committed 2 years ago

b1e96284

Merge pull request #18071 from hawkinsp:jitlink

a-googler committed 2 years ago

9ae5a433

Fix typo in setup.py

hawkinsp committed 2 years ago

927a182d

Merge pull request #18066 from jakevdp:update-mypy

a-googler committed 2 years ago

a1c4dc0a

Merge pull request #18062 from hawkinsp:jitlink

a-googler committed 2 years ago

4d247a47

CI: update mypy to v1.6.0

jakevdp committed 2 years ago

a794bebb

Add a version constraint on nvidia-nvjitlink-cu12.

hawkinsp committed 2 years ago

b8c0ca11

Migrate another subset of internal modules to use state objects

superbobry committed 2 years ago

2f70ae70

Update XLA dependency to use revision

a-googler committed 2 years ago

04422b8a

Merge pull request #18052 from jakevdp:prngkey-abstractify

a-googler committed 2 years ago

00e16de3

Optimize make_array_from_callback for fully replicated shardings by going via batched_device_put

yashk2810 committed 2 years ago

fd09b356

[random] add shaped_abstractify handler for custom PRNG key

jakevdp committed 2 years ago

e5c2a2c0

Merge pull request #18048 from superbobry:all-any-list-comp

a-googler committed 2 years ago

899cc304

[pallas] Simplify several primitives bindings.

chr1sj0nes committed 2 years ago

2c6fbc2e

[pallas] Minor cleanup of `Slice` code.

chr1sj0nes committed 2 years ago

4017a7ba

MAINT Use a generator expression with `all()` and `any()`

superbobry committed 2 years ago

5d9c39f4

Merge pull request #18044 from hawkinsp:readme

a-googler committed 2 years ago

5ed69280

Remove calling configure_library_path during jax import and get libtpu path from libtpu_module.get_library_path().

Jieying Luo committed 2 years ago

b81a3e1f

Remove take_ownership support in DLPack.

Jieying Luo committed 2 years ago

269d7ce5

Update references to aarch64 support in the README.

hawkinsp committed 2 years ago

04bee992

Merge pull request #18042 from jakevdp:dtypelike

a-googler committed 2 years ago

8f283bc9

Make jax._src.typing.DTypeLike more strictly defined

Jake VanderPlas committed 2 years ago

911f7457

Merge pull request #18035 from bb515:patch-1

a-googler committed 2 years ago

9edc8031

Define jax.typing.DTypeLike

jakevdp committed 2 years ago

117f4bdf

Merge pull request #18040 from hawkinsp:build2

a-googler committed 2 years ago

f9190006

Set -P when testing whether a package is installed during build.py.

hawkinsp committed 2 years ago

73db6ecf

Merge pull request #18038 from apaszke:mypy-fix

a-googler committed 2 years ago

00dd5b3b

Fix mypy errors in flash_attention

apaszke committed 2 years ago

3a7000ee

Older