ROCm/jax

Pull Requests Commits

Add a 2D test in memories_test.

Google-ML-Automation committed 304 days ago

7b7d36a8

[Pallas] Fix integer array indexing

ayaka14732 committed 305 days ago

9f5f6edb

Merge pull request #27938 from hawkinsp:scipy

Google-ML-Automation committed 305 days ago

c5d6a199

Bump scipy build requirement on Python 3.13.

hawkinsp committed 305 days ago

ffc33abb

Merge pull request #27937 from mattjj:while-readonly-carry-optimization

Google-ML-Automation committed 305 days ago

907725df

optimize while_loop by moving readonly carry components to be consts

mattjj committed 305 days ago

6e52b1e9

[Mosaic:TPU][Relayout] Add implicit 2nd minor

tlongeri committed 305 days ago

6d57f00b

Fix Pallas tests so they work with JAX_TEST_NUM_THREADS >= 1.

hawkinsp committed 305 days ago

b352763a

Update JAX continuous workflow to run once every 3 hours instead of 2.

nitins17 committed 305 days ago

b73bf1a0

[pallas:mgpu] Return types allowed in mgpu.inline_mgpu.

cperivol committed 305 days ago

41a8805d

Remove unused jaxlib_mlir_capi targets.

hawkinsp committed 305 days ago

59068ae6

Allow ctrl-c to cancel block_until_ready().

pschuh committed 305 days ago

3864c4f3

Update test shardings.

hawkinsp committed 305 days ago

cf8a5246

[Mosaic GPU] Implement warp-level thread semantics.

justinjfu committed 305 days ago

92be510f

Implement mutation by replacing the contents of a jax.Array with a result

pschuh committed 305 days ago

48e14dcc

Make sure direct-linearize handles res_names correctly post vma in types being enabled by default

yashk2810 committed 305 days ago

7e5966b1

[Pallas] Fix ()-shaped vectors being materialized in Pallas lowering.

justinjfu committed 305 days ago

2807ae4e

[Mosaic GPU] Skip WGMMA with cluster example on non H100 GPUs.

justinjfu committed 305 days ago

7117aa03

Merge pull request #27924 from dfm:explicit-sharding-tutorial

Google-ML-Automation committed 305 days ago

64e10ad9

Enable execution of explicit-sharding notebook in docs.

dfm committed 305 days ago

a940100a

Add documentation for JAX's CI folder

nitins17 committed 305 days ago

edc76c7a

Mark as thread-unsafe tests that modify possibly-cached jaxprs in-place.

dougalm committed 305 days ago

5f5e7423

Merge pull request #27368 from dfm:docs-on-actions

Google-ML-Automation committed 305 days ago

8482b7f6

Merge pull request #27917 from dfm:rtds-opt-in-label

Google-ML-Automation committed 305 days ago

349605c2

Run notebooks as part of docs presubmit.

dfm committed 305 days ago

9f7507f2

Merge pull request #27849 from ZacCranko:docfig

Google-ML-Automation committed 305 days ago

16ffbca5

Skip Read the Docs builds unless the 'documentation' label is added.

dfm committed 305 days ago

dc33db3f

Fix dtype failures in JaxGroupedQueryAttentionReferenceTest.

dfm committed 305 days ago

f3115d32

Merge pull request #27903 from mattjj:pvary-errors

Google-ML-Automation committed 305 days ago

9011d66a

Do a pvary in dynamic_slice_transpose_rule so that the `zeros` are varying with the correct vma as the operands were.

yashk2810 committed 305 days ago

6c0ac7a5

Older