ROCm/jax

Pull Requests Commits

ci: default ROCm pytest to 4 workers/GPU on DPX k8s

Kunreddy committed 45 days ago

d4d7090c

[ROCm] nightly: move pytest legs to the DPX k8s pools

magaonka-amd committed 45 days ago

a9c7f874

[ROCm] Run 2 pytest-xdist workers per GPU instead of 4

magaonka-amd committed 49 days ago

44256381

[ROCm] nightly: run only 4- and 8-GPU pytest on DPX pools

magaonka-amd committed 49 days ago

ad77fd3f

[ROCm] Run export_harnesses_multi_platform_test_gpu downstream

magaonka-amd committed 49 days ago

8c689b7b

Bump erf float16 ULP bound on GPU to 1.0.

hawkinsp committed 49 days ago

c54c660a

Merge pull request #40787 from vfdev-5:dump-smem-add-profiler

Google-ML-Automation committed 49 days ago

a061f067

Dynamically derive TPU process count and bounds in multiprocess tests

bhatuzdaname committed 49 days ago

b1f633f6

[Mosaic GPU] Allow FragmentedArray.splat on unreduced TiledLayout for the neutral element.

allanrenucci committed 49 days ago

7cef06fe

[mosaic_gpu] Added `allow_negative` option to `parse_indices`

superbobry committed 49 days ago

5a1d091b

Merge pull request #40489 from benSepanski:composite-dim-var-args

Google-ML-Automation committed 49 days ago

8126ffba

Merge pull request #40354 from Intel-tensorflow:kranipa/oneapi_skip_tests

Google-ML-Automation committed 49 days ago

f5f9d645

Merge pull request #40347 from Intel-tensorflow:kranipa/oneapi_scaled_matmul_lowering

Google-ML-Automation committed 49 days ago

d0f64a3f

Merge pull request #40353 from Intel-tensorflow:kranipa/oneapi_bazel_gpu_test_support

Google-ML-Automation committed 49 days ago

246a66bc

Merge pull request #40348 from Intel-tensorflow:kranipa/oneapi_ffi_tests

Google-ML-Automation committed 49 days ago

03dc7c6e

[Jax][Triton] Cache triton kernels in the instantiate FFI state to reduce overhead.

dimitar-asenov committed 49 days ago

9bd9f922

[Pallas] Enable pltpu.bitcast for SparseCore

tlongeri committed 49 days ago

3bea8093

[mosaic_gpu] `parse_indices` no longer does wrap-around for negative values

superbobry committed 49 days ago

f48d3cdf

[Jax][Triton] Merge the initialize FFI call into the execute one to reduce launch overhead.

dimitar-asenov committed 49 days ago

aa1fe391

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

Google-ML-Automation committed 50 days ago

e6a5b8c2

Merge pull request #40937 from mattjj:push-krtsommymyoq

Google-ML-Automation committed 50 days ago

75e6d85d

[remat3] reduce_precision at producers, short-circuit in save_from_both_policies, and zero-cotangent pruning

mattjj committed 50 days ago

6ca3164f

Dynamically derive TPU process count and bounds in multiprocess tests

Google-ML-Automation committed 50 days ago

5e19c830

cache Traced.closure_convert's fun so with_consts_as_arg hits the tracing cache

mattjj committed 50 days ago

41de9102

[Pallas:SC][NFC] Make expand loads explicit

naummo committed 50 days ago

cf7525f9

Merge pull request #40895 from mattjj:push-ssomtmvzwxsm

Google-ML-Automation committed 50 days ago

f17c5d3c

Merge pull request #40922 from mattjj:push-lkzysxyyktps

Google-ML-Automation committed 50 days ago

8876b255

Merge pull request #40926 from mattjj:push-kkvluuzwqsmu

Google-ML-Automation committed 50 days ago

a62f4635

Merge pull request #40921 from mattjj:push-mkrkxnrvryvy

Google-ML-Automation committed 50 days ago

9bd66afb

Merge pull request #40923 from mattjj:push-uluswnzrrztn

Google-ML-Automation committed 50 days ago

826d1808

Older