jax
d87942f3 - tweak config for custom_vjp3: put in cache keys

Commit

3 days ago

tweak config for custom_vjp3: put in cache keys

References

#34586 - tweak config for custom_vjp3: put in cache keys

#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

mattjj

mattjj

Parents

Loading