ROCm/jax

Pull Requests Commits

Update XLA dependency to use revision

Google-ML-Automation committed 264 days ago

07c1acd5

Merge pull request #30023 from Cjkkkk:fix_packed_layout_invalid_entries

Google-ML-Automation committed 264 days ago

8a1dbef3

Reverts d7badcae5a2ab3009641c1a06e9490c5ba2013ac

Google-ML-Automation committed 264 days ago

c9df14e9

Merge pull request #30062 from jax-ml:auto-dce-dynamic-jaxpr-trace

Google-ML-Automation committed 264 days ago

d7badcae

set auto_dce=False by default, enable in direct linearize

mattjj committed 264 days ago

89802205

Add `scatter*` primitives to `roofline`.

zacmustin committed 264 days ago

7a06933d

Increase mutable_array_test sharding

jakeharmon8 committed 264 days ago

d95c1dad

remove `pjit_p`, leaving only `jax.extend.core.primitives.jit_p`

froystig committed 264 days ago

fc18e57b

Add `is_ref` to `ShapeDtypeStruct` to allow doing AOT with duck types.

yashk2810 committed 264 days ago

c8b65463

[Pallas] Update changelog for new emit_pipeline features.

justinjfu committed 264 days ago

9db65dde

Add sharding to fused_attention_stablehlo_test_gpu

jakeharmon8 committed 264 days ago

59af3280

Ensure all JAX benchmarks have `block_until_ready`.

zacmustin committed 264 days ago

283fea2c

[Pallas TPU] Refactor casting logic to use bits instead of bytes and allow uint upcasts.

justinjfu committed 264 days ago

74d32ce2

[Mosaic GPU][NFC] Enable the gmem argument to `async_copy` to be an `ir.BlockArgument`

dimitar-asenov committed 264 days ago

03af274b

[Mosaic GPU] Add `CustomReturnOp` to serve as the terminator of `CustomPrimitiveOp`.

dimitar-asenov committed 264 days ago

aa24d5e3

Remove a hanging test from the TPU continuous and nightly wheel test.

quoctruong committed 264 days ago

e6a31051

[Mosaic GPU] Get rid of `func.FuncOp`s in `gpu_layout_inference_test`.

bchetioui committed 264 days ago

d110759a

[Pallas:MGPU] Allow specifying layout for TMEM refs

apaszke committed 264 days ago

7eb6720e

Internal CI Change

jakeharmon8 committed 264 days ago

386ef7f3

[Mosaic TPU] Explicitly specify the padding value in vector.transfer_read

apaszke committed 264 days ago

88af892f

[pallas:mosaic] Enabled pytype for Mosaic GPU lowering

superbobry committed 264 days ago

cc5bbb51

[Mosaic GPU][NFC] Rename `equations` to `eqns` in `gpu_layout_inference_test.py`.

bchetioui committed 264 days ago

a69deb89

[Mosaic GPU] Add a rule for `mosaic_gpu.WGMMAOp` in the equational layout inference system.

bchetioui committed 264 days ago

7804ad4c

[Mosaic GPU] Add support for reductions with data partially replicated across warps

apaszke committed 265 days ago

1f6608e3

[Mosaic GPU] Add layout inference rule for `vector.splat` in the equational layout inference.

bchetioui committed 265 days ago

d667dee4

[Mosaic GPU] Only allow extracting a non-replicated layout assignment out of a `LeastReplicatedExpression` very conservatively.

bchetioui committed 265 days ago

5a7325eb

[Mosaic GPU] Increase fuel parameter for peer id recomputation.

justinjfu committed 265 days ago

5c6d5411

[Mosaic] Add support for dynamic filter values

naummo committed 265 days ago

7dda1731

Export `jax.smap` as a public API. `smap` is an ergonomic API to write single axis shard_maps and make nesting of shard maps delightful to express :)

yashk2810 committed 265 days ago

da1088b0

Update XLA dependency to use revision

Google-ML-Automation committed 265 days ago

90e18d5f

Older