ROCm/jax

Pull Requests Commits

Add a 2D test in memories_test.

Google-ML-Automation committed 1 year ago

7b7d36a8

[Pallas] Fix integer array indexing

ayaka14732 committed 1 year ago

9f5f6edb

Merge pull request #27938 from hawkinsp:scipy

Google-ML-Automation committed 1 year ago

c5d6a199

Bump scipy build requirement on Python 3.13.

hawkinsp committed 1 year ago

ffc33abb

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

Google-ML-Automation committed 1 year ago

907725df

optimize while_loop by moving readonly carry components to be consts

mattjj committed 1 year ago

6e52b1e9

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

tlongeri committed 1 year ago

6d57f00b

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

hawkinsp committed 1 year ago

b352763a

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

nitins17 committed 1 year ago

b73bf1a0

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

cperivol committed 1 year ago

41a8805d

Remove unused jaxlib_mlir_capi targets.

hawkinsp committed 1 year ago

59068ae6

Allow ctrl-c to cancel block_until_ready().

pschuh committed 1 year ago

3864c4f3

Update test shardings.

hawkinsp committed 1 year ago

cf8a5246

[Mosaic GPU] Implement warp-level thread semantics.

justinjfu committed 1 year ago

92be510f

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

pschuh committed 1 year ago

48e14dcc

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

yashk2810 committed 1 year ago

7e5966b1

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

justinjfu committed 1 year ago

2807ae4e

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

justinjfu committed 1 year ago

7117aa03

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

Google-ML-Automation committed 1 year ago

64e10ad9

Enable execution of explicit-sharding notebook in docs.

dfm committed 1 year ago

a940100a

Add documentation for JAX's CI folder

nitins17 committed 1 year ago

edc76c7a

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

dougalm committed 1 year ago

5f5e7423

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

Google-ML-Automation committed 1 year ago

8482b7f6

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

Google-ML-Automation committed 1 year ago

349605c2

Run notebooks as part of docs presubmit.

dfm committed 1 year ago

9f7507f2

Merge pull request #27849 from ZacCranko:docfig

Google-ML-Automation committed 1 year ago

16ffbca5

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

dfm committed 1 year ago

dc33db3f

Fix dtype failures in JaxGroupedQueryAttentionReferenceTest.

dfm committed 1 year ago

f3115d32

Merge pull request #27903 from mattjj:pvary-errors

Google-ML-Automation committed 1 year 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 1 year ago

6c0ac7a5

Older