ROCm/jax

Pull Requests Commits

improve AOT input sharding/layout mismatch error message

froystig committed 236 days ago

7c885373

Update XLA dependency to use revision

Google-ML-Automation committed 236 days ago

7414f8ce

Fix creation of sharding rules for fused_attention_stablehlo.

tomnatan30 committed 236 days ago

b1af72db

[JAX] Enable Shardy by default in JAX.

tomnatan30 committed 236 days ago

df8e17a5

Update XLA dependency to use revision

Google-ML-Automation committed 237 days ago

109bcae8

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

Google-ML-Automation committed 237 days ago

860a29d7

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

Google-ML-Automation committed 237 days ago

213fa6fc

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

yashk2810 committed 237 days ago

0acbd3b3

[mutable-arrays] basic shard_map + mutable arrays support

mattjj committed 237 days ago

ee5c13b6

[mutable-arrays] remat discharge rule

mattjj committed 237 days ago

29e6f4b6

Add JAX tests for deadlock verifier

Google-ML-Automation committed 237 days ago

44b6204b

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

Google-ML-Automation committed 237 days 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 237 days ago

6635d8f3

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

superbobry committed 237 days ago

0b53b97f

Update XLA dependency to use revision

Google-ML-Automation committed 238 days ago

254c7751

Reverts 7c394f41574b1f669cbb6a47fd03bf1925a19444

voznesenskym committed 238 days ago

0c9582fb

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

sharadmv committed 238 days ago

b0d87e61

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

yashk2810 committed 238 days ago

dd36a9f7

Integrate LLVM at llvm/llvm-project@06ae0c2a1086

slackito committed 238 days 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 238 days ago

59e15b9d

Fixed a bug previously hidden by jax_use_direct_linearize=True

brills committed 238 days ago

e394e5bd

Move ShapeDtypeStruct to core.py to break circular deps

yashk2810 committed 238 days ago

fec027ec

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

bchetioui committed 238 days ago

b637d4dd

Create workflow for testing bazel cuda non rbe + py import

sgarciagoogle committed 238 days ago

3ff22786

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

yashk2810 committed 238 days ago

eb194e9a

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

liepieshov committed 238 days ago

c5c89953

[Mosaic GPU] Fix barrier indexing in tests.

bchetioui committed 238 days ago

c68b2893

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

nitins17 committed 238 days ago

bd4cbc8c

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

yashk2810 committed 238 days ago

dd59b47c

#sdy fix-forward jaxlib compatibility issue

liepieshov committed 238 days ago

218560df

Older