Using old minimization state in optimize_kl
When running optimize_kl
again starting from an old OptimizeVIState
by doing something like:
samples, state = jft.optimize_kl(new_like,
new_position,
key=minimization_subkey,
callback=plot,
**minization_config,
_optimize_vi_state=old_state,
)
I get an error because there is no "constants" key in the old OptimizeVIState
.
The trace log is the following:
File "/Users/matteani/PycharmProjects/nifty/src/re/optimize_kl.py", line 686, in optimize_kl
samples, opt_vi_st = opt_vi.update(samples, opt_vi_st)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/matteani/PycharmProjects/nifty/src/re/optimize_kl.py", line 521, in update
constants = _getitem_at_nit(config, "constants", nit)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/Users/matteani/PycharmProjects/nifty/src/re/optimize_kl.py", line 151, in _getitem_at_nit
c = config[key]
~~~~~~^^^^^
KeyError: 'constants'.
Is using the old state inside optimize_kl
not intended?
Or should we add a check in _getitem_at_nit
if the item key is in the config?
@gedenhof
Edited by Matteo Guardiani