Add MutableParams and MutableState return types for init_fns.
Haiku init_fns always return regular python dicts, so let's be explicit about
that! This means downstream use of these mappings can use eg. setitem
without the type checker complaining. We accept any Mapping for Params and State
inputs, but turn them into a Dict[str, Dict[...]] internally.
Add MutableParams and MutableState return types for init_fns.
Haiku init_fns always return regular python dicts, so let's be explicit about that! This means downstream use of these mappings can use eg. setitem without the type checker complaining. We accept any Mapping for Params and State inputs, but turn them into a Dict[str, Dict[...]] internally.