What's the scope of get_state(), set_state(), and get_parameter()

[mil-ad]

Am I right in thinking that get_state(), set_state(), and get_parameter() are only relevant inside a Haiku Module or the function that creates a Haiku Module?

If yes I think the documentation could be improved a tiny bit to make this clear. For example the way get_state() and set_state() are mentioned in the documentation of transform_with_state() suggests that you can use them with what init() returns.

[tomhennigan]

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 it is easier to think about this as if the API was self.get_state(..) instead of hk.get_state(..)).

Perhaps the best example of get/set state are in hk.ExponentialMovingAverage where (during training) the value of the moving averages are updated as part of the forward pass of the network.

If you are not inside a module method then parameters/state are put into "root" (or "no module") namespace. This is not a super common use case and it is probably easiest to not think about it. We use this in the documentation to try and keep the code snippet brief (e.g. not including a module class, instantiation etc).

FWIW all Haiku documentation is run under doctest so you can assume that the snippets in transform_with_state (and all others) would work if you ran them 😄

[mil-ad]

Aw that makes sense. This whole notion of namespaces are also not mentioned in the docs I think? Do you consider PRs from outside DM? I might have a go at making some bits in the doc a tiny bit clearer.

[tomhennigan]

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 to improve the guides section quite a bit, so if you would prefer to wait (and avoid merge conflicts) on that I'd be happy to bump this issue so you can see if the docs are clear enough.