Mutability of variational state parameters.

[skhaleefah]

Hi all,

In this page here under "Manipulating the parameters" it says that if you do not use flax.core.copy() you will only copy the outermost dictionaries but reference the inner ones and so modifying the inner ones would modify the original parameters as well.

I am on netket version 3.11.2 and flax version 0.8.1 and find this to not be the case. The screenshot demonstrates that there is no difference in referencing the parameters vs. copying them with flax.core.copy() when modifying the inner-most dictionaries. Is this supposed to be the case or am I misinterpreting something here?

Thanks in advance!

[PhilipVinc]

I think the documentation is correct in general. See the following example

In [18]: a = {'a': { 'b':1}}

In [19]: a_copy=a.copy()

In [20]: a['a']['b']=2

In [21]: a_copy
Out[21]: {'a': {'b': 2}}

In [22]: a_flax_copy=flax.core.copy(a, {})

In [23]: a['a']['b']=3

In [24]: a_flax_copy
Out[24]: {'a': {'b': 2}}

However, for the particular case you considered, NetKet is trying hard to ensure you do not modify the parameters within vs.parametersdirectly, and internally returns you a copy obtained with flax.core.copy of the parameters.
See

netket/netket/vqs/base.py

Line 71 in c9c1061

return fcore.copy(self._parameters, {})

This is because the variational state, to work correctly, requires that if you modify the parameters, you reassign them with vs.parameters = new_parameters. This is used in order to ensure that some caches (like the samples) are cleaned when you change parameters.
If you changed them directly by editing the dictionary, we would have no way of knowing, and therefore we could not clear the caches and samples automatically.