Commits ROCm/jax - SemanticDiff

Fix ROCm wheel directory naming for editable builds

phambinhfin committed 38 days ago

e1b42a8a

Fix ROCm lib names copy pattern to use wildcard version

phambinhfin committed 44 days ago

ccb51843

[Take 2] Deprecate default Auto AxisType created by make_mesh. From JAX v0.9.0, make_mesh will return `Explicit` AxisType as the default.

yashk2810 committed 44 days ago

08fcdb72

The default AxisType returned by `make_mesh` is going to change to `AxisType.Explicit` in a follow up CL. This change makes it so that the current code falls back to the old behavior.

yashk2810 committed 44 days ago

30e56531

Reverts d00382120d29056eb0d2a4df7cbbe47242482f6a

Google-ML-Automation committed 44 days ago

205126ed

Increase max_concurrent_steps for transposed ragged dot MGPU kernel after emit_pipeline bug fix

Google-ML-Automation committed 44 days ago

6b847ef9

[Mosaic] Upcast lhs to bf16 for f8 @ bf16/f32 on TPU 7x.

WindQAQ committed 44 days ago

b9abdcb8

Allow SMEM inputs/outputs in `core_map` kernels.

Google-ML-Automation committed 44 days ago

67928e13

Deprecate default Auto AxisType created by make_mesh. From JAX v0.9.0, make_mesh will return `Explicit` AxisType as the default.

yashk2810 committed 45 days ago

d0038212

[Mosaic TPU] Fix Vmem OOM message test by reducing the block size by half for the double-buffered case.

yueshengys committed 45 days ago

89b03212

Merge pull request #33160 from jax-ml:more-ref-kinds

Google-ML-Automation committed 45 days ago

1bbf0457

Fix some places where we missed propagating ref kinds

dougalm committed 45 days ago

074a6c28

[Mosaic TPU][NFC] Move the load->reshape optimization to the pre canonicalization optimization pass

apaszke committed 45 days ago

f2562d4a

[Mosaic TPU] Disable pipelining when the kernel always loads the same oeprand window

apaszke committed 45 days ago

df5d0a77

[Mosaic GPU] Add a test for remote TMA store-add

apaszke committed 45 days ago

cc8ccfd7

[Mosaic GPU] Enable bidirectional transpose relayouts in `FragmentedArray.to_layout`.

allanrenucci committed 45 days ago

bdc7cf04

[Mosaic TPU] Use the newly added tpu.pack_elementwise to speed up the reshape/store optimization

apaszke committed 45 days ago

0c370ba4

[Pallas/Mosaic GPU] Fix `emit_pipeline` to handle dynamic grids containing less than `max_concurrent_steps` iterations.

bchetioui committed 45 days ago

3f708b05

[Mosaic GPU] Remove an unnecessary system memory barrier

apaszke committed 45 days ago

f92b304d

[Mosaic TPU] Add support for slicing in reshape->store fusion optimization

apaszke committed 45 days ago

266e6398

[Mosaic TPU] Refactor the reshape->store optimization to follow the load->reshape implementation

apaszke committed 45 days ago

062e6a41

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

Google-ML-Automation committed 45 days ago

4cdb3cea

Merge pull request #33130 from davisyoshida:reshape-pull

Google-ML-Automation committed 45 days ago

c3ac1697

Merge pull request #33145 from jakevdp:fix-ref-annotations

Google-ML-Automation committed 45 days ago

75952a84

Handle pl.BoundedSlice in reshape pull rule.

davisyoshida committed 45 days ago

90dbc26b

[mutable arrays] fix type annotations for jax.ref APIs

jakevdp committed 45 days ago

7fa65c07

Remove usage of `.tar` files in CI jobs due to fast extraction of `tar.xz` files done by hermetic `xz` tool.

ybaturina committed 45 days ago

c25e095f

Add a docstring for jax.set_mesh

yashk2810 committed 45 days ago

0dfd3552

Use dtype=int instead of explicit dtype=jnp.int64 in transposed ragged dot MGPU kernel, so it works for both x64 and x32 configs.

Google-ML-Automation committed 45 days ago

7c6c2bb6

Allow reduced on the fwd pass to jnp.sin even though unreduced is not allowed as an input to sin since sin is a non-linear op.

yashk2810 committed 45 days ago

7448556f