ROCm/jax

Pull Requests Commits

Add Bazel TPU tests in Continuous workflow.

Google-ML-Automation committed 170 days ago

fb2a74cb

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

Google-ML-Automation committed 170 days ago

244192bf

Add experimental support for transferring into an existing array in an unsafe

pschuh committed 170 days ago

da357bbc

Fix frame limit feature which was omitted from the C++ port of the traceback logic.

hawkinsp committed 170 days ago

5df6868c

[pallas:mgpu] Support for integer power lowering to enable gelu.

cperivol committed 170 days ago

d16cbac6

Update tsan and asan workflows to use the new linux-x86-n4-64 runners

quoctruong committed 170 days ago

0782d4e9

Fixes flakiness in ragged_dot export test.

Google-ML-Automation committed 170 days ago

28c6fa0f

Allow deserialization with a dtype and Format (layout) via a ShapeDtypeStruct spec

rdyro committed 170 days ago

8514d7b4

Explicitly deprecate jax_default_dtype_bits config.

jakevdp committed 170 days ago

262e1b37

Switch automation to update XLA commit SHA

nitins17 committed 170 days ago

ee3b6533

Merge pull request #30516 from Cjkkkk:fix_cudnn_sdpa_shardy_for_bnth_layout

Google-ML-Automation committed 170 days ago

8bdce731

Merge pull request #30949 from mattjj:vjp3-test

Google-ML-Automation committed 170 days ago

baf9d31d

Merge pull request #30956 from MichaelHudgins:mosaic-blocking

Google-ML-Automation committed 170 days ago

04779571

[Pallas][Mosaic GPU] Allow wait_smem_to_gmem in a warp-level context.

justinjfu committed 170 days ago

15b86f8d

[Pallas][Mosaic GPU] Enable while loops in Warp lowering semantics.

justinjfu committed 170 days ago

7d6a47ee

Fix `multi_broadcast_in_dim` which was not constraining the input shardings correctly when input shape was equal to output shape

yashk2810 committed 170 days ago

90261ab8

Run B200 tests on higher risk mosaic changes

MichaelHudgins committed 170 days ago

5c779f9d

[mutable-arrays] add vjp3 test: ref grads for value primals

mattjj committed 170 days ago

42609a2b

[Pallas] Lower assertions to all kernel types.

sashabu committed 170 days ago

4243d05e

[Mosaic GPU] Simplify verification logic for MMA ops.

allanrenucci committed 170 days ago

b3163c51

[JAX] Remove dtype canonicalization from linalg dtype rules.

hawkinsp committed 170 days ago

26572b28

[pallas] Removed `BlockSpec.indexing_mode`

superbobry committed 170 days ago

81502c51

Make device_put with jax.Memory.Space.Host work outside of jit by running an identity jit.

yashk2810 committed 170 days ago

8f93935e

Fix a reference count leak when collecting tracebacks.

hawkinsp committed 170 days ago

a56a38ba

Add internal helper functions for computing default dtypes.

hawkinsp committed 170 days ago

dbe98ecb

[Pallas/Mosaic GPU] Add documentation about the `MOSAIC_GPU_LLVM_DUMP_ONLY` environment variable to the debugging documentation.

bchetioui committed 170 days ago

32f94a49

[Mosaic] Refine scalar load test for narrow bitwidths

Google-ML-Automation committed 170 days ago

4a045f73

Merge pull request #30885 from superbobry:main

Google-ML-Automation committed 170 days ago

0269e49d

Use plgpu.ndloop for the m_loop in collective matmul

Google-ML-Automation committed 170 days ago

54ea7f91

[mosaic_gpu] Enabled mypy in all Mosaic GPU sources

superbobry committed 170 days ago

ecc158c0

Newer Older