Questions tagged [dm-haiku]

What are main differences between flax and dm-haiku? From theirs descriptions: Flax, a neural network library for JAX Haiku, a neural network library for JAX inspired by Sonnet Question: Which one jax-based library should I pick to implement,…

I am writing a Markov chain Monte Carlo simulation in JAX which involves a large series of sampling steps. I currently rely on haiku's PRNGSequence to do the pseudo random number generator key bookkeeping: import haiku as hk def step(key,…

I'm trying to use the Jax library with haiku on python3.6 at conda env, I met the below error and am stuck. I have tried to update my Jax version but nothing changed how can I fix it? Traceback (most recent call last): File "train.py", line 14, in…

I'm trying to understand how did Haiku achieve 2x speedup when training ResNet50 on ImageNet https://github.com/deepmind/dm-haiku/tree/main/examples/imagenet using the Deepmind JMP lib https://github.com/deepmind/jmp, and how to replicate this with…