ROCm/jax

Pull Requests Commits

Move ROCm GPU detection to hardware_utils and add multi-GPU shm warning

alekstheod committed 11 days ago

49e73341

[Mosaic:GPU] Enable test_semaphore_signal_collective_axes for multi-device per process.

PatriosTheGreat committed 11 days ago

e66bcf62

[Pallas][GPU kernel interpreter] Log coordinates of blocks/clusters in the grid.

Google-ML-Automation committed 11 days ago

e52e8655

[Mosaic GPU] Enable warp-level lowering for `memref.CastOp` in WG semantics.

dimitar-asenov committed 11 days ago

09e3806d

[pallas] Changed error messages across Pallas to point to the new issue templates

superbobry committed 11 days ago

24655e6b

[Mosaic GPU] Generate multiple pairwise PTX conversions when possible

apaszke committed 11 days ago

3f60908b

[Mosaic:GPU] Use NCCL API on the host to collective peer addresses.

PatriosTheGreat committed 11 days ago

77c0fd0c

Re-enable the test data upload meta-workflow with a lighter runner.

belitskiy committed 11 days ago

b9728015

[Pallas][GPU kernel interpreter] Clarify type annotations, esp. of Jax/Numpy arrays, of callback arguments.

Google-ML-Automation committed 11 days ago

d7269489

[mgpu] Set attributes on temporary module used for `cond` lowering.

chr1sj0nes committed 11 days ago

6ca3f74a

[Pallas][GPU kernel interpreter] Support simulating multiple GPUs with `shard_map`.

Google-ML-Automation committed 11 days ago

be8712fb

[Mosaic GPU] Improve the warp_specialized_transpose test

dimitar-asenov committed 11 days ago

b0fff486

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

Google-ML-Automation committed 11 days ago

1b32e596

[pallas] Removed support for checkify

superbobry committed 11 days ago

ee809fb8

Expose compute_on2 temporarily

yashk2810 committed 11 days ago

a1a16a2d

[bazel] enable strict dependency checks for several targets

Jake VanderPlas committed 11 days ago

8ed483f3

Merge pull request #37448 from mike-lmctl:preserve-polyval-integer-dtype

Google-ML-Automation committed 11 days ago

cbcaa43f

Add out_sharding to remaining index update ops

yashk2810 committed 11 days ago

ba2fa98b

Preserve integer dtype in polyval evaluation

mike-lmctl committed 11 days ago

6c690e77

Increase the default number of stack frames to include in tracebacks to 100.

hawkinsp committed 11 days ago

9fdfa4b6

Fix pallas fuser BlockIndexTransform comparison to consider symbolic equivalence of block indices transforms rather than just object identity.

levskaya committed 11 days ago

91280e51

[bazel] enable strict dependency checks in several rules

Jake VanderPlas committed 11 days ago

4f909393

[pallas] Never DCE `pallas_call_p` and `mpmd_map_p`

superbobry committed 11 days ago

549eb643

Fix dot_general transpose empty mesh case

yashk2810 committed 11 days ago

d25d17ca

[Mosaic GPU] Enable warp-level lowering for `memref.Transpose` in WG semantics.

dimitar-asenov committed 11 days ago

d7708f40

[Pallas TPU] Add more MPMD TEC+SCS tests

sharadmv committed 11 days ago

55f32433

Add HTML color support to the pretty printer.

hawkinsp committed 11 days ago

967b8f36

Move concatenation tree logic to concatenate_p lowering rule.

hawkinsp committed 11 days ago

5c40943a

Merge pull request #37457 from ShresthSamyak:fix/dot-product-attention-debug-infs-float64

Google-ML-Automation committed 11 days ago

efbd24c7

Set cotangent sharding to full replicated if the corresponding primal sharding is empty in dot transpose.

yashk2810 committed 11 days ago

837f78e5

Older