sync with working jax inference code from main repo e0cb68e Contrebande Labs commited on Apr 28, 2023