search

google / objax

Apache License 2.0
768 stars 77 forks source link

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

Open alvarobartt opened 1 year ago

alvarobartt commented 1 year ago

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 commented 1 year ago

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 https://github.com/google/objax/issues/253 which is another issue I recently opened!