funsor icon indicating copy to clipboard operation
funsor copied to clipboard

Published 20 hours ago • verified publisher icon pyro-ppl

Improve the speed of tests under JAX backend

Open fehiepsi opened this issue 4 years ago • 2 comments

Currently, jax tests are pretty slow comparing to torch backend. It would be nice if we can improve the speed of those tests.

fehiepsi avatar Nov 07 '20 00:11 fehiepsi

Two places in the tests where this is most prominent: test_distribution_generic.py and test_sum_product.py.

eb8680 avatar Nov 28 '20 23:11 eb8680

I tried various ways in JAX to improve compiling time for samplers but no success. :(

fehiepsi avatar Nov 29 '20 04:11 fehiepsi