fix(tests): update tolerance levels and PRNGKey usage for improved test stability
#3740
| Job | Run time |
|---|---|
| 4m 24s | |
| 4m 32s | |
| 4m 40s | |
| 12m 33s | |
| 25m 7s | |
| 32m 37s | |
| 35m 28s | |
| 24m 51s | |
| 0s | |
| 2h 24m 12s |