Nicholas Junge

FWIW, this does not reproduce on my macbook, albeit with a (locally built) jaxlib 0.3.18, also at HEAD: ```py >>> import jax.numpy as jnp >>> import jax >>> jax.config.update("jax_enable_x64", True)...

Makes sense, and also matches what I found on SO about this error code. But to me it would seem reasonable not to remove any Bazel tools until the actual...

Bazel is installed on the `windows-latest` runner used by GitHub Actions. I'm not entirely sure how `cibuildwheel` manages to locate it from the Docker container (it does so successfully too...

Ah, so possibly a race condition between the build tool and `pip`? That seems sensible to me. I have not had time to rerun this workflow thoroughly (I spent yesterday...

Turns out that this is not a spurious issue, but rather it made my Windows wheel jobs fail consistently. I suspect that it is not `cibuildwheel`'s fault though, since it...

AFAIK there is no builtin way to load/save JAX models, as it just contains the mathematical machinery for applying transformations and differentiating. For my own uses, I did export some...

Totally. I got some time on the weekend, happy to take a look then. It may or may not be as straightforward for me though, depending on the usability of...

> Thanks for the PR @nicholasjng! > > I think this a good draft but needs some work. Your TODO list is actually quite intriguing and is more in line...

Okay. I noticed that my style transfer example is most likely >400 LOC, which would make it more complex than other existing examples. Should I settle for something easier instead?...

Apologies for the delay. I worked some more, decided to implement a custom model materializer for Haiku (that way, plugging in different jax NN libraries can also be supported by...