JAX Release?

[noelc-s]

Thank you for your contributions, I am very interested in this work. In your paper it is listed that there is also a JAX implementation, and I was just wondering if there are plans to release that any time in the near future.

[arunleob]

Apologies for the delayed response, we do have a partial JAX implementation but its not yet complete and not performant. We plan to finish it and release it sometime in January or February. I'll update this issue when it is released.