murphyk opened this issue 2 years ago · 0 comments
The following tests fail on TPU but pass on CPU, except that I added a hack to skip the tests:
https://github.com/probml/dynamax/blob/main/dynamax/generalized_gaussian_ssm/conditional_moments_gaussian_filter_test.py
https://github.com/probml/dynamax/blob/main/dynamax/nonlinear_gaussian_ssm/extended_kalman_filter_test.py
Interestingly the UKF tests pass.
Please investigate. (Might just need to turn on 64bit precision.)