ROCm/jax

Pull Requests Commits

Enable rocm devices in psend lowering.

mminutoli committed 113 days ago

75c3f830

Merge pull request #33538 from jakevdp:pull-request-template

Google-ML-Automation committed 113 days ago

3042bd8c

Add pull request template

jakevdp committed 113 days ago

d1c189bf

more guidance on AI-generated contributions

Jake VanderPlas committed 113 days ago

ca2df076

Improve the consistency in how classes are registered with nanobind.

hawkinsp committed 114 days ago

a733e40a

Remove CompileOptions.tuple_arguments, fixing an old TODO.

hawkinsp committed 114 days ago

5fae0466

[Mosaic GPU][NFC] Merge two nearly identical functions in `layout_inference`.

dimitar-asenov committed 114 days ago

aef055bf

Change Tracer so it does not directly inherit from Array.

hawkinsp committed 114 days ago

9d40174f

Automated Code Change

Google-ML-Automation committed 114 days ago

ba820fbe

Automated Code Change

Google-ML-Automation committed 114 days ago

cf2209ee

Automated Code Change

Google-ML-Automation committed 114 days ago

471ae20a

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

Google-ML-Automation committed 114 days ago

c83a7190

Unreduced/reduced can only talk about Explicit and Manual mesh axes in ShapedArray (remove auto axes from the type). This worked in v0.8.0 but fails in v0.8.1

yashk2810 committed 114 days ago

1b3c5a72

Add a test for `out: f32[8@x] = add/mul(x: f32[8@x], y: f32[8]{R:x})` in Manual and Explicit mode. Make sure the backward pass is also as expected.

yashk2810 committed 114 days ago

b08a1052

[doc] add AI contribution guidelines.

Jake VanderPlas committed 114 days ago

886f227e

Propagate unreduced for transpose and DUS correctly in shard_map

yashk2810 committed 114 days ago

1fb2ec4b

[Mosaic GPU][NFC] Rename `EquationSystem` to `ConstraintSystem`.

allanrenucci committed 114 days ago

9504a479

Delete jax._src.dtypes.bit_width from internal API.

justinjfu committed 114 days ago

f1601fb7

Merge pull request #33509 from jax-ml:dependabot/github_actions/actions/checkout-6.0.0

Google-ML-Automation committed 114 days ago

5b5f09ec

Merge pull request #33492 from gnecula:export_memories

Google-ML-Automation committed 114 days ago

156e076c

Bump actions/checkout from 5.0.0 to 6.0.0

dependabot[bot] committed 114 days ago

Verified 6c28d4d9

[Pallas:MGPU] Fix `plgpu.inline_mgpu` support for scalar results.

allanrenucci committed 115 days ago

b3429173

Marked the parameters of safe_zip, safe_map and foreach as positional-only

superbobry committed 115 days ago

0adb584b

[Mosaic GPU] Support inference of additional layouts for `BroadcastInDimOp`.

allanrenucci committed 115 days ago

7b052da3

[Mosaic GPU] Add more info to exception messages in core.py

golechwierowicz committed 115 days ago

4d72ab29

[Mosaic GPU][NFC] Move result inference logic to .cc file.

allanrenucci committed 115 days ago

3d79bc81

Update Jax B200 single gpu tests to use MIG partitions.

IllogicalMoose committed 115 days ago

efdc83d7

[Mosaic GPU][NFC] Make `Equation` a type of `Constraint`.

allanrenucci committed 115 days ago

2805e6f6

[Mosaic TPU] Remove pre-canonicalization optimization pass from tpu.td

apaszke committed 115 days ago

0b0073cb

Tiny usability tweak to allow user to pass a list instead of a tuple, e.g. `pltpu.VMEM([8], jnp.int32)`.

brianwa84 committed 115 days ago

93c57481

Older