Neil Girdhar

I'm working with @shoyer's pull request, and entering the context manager, but I'm getting ``` ValueError: hk.while_loop does not support initialization (since we cannot statically determine if your loop will...

(I'm sick in bed, so apologies if this reply doesn't make sense.) > IIRC cond has a requirement that the output structure of each branch must be the same, so...

Since I'm blocked, I'm going to start implementing `hk.vjp`. I've been thinking about the design of this, and wanted to see if anyone had time to provide feedback before I...

I've implemented it [here](https://github.com/deepmind/dm-haiku/compare/master...NeilGirdhar:various#diff-d9aa1c2f511d653fc292d5c779d49796ff2ea087ac471c72cf74937c25196d12R694), but it's not working for me yet.

Yup, I agree, and that makes sense. I wanted to update you on my usage of Haiku going forward and maybe get your thoughts. When I first decided to start...

This is a Liskov error because method overrides should accept a superset of whatever their overridden methods accept. Otherwise, the child class cannot be used polymorphically.

> This seems like quite a lot of duplication, and feels a bit odd since the actual attributes aren't defined in any of those classes I agree. Ideally RTD should...

Thank you so much for the lightning fast reply! The first thing I tried was profiling, but I just couldn't make sense of all the information. I even tried using...

> Well, I wouldn't suggest using scan, because I'm pretty sure XLA will do that optimization for you when you make the trip count static (as you did in the...

I made your two suggested changes (below), and the runtime went from 0.69s down to 0.13s: - I replaced `jax.while` with `jax.scan`, and - I call normal once to generate...