ROCm/jax

Pull Requests Commits

[ROCm]: Triton updates

Rahul Batra committed 2 years ago

bae76ab7

[ROCm]: Disable sparse tests

invalid-email-address committed 2 years ago

9ab0f61f

Disable linalg_test::EighTinyNorm

Rahul Batra committed 2 years ago

8229fd71

Revert "Switch jaxlib to use nanobind instead of pybind11."

Rahul Batra committed 2 years ago

3fae61c9

[ROCm]: More Fixes

Rahul Batra committed 2 years ago

cf073db4

[ROCm]: Update fused attention test tolerances

Rahul Batra committed 2 years ago

6f3bd2e5

[ROCm] - Port Jax-Triton API to ROCm

Rahul Batra committed 2 years ago

424a72eb

[ROCm]: CI scripts update to handle ROCm versions

Rahul Batra committed 2 years ago

8b75de22

[ROCm]: Fix duplicate deps include

Rahul Batra committed 2 years ago

06a902d3

[ROCm]: Update docker file to ROCm5.6

Rahul Batra committed 2 years ago

a183e7c6

[ROCm]: Add hipblaslt in Dockerfile

Rahul Batra committed 2 years ago

6a1eea07

Merge pull request #17362 from skye:version

a-googler committed 2 years ago

4268ed74

Update versions for jax/jaxlib 0.4.15 release

skye committed 2 years ago

2c5e8577

Update XLA dependency to use revision

a-googler committed 2 years ago

c1680aac

[MOSAIC] apply_vector_layout C++ rewrite (1) VectorLayout functions

tlongeri committed 2 years ago

d02b59e4

Abstract the array_serialization error message to a global variable so that it can be overridden.

yashk2810 committed 2 years ago

6b574708

Raise a good error message when mesh is not provided to jax.jit when using spmd_axis_name parameter of jax.vmap

yashk2810 committed 2 years ago

e785f894

Deprecate a number of exports from jax.interpreters.xla.

hawkinsp committed 2 years ago

9be96c1d

[JAX] Implement importing external dlpack-aware Python arrays.

skye committed 2 years ago

ecee8f91

Remove tests for jax.numpy.in1d, which is deprecated.

hawkinsp committed 2 years ago

e3694455

Remove jax.interpreters.xla.register_collective_primitive.

hawkinsp committed 2 years ago

93900245

Merge pull request #17351 from jakevdp:mypy-fix

a-googler committed 2 years ago

289ccad2

Add segment_ids support to pallas flash attention on GPU.

wang12tao committed 2 years ago

5a578cba

jakevdp committed 2 years ago

f1fc2adf

Any devices passed to jax.sharding.Mesh are required to be hashable.

yashk2810 committed 2 years ago

6072d599

Merge pull request #17337 from jakevdp:jax-extend-doc

a-googler committed 2 years ago

ff5b480c

DOC: add missing docs for jax.random functions

jakevdp committed 2 years ago

f935c006

Make mlir.custom_call() more general and expose it as jax.interpreters.mlir.custom_call().

hawkinsp committed 2 years ago

d0a6813e

Update XLA dependency to use revision

a-googler committed 2 years ago

ba8c4c30

Use the default CSR matmul algorithm.

hawkinsp committed 2 years ago

46ac9e21

Older