Open miltonllera opened 6 months ago
Hi Milton,
Sounds great. Please, write de PR for this, we will be happy to review and accept it.
Sorry I've been slow on this.
I've submitted a PR, but I ran into some issues (as can be seen by the failed tests). The problem is the version of Haiku in requirements.txt (0.0.10) doesn't support new style rng keys. It is supported in the latest version (0.0.13) as I have tested this in my own machine. Updating the requirements should thus allow the examples/tests to remain the same.
Jax recently started moving towards new style RNG-keys which entails changing calls to
jax.random.PRNGKey
for calls tojax.key
. Additionally, seeds tonumpy.random.RandomState
from new-style keys (as incompute_cvt_centroids
) must be first processed withjax.random.key_data
otherwise we get a type error during conversion.I am happy to write a pull request for this.