ROCm/jax

Pull Requests Commits

Add org-based guards to all CI caller workflows

mminutoli committed 125 days ago

8ec727f8

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

hawkinsp committed 125 days ago

9fdfa4b6

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

levskaya committed 125 days ago

91280e51

[bazel] enable strict dependency checks in several rules

Jake VanderPlas committed 125 days ago

4f909393

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

superbobry committed 125 days ago

549eb643

Fix dot_general transpose empty mesh case

yashk2810 committed 125 days ago

d25d17ca

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

dimitar-asenov committed 125 days ago

d7708f40

[Pallas TPU] Add more MPMD TEC+SCS tests

sharadmv committed 125 days ago

55f32433

Add HTML color support to the pretty printer.

hawkinsp committed 125 days ago

967b8f36

Move concatenation tree logic to concatenate_p lowering rule.

hawkinsp committed 125 days ago

5c40943a

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

Google-ML-Automation committed 125 days ago

efbd24c7

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

yashk2810 committed 125 days ago

837f78e5

Merge pull request #36739 from ROCm:rocm-ci-cont-build-fix

Google-ML-Automation committed 125 days ago

ab8fa291

[JAX] Add CHLO lowering for cumulative reductions on GPU.

chsigg committed 125 days ago

eb89e727

fix: cap _get_large_negative at float32 range to avoid -inf under debug_infs with float64 inputs

ShresthSamyak committed 125 days ago

e935bd11

Merge pull request #37468 from jakevdp:partial-jit

Google-ML-Automation committed 125 days ago

11214297

[docs] use jit decorator factory throughout

jakevdp committed 125 days ago

ee6cce56

[Mosaic GPU] Handle all possible register values of `Constant` in `Divides` constraint.

golechwierowicz committed 125 days ago

f9669ddd

[Mosaic] Fix: Remove trailing semicolon in definition

Google-ML-Automation committed 125 days ago

b6dac3c2

Reverts 8adf5dd5e4a6bc08d26a4f90a5fb682e481f219d

superbobry committed 125 days ago

1b119048

[Mosaic GPU] Add MinorDimByteAligned constraint.

golechwierowicz committed 125 days ago

51f44510

[Mosaic GPU][NFC] Try calling do_convert with more than 2 elements at a time

apaszke committed 125 days ago

0ff3a512

[Pallas] Move debug_check flag to pallas_core to avoid import cycles

apaszke committed 125 days ago

63c68159

[pallas:mosaic] Ensure kernels with truthy `has_side_effects=` are not DCEd by JAX

superbobry committed 125 days ago

8adf5dd5

Merge pull request #37452 from superbobry:maint

Google-ML-Automation committed 126 days ago

017dcf3a

[Pallas:MGPU] Add optional `cluster_axes` parameter to `dynamic_scheduling_loop`.

allanrenucci committed 126 days ago

879be7c2

Upgraded to Pyrefly 0.64.0

superbobry committed 126 days ago

87c67a3c

Fix a quoting issue, causing bazel to receive two flags as a single argument

apaszke committed 126 days ago

0ccce351

[Mosaic GPU][NFC] Have pairwise_convert take the PTX instruction

apaszke committed 126 days ago

05cff3e5

[Pallas:TPU] Add a missing test skip (not all TPUs have SparseCores)

apaszke committed 126 days ago

cebcaf8c

Older