ROCm/jax

Pull Requests Commits

[pallas] Removed a few previously deprecated APIs

superbobry committed 11 days ago

0c9d6eba

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

Google-ML-Automation committed 12 days ago

7d859294

Integrate StableHLO at openxla/stablehlo@63933593

Srishti Srivastava committed 12 days ago

d0e43456

Make jit non-strict (modulo other HOPs) for now until we decide what to do with all eqns.

yashk2810 committed 12 days ago

2d323591

Update `remat`'s signature to accept `bool | Sequence[bool]` to match `jax.checkpoint`.

jbspooner committed 12 days ago

79be479f

Merge pull request #40406 from mattjj:push-zkxnywuwkxpl

Google-ML-Automation committed 12 days ago

e29c021e

[Mosaic GPU] Add test coverage for integer splat product reduction

allanrenucci committed 12 days ago

daef08a1

Merge pull request #40404 from mattjj:push-upykpxokuuzt

Google-ML-Automation committed 12 days ago

93dca8e2

[remat3] misc fixes, plus checkpoint_name_fwd

mattjj committed 12 days ago

ca80d466

[docs] phrasing tweaks throughout

mattjj committed 12 days ago

1da28a0c

Guard PyArray_Storage::py_arrays with mu.

hawkinsp committed 12 days ago

30fa0a09

Bump XLA commit.

hawkinsp committed 12 days ago

dfaa9aec

Guard PyArray_Storage::ifrt_array with mu.

hawkinsp committed 12 days ago

3092863c

Merge pull request #40389 from ROCm:fix/rocm-multidim-real-fft-twiddle-cache

Google-ML-Automation committed 12 days ago

7d3d882c

Merge pull request #40370 from ROCm:user/mminutoli/drop_seccomp_unconfined_upstream

Google-ML-Automation committed 12 days ago

3662a88a

Make _arrays a read-only property of Array.

hawkinsp committed 12 days ago

39ec59b4

Don't barrier over literals everywhere in program_order

yashk2810 committed 12 days ago

cef889fc

Make Array.aval immutable.

hawkinsp committed 12 days ago

cd81d48f

[CI - Windows] Configure Bazel's JVM DNS to better survive short-lived lookup flakes.

belitskiy committed 12 days ago

620b2c55

Remove now-unused PyArray::ifrt_array() method.

hawkinsp committed 12 days ago

0636f0f2

[mgpu] Add `collective_axes` parameter to `try_cluster_cancel`.

chr1sj0nes committed 12 days ago

6f42033c

[Mosaic GPU] Fix missing `orders_tensor_cores=True` in test.

dimitar-asenov committed 12 days ago

f28bee43

Bump pinned XLA commit to get the LLVM integration patch fix.

belitskiy committed 12 days ago

bbb3fc3a

[Mosaic GPU] Replace NVVM ops with inline PTX in `dialect_lowering.py`.

bchetioui committed 12 days ago

1c7b170c

[Mosaic GPU] Add a `ScopedDiagnosticHandler` to propagate deserialization failures.

bchetioui committed 12 days ago

2334ef6d

[Mosaic GPU] Replace NVVM ops with inline PTX in `core.py`.

bchetioui committed 12 days ago

9ba3b04d

Merge pull request #40279 from saladpalad:mosaic-gpu-multidim-collective-commit

Google-ML-Automation committed 12 days ago

d0e27424

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

Google-ML-Automation committed 13 days ago

8126382e

[Mosaic GPU] Clean up `dialect_lowering`'s `LoweringContext` to always hold a `LaunchContext`.

bchetioui committed 13 days ago

2f838644

[ROCm] Work around a rocFFT twiddle cache bug in multi-dimensional real FFTs

magaonka-amd committed 13 days ago

9034a8a2

Older