objax.variable.VarCollection.update not compliant with key-value assignment

[alvarobartt]

Hi everyone! Thanks for the awesome work with objax and the JAX environment, and happy holidays!

I'm trying to load some VarCollection and/or Dict[str, jnp.DeviceArray] params into the model.vars() which is a VarCollection class, and I can do so by:

for key, value in new_params.items():
    model.vars()[key].assign(value)

But I'd expect objax.variable.VarCollection.update to work the same way e.g.

model.vars().update(new_params)

And the later doesn't work while the first one does, not sure if it's because that's not the intended behavior for VarCollection.update or if I'm doing anything wrong... But just the first one works, which for the moment is fine for what I need, but wanted to mention this just in case there's something not working as expected.

[alvarobartt]

If the intended behavior of objax.variable.VarCollection.update is the one mentioned above, I'll be glad to work on a fix for it! 🤗 Also same for #253 which is another issue I recently opened!