ROCm/jax

Pull Requests Commits

Add valid repo resolution name for making it work as an external project

alekstheod committed 27 days ago

0a2f4cca

Adjust name to external_deps_repository

alekstheod committed 27 days ago

a50a4dec

Rename test_deps to external_deps.bzl

alekstheod committed 27 days ago

d02fcbae

[XLA:MGPU] Port the entire TiledLayout class to C++.

golechwierowicz committed 27 days ago

accb7199

Move deps to gpu_test_deps

alekstheod committed 27 days ago

e615e898

alekstheod committed 27 days ago

87707e6e

Update XLA dependency to use revision http://github.com/openxla/xla/commit/f9ede68ed63446045a0ab73660a80348c405b394

Google-ML-Automation committed 27 days ago

05d9eec3

Address review comments

alekstheod committed 27 days ago

a8b2837d

Address review comments

alekstheod committed 27 days ago

f7894478

Add reshard rule for mosaic gpu

justinjfu committed 27 days ago

a67c84af

Add JAX version metadata to xspace, creating a metadata registry API to add metadata

Google-ML-Automation committed 27 days ago

8ca2421b

Raise NotImplementedError for pcast primitives in their impl rule. `pcast` should be called under shard_map. Fixes https://github.com/jax-ml/jax/issues/34801

yashk2810 committed 27 days ago

3ddf0685

Merge pull request #34734 from jakevdp:indexing-dynamic-slice

Google-ML-Automation committed 28 days ago

c9ef1c10

[indexing] support pallas-like dynamic slicing in jax.numpy indexing

jakevdp committed 28 days ago

6156decd

[pallas:sc] `fetch_and_add` now support multiple indices

superbobry committed 28 days ago

b23e3a44

[pmap] Add reshard jet rule.

danielsuo committed 28 days ago

dc1eedbe

Increase shard count for the newly added scaled_dot_test on GPU and TPU.

belitskiy committed 28 days ago

c9dbcaa9

Merge pull request #34798 from jakevdp:fix-lint

Google-ML-Automation committed 28 days ago

e1055fdd

Fix lint at head

jakevdp committed 28 days ago

6f6141d2

[mosaic] Exported `Tiling` under `jax.experimental.mosaic`

superbobry committed 28 days ago

30d7f945

[Pallas/interpreter] Add interpreter parameters for logging during interpretation of GPU kernels.

Google-ML-Automation committed 28 days ago

65b08f83

Simplify make_jaxpr(return_shape=True) by just returning traced.out_info. Fixes https://github.com/jax-ml/jax/issues/34193

yashk2810 committed 28 days ago

5272ac6b

Merge pull request #34790 from hrideymarwah15:fix-typos-documentation

Google-ML-Automation committed 28 days ago

15accf5f

Fix duplicate word typos in documentation

hrideymarwah15 committed 28 days ago

191a2427

[lapack] Parallelize SVD computations across the batch dimension.

basioli-k committed 28 days ago

460d6b5c

Address review comments

alekstheod committed 28 days ago

2b34afd5

Update XLA dependency to use revision http://github.com/openxla/xla/commit/2734172e27d22c1e774b4d27365074e09e4c1c6d

Google-ML-Automation committed 28 days ago

5321e425

Expand ROCm testing matrix and improve job naming

mminutoli committed 28 days ago

f9d2e1b5

Add JAXCI_ROCM_VERSION environment variable

mminutoli committed 28 days ago

6ad760e6

Add ROCm pytest workflow to nightly/release CI

mminutoli committed 28 days ago

f952348c

Newer Older