ROCm/jax

Pull Requests Commits

[ROCm] Added back compatibility test for hipsolver_sytrd and marked it as stable.

tsrw2048 committed 20 hours ago

a065714d

Merge pull request #34809 from justinjfu:explicit_sharding

Google-ML-Automation committed 21 hours ago

0bcf6e59

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

golechwierowicz committed 1 day ago

accb7199

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

Google-ML-Automation committed 1 day ago

05d9eec3

Add reshard rule for mosaic gpu

justinjfu committed 1 day ago

a67c84af

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

Google-ML-Automation committed 1 day 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 1 day ago

3ddf0685

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

Google-ML-Automation committed 1 day ago

c9ef1c10

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

jakevdp committed 1 day ago

6156decd

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

superbobry committed 1 day ago

b23e3a44

[pmap] Add reshard jet rule.

danielsuo committed 1 day ago

dc1eedbe

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

belitskiy committed 1 day ago

c9dbcaa9

Merge pull request #34798 from jakevdp:fix-lint

Google-ML-Automation committed 1 day ago

e1055fdd

Fix lint at head

jakevdp committed 1 day ago

6f6141d2

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

superbobry committed 1 day ago

30d7f945

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

Google-ML-Automation committed 1 day 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 1 day ago

5272ac6b

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

Google-ML-Automation committed 1 day ago

15accf5f

Fix duplicate word typos in documentation

hrideymarwah15 committed 1 day ago

191a2427

[lapack] Parallelize SVD computations across the batch dimension.

basioli-k committed 1 day ago

460d6b5c

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

Google-ML-Automation committed 2 days ago

5321e425

Improve OverFlow error raised by jit

yashk2810 committed 2 days ago

b9fdde25

Fix `get_sharding_for_vmap` to preserve unreduced

yashk2810 committed 2 days ago

14ce7f28

Update transpose rules of all primitives that define an unreduced_rule to call `.to_cotangent_aval()` on residuals.

yashk2810 committed 2 days ago

57a4f6a0

Propagate memory spaces from inputs to outputs when aliased.

brianwa84 committed 2 days ago

5f5be6f1

Switch all Windows RBE builds to use `ml-oss-rbe-general`

nitins17 committed 2 days ago

41256530

Save CPU version of collective metadata to resources so it can be used during execution.

Google-ML-Automation committed 2 days ago

b2c727e3

Merge pull request #34770 from ROCm:add_support_of_building_jax_under_umbrella_workspace_new2

Google-ML-Automation committed 2 days ago

68992a18

Address review comments

alekstheod committed 2 days ago

60f50910

[mosaic] Added a new tpu.fetch_and_add_sync op

superbobry committed 2 days ago

2f1645f1

Older