ROCm/jax

Pull Requests Commits

kiran-thumma committed 1 year ago

Verified 590ac67c

Merge pull request #169 from ROCm/ci_qa_sync

Ruturaj4 committed 1 year ago

Verified 8985f20e

[ROCm] Update XLA commit hash (#168)

Ruturaj4 committed 1 year ago

fab53b06

[ROCm] Remove exit on error (#167)

Ruturaj4 committed 1 year ago

20362f83

Add argument to override base docker in dockerfile

Ruturaj4 committed 1 year ago

010aea0c

[ROCm] ci build and dockerfile changes

Ruturaj4 committed 1 year ago

3b3083bb

[ROCm] Change run_multi_gpu set opts

Ruturaj4 committed 1 year ago

b8d2dede

BUILD file fixes and linter run

Ruturaj4 committed 1 year ago

a5a14175

[ROCm] fix jax and wheelhouse relative paths

Ruturaj4 committed 1 year ago

addca5d5

[ROCm] ci build fixes

Ruturaj4 committed 1 year ago

39c04506

JAX release 0.4.35

hawkinsp committed 1 year ago

81991d87

Use libtpu releases rather than libtpu-nightly for jax[tpu].

hawkinsp committed 1 year ago

e4f3f8f0

Merge pull request #24465 from jakevdp:fix-mypy

Google-ML-Automation committed 1 year ago

1c6b0a91

Merge pull request #21524 from andportnoy:aportnoy/unknown-platform-lowering-warning

Google-ML-Automation committed 1 year ago

9a2dd19a

Merge pull request #24452 from jakevdp:insert-doc

Google-ML-Automation committed 1 year ago

1e41d5ef

Merge pull request #24467 from andportnoy:patch-2

Google-ML-Automation committed 1 year ago

1a2737b7

jakevdp committed 1 year ago

84985021

Merge pull request #24463 from jakevdp:ml-dtypes-version

Google-ML-Automation committed 1 year ago

92aa9a39

[Pallas TPU] Fix lowering for `jnp.remainder`

ayaka14732 committed 1 year ago

c60bafcc

Add back the import of `jtu` in `flash_attention.py`

andportnoy committed 1 year ago

Verified 63789849

Raise an error when registering a lowering for an unknown platform

andportnoy committed 1 year ago

2aaa108f

[Pallas TPU] Add lowerings for bf16 `jnp.ceil` and `jnp.floor` in TPU v6+

ayaka14732 committed 1 year ago

2b7b0742

Better docs for jnp.insert

jakevdp committed 1 year ago

48dd153e

Merge pull request #24412 from jakevdp:fromfunction-doc

Google-ML-Automation committed 1 year ago

2596a405

Require ml_dtypes >= 0.4.0

jakevdp committed 1 year ago

74fc7360

Better docs for jnp.fromfunction

jakevdp committed 1 year ago

7e38cbd6

Merge pull request #24442 from jakevdp:lexsort-doc

Google-ML-Automation committed 1 year ago

587832f2

Merge pull request #24425 from dfm:rename-vmap-methods

Google-ML-Automation committed 1 year ago

a2e4aff8

[pallas] Support for setting explicit backends to pallas_call.

cperivol committed 1 year ago

4f935636

[Pallas:MGPU] Add support for grid dims in GPUMesh

apaszke committed 1 year ago

2db03ba5

Older