the training_kwargs in VariationalStates are unused

[vigsterkr]

There's a dictionary parameter for MCState called training_kwargs that based on my understanding is for setting the variables that are going to be passed when training a model. this parameter is being saved as a frozen dictionary in MCState but i don't see it being used anywhere in the code-base.

this parameter is to be deprecated or just the feature was never finished? as i'd like to use a dropout layer in my model and would like to be able to pass an is_training boolean parameter when the model is being trained.

[PhilipVinc]

just the feature was never finished?

That one.
I was planning to use it as a starting point to implement things like batch norm, but never got around to finish it.
Would love for someone to actually implement it at some point.

You need to store this dictionary in the constructor.
Then pass it to the jitted kernel in the body of vstate.expect_and_grad/forces in here as a static argument

[vigsterkr]

cool! sure thing i’ll give it a go and will tag you in the PROn 30 May 2023, at 19:36, Filippo Vicentini ***@***.***> wrote: just the feature was never finished? That one. I was planning to use it as a starting point to implement things like batch norm, but never got around to finish it. Would love for someone to actually implement it at some point. You need to store this dictionary in the constructor. Then pass it to the jitted kernel in the body of vstate.expect_and_grad/forces in here as a static argument —Reply to this email directly, view it on GitHub, or unsubscribe.You are receiving this because you authored the thread.Message ID: ***@***.***>

[vigsterkr]

@PhilipVinc a bit unrelated but on the same topic: when using haiku as for model fw, if you use a layer that requires a PRNG (say dropout layer), it'll blow up on hk.next_rng_key() when calling apply() because the prng is not being passed. do you have any quick-fix ideas for that?