ROCm/jax

Pull Requests Commits

improve AOT input sharding/layout mismatch error message

froystig committed 1 year ago

7c885373

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

7414f8ce

Fix creation of sharding rules for fused_attention_stablehlo.

tomnatan30 committed 1 year ago

b1af72db

[JAX] Enable Shardy by default in JAX.

tomnatan30 committed 1 year ago

df8e17a5

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

109bcae8

Merge pull request #29370 from mattjj:mutable-array-custom-vjp-scan3

Google-ML-Automation committed 1 year ago

860a29d7

Merge pull request #30340 from mattjj:mutable-array-shmap

Google-ML-Automation committed 1 year ago

213fa6fc

Separate out `vma` and `unreduced/reduced` in the ShapedArray str_short().

yashk2810 committed 1 year ago

0acbd3b3

[mutable-arrays] basic shard_map + mutable arrays support

mattjj committed 1 year ago

ee5c13b6

[mutable-arrays] remat discharge rule

mattjj committed 1 year ago

29e6f4b6

Add JAX tests for deadlock verifier

Google-ML-Automation committed 1 year ago

44b6204b

Fix CI config for non-rbe CUDA tests with py_import.

Google-ML-Automation committed 1 year ago

0e60c6d4

#sdy Fix forward of making JAX changes so we can fall back to GSPMD in JAX export if the loaded module was lowered for GSPMD.

tomnatan30 committed 1 year ago

6635d8f3

[mosaic] `tpu_custom_call.CostEstimate` is now a typed dict

superbobry committed 1 year ago

0b53b97f

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

254c7751

Reverts 7c394f41574b1f669cbb6a47fd03bf1925a19444

voznesenskym committed 1 year ago

0c9582fb

[Pallas/TPU] Interpret mode: Don't use source_info context manager if it's None

sharadmv committed 1 year ago

b0d87e61

Use `operand.aval.sharding` in convert_element_type's transpose rule

yashk2810 committed 1 year ago

dd36a9f7

Integrate LLVM at llvm/llvm-project@06ae0c2a1086

slackito committed 1 year ago

8c8b013a

Remove the assert from scatter_addsub_transpose_rule. Fixes https://github.com/jax-ml/jax/issues/30327. We don't need to use the `out_sharding` in transpose rule anyways that's bound to the primitive.

yashk2810 committed 1 year ago

59e15b9d

Fixed a bug previously hidden by jax_use_direct_linearize=True

brills committed 1 year ago

e394e5bd

Move ShapeDtypeStruct to core.py to break circular deps

yashk2810 committed 1 year ago

fec027ec

[Pallas/Mosaic GPU] Disable broken export test below cc 9.0.

bchetioui committed 1 year ago

b637d4dd

Create workflow for testing bazel cuda non rbe + py import

sgarciagoogle committed 1 year ago

3ff22786

Replace `jax.stages.OutInfo` with `jax.ShapeDtypeStruct`

yashk2810 committed 1 year ago

eb194e9a

#sdy condition jax export tests on compatibility version when using shardy

liepieshov committed 1 year ago

c5c89953

[Mosaic GPU] Fix barrier indexing in tests.

bchetioui committed 1 year ago

c68b2893

Exit the CUDA test job early if not all wheels can be downloaded.

nitins17 committed 1 year ago

bd4cbc8c

Add `eval_shape` as a method of `Traced`.

yashk2810 committed 1 year ago

dd59b47c

#sdy fix-forward jaxlib compatibility issue

liepieshov committed 1 year ago

218560df

Older