Commits ROCm/jax - SemanticDiff

Maintain the requirement for CUDA builds

phambinhfin committed 49 days ago

b8bc2546

Make nvidia version data optional for ROCm builds

phambinhfin committed 55 days ago

33e668f9

[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 55 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 55 days ago

30e56531

Reverts d00382120d29056eb0d2a4df7cbbe47242482f6a

Google-ML-Automation committed 55 days ago

205126ed

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

Google-ML-Automation committed 55 days ago

6b847ef9

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

WindQAQ committed 55 days ago

b9abdcb8

Allow SMEM inputs/outputs in `core_map` kernels.

Google-ML-Automation committed 55 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 55 days ago

d0038212

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

yueshengys committed 55 days ago

89b03212

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

Google-ML-Automation committed 55 days ago

1bbf0457

Fix some places where we missed propagating ref kinds

dougalm committed 55 days ago

074a6c28

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

apaszke committed 56 days ago

f2562d4a

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

apaszke committed 56 days ago

df5d0a77

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

apaszke committed 56 days ago

cc8ccfd7

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

allanrenucci committed 56 days ago

bdc7cf04

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

apaszke committed 56 days ago

0c370ba4

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

bchetioui committed 56 days ago

3f708b05

[Mosaic GPU] Remove an unnecessary system memory barrier

apaszke committed 56 days ago

f92b304d

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

apaszke committed 56 days ago

266e6398

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

apaszke committed 56 days ago

062e6a41

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

Google-ML-Automation committed 56 days ago

4cdb3cea

Merge pull request #33130 from davisyoshida:reshape-pull

Google-ML-Automation committed 56 days ago

c3ac1697

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

Google-ML-Automation committed 56 days ago

75952a84

Handle pl.BoundedSlice in reshape pull rule.

davisyoshida committed 56 days ago

90dbc26b

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

jakevdp committed 56 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 56 days ago

c25e095f

Add a docstring for jax.set_mesh

yashk2810 committed 56 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 56 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 56 days ago

7448556f