jax
490b6472 - Artificially limit SMEM capacity for tests

Commit

5 days ago

Artificially limit SMEM capacity for tests This removes a failure mode where new tests that pass on 10.x devices would fail on 12.y devices due to different SMEM capacity. Also adjust test cases to fit in 12.x SMEM limits.

References

#34288 - Apply artificial shared memory limit to Pallas/Mosaic GPU tests

#34474 - [ROCm] Added support for GESVDJ on ROCm devices

#34494 - [ROCm] Enabled ROCm devices to default to Jacobi SVD on smaller matrices.

#34561 - [ROCm] Enable ToeplitzSymmetricConstruction and condition number tests

#34567 - [ROCm]Unskip testMultivariateNormalSingularCovariance on ROCm

#34568 - [ROCm] Skip test_tridiagonal_solve on ROCm

#34575 - [ROCm] Enabled the Triton Pallas tests to run for ROCm.

#34577 - [ROCm] Enabled RNN unit test "test_no_workspace_overflow" for ROCm.

#34578 - [ROCm] Unskip supported dtypes for testConvolutionsPreferredElementType on ROCm

#34598 - Implement approx_tanh for ROCm using OCML tanh function

#34600 - Enable scaled_matmul tests on ROCm

#34601 - Enable neural network tests on ROCm

#34602 - Fix and enable Pallas ops tests on ROCm

#34603 - Enable miscellaneous tests on ROCm

#34615 - [ROCm] Skip test_batch_axis_sharding_jvp on ROCm

Author

olupton

Parents

cb24d7fc

jax 490b6472 - Artificially limit SMEM capacity for tests

jax
490b6472 - Artificially limit SMEM capacity for tests