ROCm/jax

Pull Requests Commits

[pallas:triton] Fix atomic min/max lowering for unsigned integers and float types

draganmladjenovic committed 2 days ago

0f6e6d7f

[Pallas:TPU] Add a test checking that dynamic shape exports require the pallas_export_experimental flag

apaszke committed 2 days ago

f2017f78

[Pallas:TPU] Add tests for DMAs with regular semaphores

apaszke committed 2 days ago

9e11f1c0

[Mosaic GPU] Add support for warp-level f8 MMAs

apaszke committed 2 days ago

7192cf7d

Merge pull request #35412 from levskaya:pyrefly

Google-ML-Automation committed 2 days ago

4826da47

Upgrade Abseil to LTS 20260107.1

akuegel committed 2 days ago

ee4cea11

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

Google-ML-Automation committed 2 days ago

d0d6f426

[Mosaic:GPU] Simplify collective metadata initialization.

PatriosTheGreat committed 2 days ago

6f06ab80

Reverts 6869e28a4bde0f6d3a06a7c39858fe40cc44ff60

Google-ML-Automation committed 2 days ago

efa6bd9f

Merge pull request #35414 from levskaya:pyrefly3

Google-ML-Automation committed 2 days ago

dc3f50ca

[mosaic] Removed previously deprecated tpu.wait_dma

superbobry committed 2 days ago

6869e28a

[typing] pyrefly fixes in jax/_src/scipy/optimize

levskaya committed 2 days ago

3b399a8f

[typing] pyrefly fixes in jax/_src/scipy/stats

levskaya committed 2 days ago

b0d91345

Explicitly check for disallowing TransformedRefs in higher-order JAX primitives which currently don't support them.

rdyro committed 2 days ago

6a635c9a

[windows] Skip triton export tests in pallas_shape_poly_test if triton can't be imported.

danielsuo committed 2 days ago

7f250e0a

Removed more dead code

superbobry committed 2 days ago

480847f8

Merge pull request #35325 from jakevdp:pyrefly-fixes

Google-ML-Automation committed 2 days ago

742ad02a

[typing] more pyrefly fixes in jax/_src/lax/

jakevdp committed 2 days ago

0461a3bb

[Pallas][Mosaic TPU] Add flag to disable semaphore checks per kernel.

bythew3i committed 2 days ago

77805318

Merge pull request #35178 from ezhulenev:jax_sort_devices_by_process_index

Google-ML-Automation committed 2 days ago

015f1cbd

Merge pull request #34601 from ROCm:phambinh/ci_rocm_nn_tests

Google-ML-Automation committed 2 days ago

1f3139e0

Merge pull request #35405 from jakevdp:pyrefly-ignore-statements

Google-ML-Automation committed 2 days ago

cda29303

[typing] add type annotation for `jax.smap`.

Jake VanderPlas committed 2 days ago

b8171e35

[pyrefly] remove ignore statements for issues fixed in v0.54

jakevdp committed 2 days ago

19e71a58

Merge pull request #35015 from IvyZX:scguide

Google-ML-Automation committed 2 days ago

100b05c8

Merge pull request #35404 from jakevdp:pyrefly-054

Google-ML-Automation committed 2 days ago

f5c39fed

[lax_numpy] Skip setxor1 tests temporarily.

danielsuo committed 2 days ago

ac537c76

[windows] Skip ops_test.test_binary_scalar if no TPU.

danielsuo committed 2 days ago

5c0ec3e8

[ci] Update pyrefly pre-commit to v0.54.0

jakevdp committed 2 days ago

a211a318

[Pallas:MGPU] Fix a race in dynamic_scheduling_loop

apaszke committed 2 days ago

84a33c08

Older