Tom Hennigan

FWIW I've spotted the same thing recently, still not sure what is allocating but I don't think this allocation is coming via the CUDA runtime or driver APIs (i've tried...

Hi @NeilGirdhar and @trevorcai , I have some reservations about lifting this restriction. I think we should try as much as possible to avoid Haiku "leaking" into more of your...

I really like the idea of enabling libraries to override/monkey patch jit et al and I think I've discussed this before with @shoyer and @jekbradbury although I can't find the...

Ah here is the issue google/jax#4117.

I think that sounds good, It looks like the PR mostly needs rebasing and someone to merge it so I don't think you'll need to wait for too long. Wrt...

Sorry for the long delay, it's been a busy few weeks and I've not had the headspace to dig into this. > First off, what about parameters that are initialized...

Hi, in this case you can avoid using `hk.lift` if you would prefer. We're missing documentation here (I will fix shortly, an internal user pointed this out earlier today :smile:)...

Compile time for a program of that size should only be 2-3 seconds, and this is a cost you should only pay once per process (unless you are frequently changing...

Hi @mil-ad , thank you for the feedback! `*_state` mirror `get_parameter` in that the state is associated with the module you are currently inside (I have heard from others that...

We sure do 😄 As a heads up, a few of us have actually been iterating on the docs a bit recently (e.g. ae0d060, 81d1b4b) and we have some plans...