ROCm/jax

Pull Requests Commits

Merge pull request #29794 from mattjj:remove-attrs

Google-ML-Automation committed 182 days ago

e179861e

[Mosaic GPU] Make the Pallas Blackwell matmul kernel persistent.

Rifur13 committed 182 days ago

0a483640

[pallas] Relax check for block spec equality in fuser.

chr1sj0nes committed 182 days ago

fc74a0aa

Move jax._src.ops to jax._src.numpy build rule

Jake VanderPlas committed 182 days ago

a75061dd

Update XLA dependency to use revision

Google-ML-Automation committed 183 days ago

1394a92d

Merge pull request #29797 from mattjj:mutable-array-jit-ad

Google-ML-Automation committed 183 days ago

62b8ec5c

[attrs] remove attrs and boxes v1 (attrs_tracked etc)

mattjj committed 183 days ago

8b526c54

[mutable-arrays] fix mutable array AD bug with HOPs

mattjj committed 183 days ago

14194f84

[ra2a] improve ragged_all_to_all batching rule to perform one call

mattjj committed 183 days ago

d1e15f61

Add a comment to make the current McJAX-E code a bit clearer

yashk2810 committed 183 days ago

77a2b616

[Mosaic] Move gather/scatter differentiation to a separate method

naummo committed 183 days ago

a6ca75d7

add psend/precv test for 2 GPUs and re-enable tests for 8 GPUs

Google-ML-Automation committed 183 days ago

0b996b7f

[Pallas][Mosaic GPU] Expose tcgen05.commit_arrive as a standalone primitive.

justinjfu committed 183 days ago

978b76b3

Reverts 4d142d89e8e260266258ebf55e15620a4eeea8d4

WindQAQ committed 183 days ago

ebf7d649

Merge pull request #29786 from Cjkkkk:disable_nn_sdpa_test_on_gb300

Google-ML-Automation committed 183 days ago

baa9ffba

disable gb300 runs until cudnn supports it properly

Cjkkkk committed 183 days ago

6b81ef5f

Add _XlaShardingV2 to tf.XlaShardOp and use it for tf2xla lowering.

bixia1 committed 183 days ago

169a123e

Update the next release version to 0.7.0

yashk2810 committed 183 days ago

0b0a651c

Prepare TPU interpret mode for AbstractRefs having a memory_space attr.

jburnim committed 183 days ago

b1f4ba9d

Make the `device_assignment` argument of `.compile` kwarg only

yashk2810 committed 183 days ago

f0f16ce0

[Mosaic] Support bf16 <-> f8 casting.

WindQAQ committed 183 days ago

4d142d89

Remove donated_invars from PjitParams as a field since it already exists in PjitParams

yashk2810 committed 183 days ago

f392d680

Remove name_stack argument from lower_jaxpr_to_module.

hawkinsp committed 183 days ago

5f805ce2

Merge pull request #29738 from gnecula:composite_consts

Google-ML-Automation committed 183 days ago

59904962

Integrate Triton up to [0a4aa69](https://github.com/openai/triton/commits/0a4aa6960599b17290eb942d71d14afc1775b175)

chsigg committed 183 days ago

15fb17cc

Add modes for `config.jax_dump_ir_to` via comma-delimited `config.jax_dump_ir_modes`.

danielsuo committed 183 days ago

18259e7c

[Mosaic GPU] Add a general broadcast_in_dim method that works for any TiledLayout

apaszke committed 183 days ago

172d44b8

[Mosaic:TPU] 16-bit tpu.dynamic_gather

tlongeri committed 183 days ago

6a6ceee6

Update XLA dependency to use revision

Google-ML-Automation committed 184 days ago

36c4581f

[Mosaic:TPU] tpu.dynamic_gather for bigger shapes, and remove wrapping property

tlongeri committed 184 days ago

a2f28259

Older