ROCm/jax

Pull Requests Commits

Fix run_multi_gpu script multi-gpu issue and refactor code

Ruturaj4 committed 337 days ago

8e88adcd

[sharding_in_types] Handle ShapeDtypeStruct inputs with sharding_in_types by registering the sharding on the aval properly created by SDS in it's pytype_aval_mapping.

yashk2810 committed 338 days ago

c72ed260

Pin actions/checkout to a commit

nitins17 committed 338 days ago

a1bbad68

Merge pull request #25875 from jax-ml:issue-25517

Google-ML-Automation committed 338 days ago

29a3dded

Mention expected tangent aval in error message, see #25517.

dougalm committed 338 days ago

7d11d12b

#sdy dynamically choose which `custom_partitioning` API to use based on the current

bartchr808 committed 338 days ago

74e912c3

Merge pull request #25798 from gnecula:fix_fori_error

Google-ML-Automation committed 338 days ago

4f2f5fa5

enable partitionable threefry by default

froystig committed 338 days ago

a60ead6f

[Pallas TPU] Add vector support to `pl.debug_print`

ayaka14732 committed 339 days ago

9ba1fd28

[Mosaic GPU] Fix layout API bugs.

justinjfu committed 339 days ago

f69592ae

Merge pull request #25833 from jax-ml:more-linearize-fixes

Google-ML-Automation committed 339 days ago

dabe27bc

Even more linearize fixes

dougalm committed 339 days ago

96769f96

Internal change.

dfm committed 339 days ago

e72c1484

Update XLA dependency to use revision

Google-ML-Automation committed 339 days ago

9ee25e7e

[Pallas TPU] Skip cast test incompatible with older libtpu builds

apaszke committed 339 days ago

aa51f2af

[Mosaic TPU] Add support for arith.fptosi with non-32bit source and target types

apaszke committed 339 days ago

391bad8f

[better_errors] Ensure that tracer errors in for_loop points to use code

gnecula committed 339 days ago

3faff78c

Add JAX unit test for Shardy which causes the compiler to introduce the `mlir::tensor::TensorDialect`. This was causing the compiler to crash.

bartchr808 committed 339 days ago

c14e5b43

Use thread-safe initialization of LAPACK kernels.

hawkinsp committed 339 days ago

91ffb640

[Mosaic TPU] Enable non-sublane-aligned bf16 2D load/stores for earlier TPU gens

tlongeri committed 339 days ago

78520455

[Mosaic TPU][NFC] Remove redundant num_subelems attribute from CreateSubelementMaskOp

tlongeri committed 340 days ago

09302899

Update XLA dependency to use revision

Google-ML-Automation committed 340 days ago

ed349e45

[shardy] Fix cases in shardy where you have a nullary function with partially specified out_shardings (i.e. some out_sharding's are None and others are NamedShardings).

yashk2810 committed 340 days ago

6b253b2f

Update XLA dependency to use revision

Google-ML-Automation committed 341 days ago

4e729316

Fix pmap with sharded typed prng key

yashk2810 committed 342 days ago

aeac6b03

[sharding_in_types] Introduce `auto_mode`, `user_mode`, `auto_mode_ctx` and `user_mode_ctx` as **private** APIs to make writing auto/user sharding in types code way easier and noise-free.

yashk2810 committed 342 days ago

a817f532

[Mosaic GPU] Allow multiple gmem indexers on copies.

justinjfu committed 342 days ago

8e86bede

[Mosaic GPU] Enable loop carries in the pipeline emitter.

justinjfu committed 342 days ago

73b64b8e

[Mosaic][TPU] Add a compatibility mode to Mosaic's canonicalization pass, skipping over elementwise and matmul op insertions and/or type compat casts.

Google-ML-Automation committed 342 days ago

a16fbffc

Merge pull request #25817 from jakevdp:deprecation-utils

Google-ML-Automation committed 342 days ago

743872df

Older