KeyError in conditional state (solved)

[bionicles]

in a sparsely gated mixture of experts where each expert has state (a memory), there is a KeyError if different experts are activated on apply as happened to be activated on init -- to solve this, you can pass an 'init' flag into your custom module and if it is True, then you just use all the experts on that call. If that's a memory issue, you can use them one by one. Just make sure init hits all conditional branches of the modules with state

(let me know if there's an easier solution)