Nicholas Vadivelu
Nicholas Vadivelu
> Hidden state is notoriously hard to reason about and I think all ML frameworks are struggling with it currently. See for example the widely made mistake of mixing the...
I think `flatten_dict` and `unflatten_dict` are helpful utilities, perhaps they can be moved to `jax_utils`?
I'd be happy to work on this, but what's the process to actually deprecate then remove this module? Do I just add deprecation warnings to all the functions?
Got it, will do next week!
Thanks for raising this @cgarciae, definitely a relevant use case. I would prefer having a `use_running_average` member variable in `ConvBlock`. Perhaps in the future we can add a `use_running_average=None` argument...
Thanks for the PR! I'm going to leave this PR open for now until the library has a few more users/iterations before merging it. It looks great for now, and...
Cool library, but I think it is a bit too new right now. Would love to add this in once there are more examples, users, etc. Thanks!
Thanks for putting this together @davisyoshida! Porting the weights is indeed very helpful. Out of curiousity, have you tried contributing this directly to Haiku (just the model, not the weights)?...
Haiku has excellent model summary functionality (documentation [here](https://dm-haiku.readthedocs.io/en/latest/api.html#haiku.experimental.tabulate)). I think equivalents to `haiku.experimental.tabulate` and `haiku.experimental.eval_summary` would be very helpful.
Nice library! I will leave this PR open and wait until data-snax accrues some users and stability before adding it to the list.