We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
AbstractVariationalDriver.py bug in updating step_count
![image](https://user-images.githubusercontent.com/47376937/125520970-8a189955-4de6-45b2-acfc-45eab4c82d9e.png)

I think this error can be fixed either by writing a setter function for the `@property step_count` in the same way that `@optimizer.setter` exists in the same file. 

Otherwise, maybe this can be fixed by changing ` self.step_count = 0` to ` self._step_count = 0`.

![image](https://user-images.githubusercontent.com/47376937/125521480-ca41de24-c283-4d8c-bf16-ba25a93b2399.png)


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `netket/driver/abstract_variational_driver.py`
Content:
```
1 # Copyright 2021 The NetKet Authors - All rights reserved.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #    http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 
15 import abc
16 import numbers
17 from functools import partial
18 
19 from tqdm import tqdm
20 
21 import jax
22 from jax.tree_util import tree_map
23 
24 from netket.logging import JsonLog
25 from netket.utils import mpi
26 
27 
28 def _to_iterable(maybe_iterable):
29     """
30     _to_iterable(maybe_iterable)
31 
32     Ensure the result is iterable. If the input is not iterable, it is wrapped into a tuple.
33     """
34     if hasattr(maybe_iterable, "__iter__"):
35         surely_iterable = maybe_iterable
36     else:
37         surely_iterable = (maybe_iterable,)
38 
39     return surely_iterable
40 
41 
42 # Note: to implement a new Driver (see also _vmc.py for an example)
43 # If you want to inherit the nice interface of AbstractMCDriver, you should
44 # subclass it, defining the following methods:
45 # - Either _forward_and_backward or individually _forward, _backward, that should
46 #   compute the loss function and the gradient. If the driver is minimizing or
47 #   maximising some loss function, this quantity should be assigned to self._stats
48 #   in order to monitor it.
49 # - _estimate_stats should return the MC estimate of a single operator
50 # - reset should reset the driver (usually the sampler).
51 # - info should return a string with an overview of the driver.
52 # - The __init__ method shouldbe called with the machine and the optimizer. If this
53 #   driver is minimising a loss function and you want it's name to show up automatically
54 #   in the progress bar/ouput files you should pass the optional keyword argument
55 #   minimized_quantity_name.
56 class AbstractVariationalDriver(abc.ABC):
57     """Abstract base class for NetKet Variational Monte Carlo drivers"""
58 
59     def __init__(self, variational_state, optimizer, minimized_quantity_name=""):
60         self._mynode = mpi.node_number
61         self._mpi_nodes = mpi.n_nodes
62         self._loss_stats = None
63         self._loss_name = minimized_quantity_name
64         self._step_count = 0
65 
66         self._variational_state = variational_state
67         self.optimizer = optimizer
68 
69     def _forward_and_backward(self):
70         """
71         Performs the forward and backward pass at the same time.
72         Concrete drivers should either override this method, or override individually
73         _forward and _backward.
74 
75         Returns:
76             the update for the weights.
77         """
78         self._forward()
79         dp = self._backward()
80         return dp
81 
82     def _forward(self):
83         """
84         Performs the forward pass, computing the loss function.
85         Concrete should either implement _forward and _backward or the joint method
86         _forward_and_backward.
87         """
88         raise NotImplementedError()
89 
90     def _backward(self):
91         """
92         Performs the backward pass, computing the update for the parameters.
93         Concrete should either implement _forward and _backward or the joint method
94         _forward_and_backward.
95         """
96         raise NotImplementedError()
97 
98     def _estimate_stats(self, observable):
99         """
100         Returns the MCMC statistics for the expectation value of an observable.
101         Must be implemented by super-classes of AbstractVMC.
102 
103         :param observable: A quantum operator (netket observable)
104         :return:
105         """
106         return self.state.expect(observable)
107 
108     def reset(self):
109         """
110         Resets the driver.
111         Concrete drivers should also call super().reset() to ensure that the step
112         count is set to 0.
113         """
114         self.state.reset()
115         self.step_count = 0
116         pass
117 
118     @abc.abstractmethod
119     def info(self, depth=0):
120         """
121         Returns an info string used to print information to screen about this driver.
122         """
123         pass
124 
125     @property
126     def state(self):
127         """
128         Returns the machine that is optimized by this driver.
129         """
130         return self._variational_state
131 
132     @property
133     def optimizer(self):
134         """
135         The optimizer used to update the parameters at every iteration.
136         """
137         return self._optimizer
138 
139     @optimizer.setter
140     def optimizer(self, optimizer):
141         self._optimizer = optimizer
142         self._optimizer_state = optimizer.init(self.state.parameters)
143 
144     @property
145     def step_count(self):
146         """
147         Returns a monotonic integer labelling all the steps performed by this driver.
148         This can be used, for example, to identify the line in a log file.
149         """
150         return self._step_count
151 
152     def iter(self, n_steps: int, step: int = 1):
153         """
154         Returns a generator which advances the VMC optimization, yielding
155         after every `step_size` steps.
156 
157         Args:
158             n_iter: The total number of steps to perform.
159             step_size: The number of internal steps the simulation
160                 is advanced every turn.
161 
162         Yields:
163             int: The current step.
164         """
165         for _ in range(0, n_steps, step):
166             for i in range(0, step):
167                 dp = self._forward_and_backward()
168                 if i == 0:
169                     yield self.step_count
170 
171                 self._step_count += 1
172                 self.update_parameters(dp)
173 
174     def advance(self, steps: int = 1):
175         """
176         Performs `steps` optimization steps.
177 
178         steps: (Default=1) number of steps
179         """
180         for _ in self.iter(steps):
181             pass
182 
183     def run(
184         self,
185         n_iter,
186         out=None,
187         obs=None,
188         show_progress=True,
189         save_params_every=50,  # for default logger
190         write_every=50,  # for default logger
191         step_size=1,  # for default logger
192         callback=lambda *x: True,
193     ):
194         """
195         Executes the Monte Carlo Variational optimization, updating the weights of the network
196         stored in this driver for `n_iter` steps and dumping values of the observables `obs`
197         in the output `logger`. If no logger is specified, creates a json file at `out`,
198         overwriting files with the same prefix.
199 
200         By default uses :ref:`netket.logging.JsonLog`. To know about the output format
201         check it's documentation. The logger object is also returned at the end of this function
202         so that you can inspect the results without reading the json output.
203 
204         Args:
205             n_iter: the total number of iterations
206             out: A logger object, or an iterable of loggers, to be used to store simulation log and data.
207                 If this argument is a string, it will be used as output prefix for the standard JSON logger.
208             obs: An iterable containing all observables that should be computed
209             save_params_every: Every how many steps the parameters of the network should be
210                 serialized to disk (ignored if logger is provided)
211             write_every: Every how many steps the json data should be flushed to disk (ignored if
212                 logger is provided)
213             step_size: Every how many steps should observables be logged to disk (default=1)
214             show_progress: If true displays a progress bar (default=True)
215             callback: Callable or list of callable callback functions to stop training given a condition
216         """
217 
218         if not isinstance(n_iter, numbers.Number):
219             raise ValueError(
220                 "n_iter, the first positional argument to `run`, must be a number!"
221             )
222 
223         if obs is None:
224             obs = {}
225 
226         if out is None:
227             out = tuple()
228             print(
229                 "No output specified (out=[apath|nk.logging.JsonLogger(...)])."
230                 "Running the optimization but not saving the output."
231             )
232 
233         # Log only non-root nodes
234         if self._mynode == 0:
235             # if out is a path, create an overwriting Json Log for output
236             if isinstance(out, str):
237                 loggers = (JsonLog(out, "w", save_params_every, write_every),)
238             else:
239                 loggers = _to_iterable(out)
240         else:
241             loggers = tuple()
242             show_progress = False
243 
244         callbacks = _to_iterable(callback)
245         callback_stop = False
246 
247         with tqdm(total=n_iter, disable=not show_progress) as pbar:
248             old_step = self.step_count
249             first_step = True
250 
251             for step in self.iter(n_iter, step_size):
252 
253                 log_data = self.estimate(obs)
254 
255                 # if the cost-function is defined then report it in the progress bar
256                 if self._loss_stats is not None:
257                     pbar.set_postfix_str(self._loss_name + "=" + str(self._loss_stats))
258                     log_data[self._loss_name] = self._loss_stats
259 
260                 # Execute callbacks before loggers because they can append to log_data
261                 for callback in callbacks:
262                     if not callback(step, log_data, self):
263                         callback_stop = True
264 
265                 for logger in loggers:
266                     logger(self.step_count, log_data, self.state)
267 
268                 if len(callbacks) > 0:
269                     if mpi.mpi_any(callback_stop):
270                         break
271 
272                 # Reset the timing of tqdm after the first step, to ignore compilation time
273                 if first_step:
274                     first_step = False
275                     pbar.unpause()
276 
277                 # Update the progress bar
278                 pbar.update(self.step_count - old_step)
279                 old_step = self.step_count
280 
281             # Final update so that it shows up filled.
282             pbar.update(self.step_count - old_step)
283 
284         # flush at the end of the evolution so that final values are saved to
285         # file
286         for logger in loggers:
287             logger.flush(self.state)
288 
289         return loggers
290 
291     def estimate(self, observables):
292         """
293         Return MCMC statistics for the expectation value of observables in the
294         current state of the driver.
295 
296         Args:
297             observables: A pytree of operators for which statistics should be computed.
298 
299         Returns:
300             A pytree of the same structure as the input, containing MCMC statistics
301             for the corresponding operators as leaves.
302         """
303         return tree_map(self._estimate_stats, observables)
304 
305     def update_parameters(self, dp):
306         """
307         Updates the parameters of the machine using the optimizer in this driver
308 
309         Args:
310             dp: the pytree containing the updates to the parameters
311         """
312         self._optimizer_state, self.state.parameters = apply_gradient(
313             self._optimizer.update, self._optimizer_state, dp, self.state.parameters
314         )
315 
316 
317 @partial(jax.jit, static_argnums=0)
318 def apply_gradient(optimizer_fun, optimizer_state, dp, params):
319     import optax
320 
321     updates, new_optimizer_state = optimizer_fun(dp, optimizer_state, params)
322 
323     new_params = optax.apply_updates(params, updates)
324     return new_optimizer_state, new_params
325
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/netket/driver/abstract_variational_driver.py b/netket/driver/abstract_variational_driver.py
--- a/netket/driver/abstract_variational_driver.py
+++ b/netket/driver/abstract_variational_driver.py
@@ -112,7 +112,7 @@
         count is set to 0.
         """
         self.state.reset()
-        self.step_count = 0
+        self._step_count = 0
         pass
 
     @abc.abstractmethod

{"golden_diff": "diff --git a/netket/driver/abstract_variational_driver.py b/netket/driver/abstract_variational_driver.py\n--- a/netket/driver/abstract_variational_driver.py\n+++ b/netket/driver/abstract_variational_driver.py\n@@ -112,7 +112,7 @@\n         count is set to 0.\n         \"\"\"\n         self.state.reset()\n-        self.step_count = 0\n+        self._step_count = 0\n         pass\n \n     @abc.abstractmethod\n", "issue": "AbstractVariationalDriver.py bug in updating step_count\n![image](https://user-images.githubusercontent.com/47376937/125520970-8a189955-4de6-45b2-acfc-45eab4c82d9e.png)\r\n\r\nI think this error can be fixed either by writing a setter function for the `@property step_count` in the same way that `@optimizer.setter` exists in the same file. \r\n\r\nOtherwise, maybe this can be fixed by changing ` self.step_count = 0` to ` self._step_count = 0`.\r\n\r\n![image](https://user-images.githubusercontent.com/47376937/125521480-ca41de24-c283-4d8c-bf16-ba25a93b2399.png)\r\n\n", "before_files": [{"content": "# Copyright 2021 The NetKet Authors - All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#    http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport abc\nimport numbers\nfrom functools import partial\n\nfrom tqdm import tqdm\n\nimport jax\nfrom jax.tree_util import tree_map\n\nfrom netket.logging import JsonLog\nfrom netket.utils import mpi\n\n\ndef _to_iterable(maybe_iterable):\n    \"\"\"\n    _to_iterable(maybe_iterable)\n\n    Ensure the result is iterable. If the input is not iterable, it is wrapped into a tuple.\n    \"\"\"\n    if hasattr(maybe_iterable, \"__iter__\"):\n        surely_iterable = maybe_iterable\n    else:\n        surely_iterable = (maybe_iterable,)\n\n    return surely_iterable\n\n\n# Note: to implement a new Driver (see also _vmc.py for an example)\n# If you want to inherit the nice interface of AbstractMCDriver, you should\n# subclass it, defining the following methods:\n# - Either _forward_and_backward or individually _forward, _backward, that should\n#   compute the loss function and the gradient. If the driver is minimizing or\n#   maximising some loss function, this quantity should be assigned to self._stats\n#   in order to monitor it.\n# - _estimate_stats should return the MC estimate of a single operator\n# - reset should reset the driver (usually the sampler).\n# - info should return a string with an overview of the driver.\n# - The __init__ method shouldbe called with the machine and the optimizer. If this\n#   driver is minimising a loss function and you want it's name to show up automatically\n#   in the progress bar/ouput files you should pass the optional keyword argument\n#   minimized_quantity_name.\nclass AbstractVariationalDriver(abc.ABC):\n    \"\"\"Abstract base class for NetKet Variational Monte Carlo drivers\"\"\"\n\n    def __init__(self, variational_state, optimizer, minimized_quantity_name=\"\"):\n        self._mynode = mpi.node_number\n        self._mpi_nodes = mpi.n_nodes\n        self._loss_stats = None\n        self._loss_name = minimized_quantity_name\n        self._step_count = 0\n\n        self._variational_state = variational_state\n        self.optimizer = optimizer\n\n    def _forward_and_backward(self):\n        \"\"\"\n        Performs the forward and backward pass at the same time.\n        Concrete drivers should either override this method, or override individually\n        _forward and _backward.\n\n        Returns:\n            the update for the weights.\n        \"\"\"\n        self._forward()\n        dp = self._backward()\n        return dp\n\n    def _forward(self):\n        \"\"\"\n        Performs the forward pass, computing the loss function.\n        Concrete should either implement _forward and _backward or the joint method\n        _forward_and_backward.\n        \"\"\"\n        raise NotImplementedError()\n\n    def _backward(self):\n        \"\"\"\n        Performs the backward pass, computing the update for the parameters.\n        Concrete should either implement _forward and _backward or the joint method\n        _forward_and_backward.\n        \"\"\"\n        raise NotImplementedError()\n\n    def _estimate_stats(self, observable):\n        \"\"\"\n        Returns the MCMC statistics for the expectation value of an observable.\n        Must be implemented by super-classes of AbstractVMC.\n\n        :param observable: A quantum operator (netket observable)\n        :return:\n        \"\"\"\n        return self.state.expect(observable)\n\n    def reset(self):\n        \"\"\"\n        Resets the driver.\n        Concrete drivers should also call super().reset() to ensure that the step\n        count is set to 0.\n        \"\"\"\n        self.state.reset()\n        self.step_count = 0\n        pass\n\n    @abc.abstractmethod\n    def info(self, depth=0):\n        \"\"\"\n        Returns an info string used to print information to screen about this driver.\n        \"\"\"\n        pass\n\n    @property\n    def state(self):\n        \"\"\"\n        Returns the machine that is optimized by this driver.\n        \"\"\"\n        return self._variational_state\n\n    @property\n    def optimizer(self):\n        \"\"\"\n        The optimizer used to update the parameters at every iteration.\n        \"\"\"\n        return self._optimizer\n\n    @optimizer.setter\n    def optimizer(self, optimizer):\n        self._optimizer = optimizer\n        self._optimizer_state = optimizer.init(self.state.parameters)\n\n    @property\n    def step_count(self):\n        \"\"\"\n        Returns a monotonic integer labelling all the steps performed by this driver.\n        This can be used, for example, to identify the line in a log file.\n        \"\"\"\n        return self._step_count\n\n    def iter(self, n_steps: int, step: int = 1):\n        \"\"\"\n        Returns a generator which advances the VMC optimization, yielding\n        after every `step_size` steps.\n\n        Args:\n            n_iter: The total number of steps to perform.\n            step_size: The number of internal steps the simulation\n                is advanced every turn.\n\n        Yields:\n            int: The current step.\n        \"\"\"\n        for _ in range(0, n_steps, step):\n            for i in range(0, step):\n                dp = self._forward_and_backward()\n                if i == 0:\n                    yield self.step_count\n\n                self._step_count += 1\n                self.update_parameters(dp)\n\n    def advance(self, steps: int = 1):\n        \"\"\"\n        Performs `steps` optimization steps.\n\n        steps: (Default=1) number of steps\n        \"\"\"\n        for _ in self.iter(steps):\n            pass\n\n    def run(\n        self,\n        n_iter,\n        out=None,\n        obs=None,\n        show_progress=True,\n        save_params_every=50,  # for default logger\n        write_every=50,  # for default logger\n        step_size=1,  # for default logger\n        callback=lambda *x: True,\n    ):\n        \"\"\"\n        Executes the Monte Carlo Variational optimization, updating the weights of the network\n        stored in this driver for `n_iter` steps and dumping values of the observables `obs`\n        in the output `logger`. If no logger is specified, creates a json file at `out`,\n        overwriting files with the same prefix.\n\n        By default uses :ref:`netket.logging.JsonLog`. To know about the output format\n        check it's documentation. The logger object is also returned at the end of this function\n        so that you can inspect the results without reading the json output.\n\n        Args:\n            n_iter: the total number of iterations\n            out: A logger object, or an iterable of loggers, to be used to store simulation log and data.\n                If this argument is a string, it will be used as output prefix for the standard JSON logger.\n            obs: An iterable containing all observables that should be computed\n            save_params_every: Every how many steps the parameters of the network should be\n                serialized to disk (ignored if logger is provided)\n            write_every: Every how many steps the json data should be flushed to disk (ignored if\n                logger is provided)\n            step_size: Every how many steps should observables be logged to disk (default=1)\n            show_progress: If true displays a progress bar (default=True)\n            callback: Callable or list of callable callback functions to stop training given a condition\n        \"\"\"\n\n        if not isinstance(n_iter, numbers.Number):\n            raise ValueError(\n                \"n_iter, the first positional argument to `run`, must be a number!\"\n            )\n\n        if obs is None:\n            obs = {}\n\n        if out is None:\n            out = tuple()\n            print(\n                \"No output specified (out=[apath|nk.logging.JsonLogger(...)]).\"\n                \"Running the optimization but not saving the output.\"\n            )\n\n        # Log only non-root nodes\n        if self._mynode == 0:\n            # if out is a path, create an overwriting Json Log for output\n            if isinstance(out, str):\n                loggers = (JsonLog(out, \"w\", save_params_every, write_every),)\n            else:\n                loggers = _to_iterable(out)\n        else:\n            loggers = tuple()\n            show_progress = False\n\n        callbacks = _to_iterable(callback)\n        callback_stop = False\n\n        with tqdm(total=n_iter, disable=not show_progress) as pbar:\n            old_step = self.step_count\n            first_step = True\n\n            for step in self.iter(n_iter, step_size):\n\n                log_data = self.estimate(obs)\n\n                # if the cost-function is defined then report it in the progress bar\n                if self._loss_stats is not None:\n                    pbar.set_postfix_str(self._loss_name + \"=\" + str(self._loss_stats))\n                    log_data[self._loss_name] = self._loss_stats\n\n                # Execute callbacks before loggers because they can append to log_data\n                for callback in callbacks:\n                    if not callback(step, log_data, self):\n                        callback_stop = True\n\n                for logger in loggers:\n                    logger(self.step_count, log_data, self.state)\n\n                if len(callbacks) > 0:\n                    if mpi.mpi_any(callback_stop):\n                        break\n\n                # Reset the timing of tqdm after the first step, to ignore compilation time\n                if first_step:\n                    first_step = False\n                    pbar.unpause()\n\n                # Update the progress bar\n                pbar.update(self.step_count - old_step)\n                old_step = self.step_count\n\n            # Final update so that it shows up filled.\n            pbar.update(self.step_count - old_step)\n\n        # flush at the end of the evolution so that final values are saved to\n        # file\n        for logger in loggers:\n            logger.flush(self.state)\n\n        return loggers\n\n    def estimate(self, observables):\n        \"\"\"\n        Return MCMC statistics for the expectation value of observables in the\n        current state of the driver.\n\n        Args:\n            observables: A pytree of operators for which statistics should be computed.\n\n        Returns:\n            A pytree of the same structure as the input, containing MCMC statistics\n            for the corresponding operators as leaves.\n        \"\"\"\n        return tree_map(self._estimate_stats, observables)\n\n    def update_parameters(self, dp):\n        \"\"\"\n        Updates the parameters of the machine using the optimizer in this driver\n\n        Args:\n            dp: the pytree containing the updates to the parameters\n        \"\"\"\n        self._optimizer_state, self.state.parameters = apply_gradient(\n            self._optimizer.update, self._optimizer_state, dp, self.state.parameters\n        )\n\n\n@partial(jax.jit, static_argnums=0)\ndef apply_gradient(optimizer_fun, optimizer_state, dp, params):\n    import optax\n\n    updates, new_optimizer_state = optimizer_fun(dp, optimizer_state, params)\n\n    new_params = optax.apply_updates(params, updates)\n    return new_optimizer_state, new_params\n", "path": "netket/driver/abstract_variational_driver.py"}], "after_files": [{"content": "# Copyright 2021 The NetKet Authors - All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#    http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport abc\nimport numbers\nfrom functools import partial\n\nfrom tqdm import tqdm\n\nimport jax\nfrom jax.tree_util import tree_map\n\nfrom netket.logging import JsonLog\nfrom netket.utils import mpi\n\n\ndef _to_iterable(maybe_iterable):\n    \"\"\"\n    _to_iterable(maybe_iterable)\n\n    Ensure the result is iterable. If the input is not iterable, it is wrapped into a tuple.\n    \"\"\"\n    if hasattr(maybe_iterable, \"__iter__\"):\n        surely_iterable = maybe_iterable\n    else:\n        surely_iterable = (maybe_iterable,)\n\n    return surely_iterable\n\n\n# Note: to implement a new Driver (see also _vmc.py for an example)\n# If you want to inherit the nice interface of AbstractMCDriver, you should\n# subclass it, defining the following methods:\n# - Either _forward_and_backward or individually _forward, _backward, that should\n#   compute the loss function and the gradient. If the driver is minimizing or\n#   maximising some loss function, this quantity should be assigned to self._stats\n#   in order to monitor it.\n# - _estimate_stats should return the MC estimate of a single operator\n# - reset should reset the driver (usually the sampler).\n# - info should return a string with an overview of the driver.\n# - The __init__ method shouldbe called with the machine and the optimizer. If this\n#   driver is minimising a loss function and you want it's name to show up automatically\n#   in the progress bar/ouput files you should pass the optional keyword argument\n#   minimized_quantity_name.\nclass AbstractVariationalDriver(abc.ABC):\n    \"\"\"Abstract base class for NetKet Variational Monte Carlo drivers\"\"\"\n\n    def __init__(self, variational_state, optimizer, minimized_quantity_name=\"\"):\n        self._mynode = mpi.node_number\n        self._mpi_nodes = mpi.n_nodes\n        self._loss_stats = None\n        self._loss_name = minimized_quantity_name\n        self._step_count = 0\n\n        self._variational_state = variational_state\n        self.optimizer = optimizer\n\n    def _forward_and_backward(self):\n        \"\"\"\n        Performs the forward and backward pass at the same time.\n        Concrete drivers should either override this method, or override individually\n        _forward and _backward.\n\n        Returns:\n            the update for the weights.\n        \"\"\"\n        self._forward()\n        dp = self._backward()\n        return dp\n\n    def _forward(self):\n        \"\"\"\n        Performs the forward pass, computing the loss function.\n        Concrete should either implement _forward and _backward or the joint method\n        _forward_and_backward.\n        \"\"\"\n        raise NotImplementedError()\n\n    def _backward(self):\n        \"\"\"\n        Performs the backward pass, computing the update for the parameters.\n        Concrete should either implement _forward and _backward or the joint method\n        _forward_and_backward.\n        \"\"\"\n        raise NotImplementedError()\n\n    def _estimate_stats(self, observable):\n        \"\"\"\n        Returns the MCMC statistics for the expectation value of an observable.\n        Must be implemented by super-classes of AbstractVMC.\n\n        :param observable: A quantum operator (netket observable)\n        :return:\n        \"\"\"\n        return self.state.expect(observable)\n\n    def reset(self):\n        \"\"\"\n        Resets the driver.\n        Concrete drivers should also call super().reset() to ensure that the step\n        count is set to 0.\n        \"\"\"\n        self.state.reset()\n        self._step_count = 0\n        pass\n\n    @abc.abstractmethod\n    def info(self, depth=0):\n        \"\"\"\n        Returns an info string used to print information to screen about this driver.\n        \"\"\"\n        pass\n\n    @property\n    def state(self):\n        \"\"\"\n        Returns the machine that is optimized by this driver.\n        \"\"\"\n        return self._variational_state\n\n    @property\n    def optimizer(self):\n        \"\"\"\n        The optimizer used to update the parameters at every iteration.\n        \"\"\"\n        return self._optimizer\n\n    @optimizer.setter\n    def optimizer(self, optimizer):\n        self._optimizer = optimizer\n        self._optimizer_state = optimizer.init(self.state.parameters)\n\n    @property\n    def step_count(self):\n        \"\"\"\n        Returns a monotonic integer labelling all the steps performed by this driver.\n        This can be used, for example, to identify the line in a log file.\n        \"\"\"\n        return self._step_count\n\n    def iter(self, n_steps: int, step: int = 1):\n        \"\"\"\n        Returns a generator which advances the VMC optimization, yielding\n        after every `step_size` steps.\n\n        Args:\n            n_iter: The total number of steps to perform.\n            step_size: The number of internal steps the simulation\n                is advanced every turn.\n\n        Yields:\n            int: The current step.\n        \"\"\"\n        for _ in range(0, n_steps, step):\n            for i in range(0, step):\n                dp = self._forward_and_backward()\n                if i == 0:\n                    yield self.step_count\n\n                self._step_count += 1\n                self.update_parameters(dp)\n\n    def advance(self, steps: int = 1):\n        \"\"\"\n        Performs `steps` optimization steps.\n\n        steps: (Default=1) number of steps\n        \"\"\"\n        for _ in self.iter(steps):\n            pass\n\n    def run(\n        self,\n        n_iter,\n        out=None,\n        obs=None,\n        show_progress=True,\n        save_params_every=50,  # for default logger\n        write_every=50,  # for default logger\n        step_size=1,  # for default logger\n        callback=lambda *x: True,\n    ):\n        \"\"\"\n        Executes the Monte Carlo Variational optimization, updating the weights of the network\n        stored in this driver for `n_iter` steps and dumping values of the observables `obs`\n        in the output `logger`. If no logger is specified, creates a json file at `out`,\n        overwriting files with the same prefix.\n\n        By default uses :ref:`netket.logging.JsonLog`. To know about the output format\n        check it's documentation. The logger object is also returned at the end of this function\n        so that you can inspect the results without reading the json output.\n\n        Args:\n            n_iter: the total number of iterations\n            out: A logger object, or an iterable of loggers, to be used to store simulation log and data.\n                If this argument is a string, it will be used as output prefix for the standard JSON logger.\n            obs: An iterable containing all observables that should be computed\n            save_params_every: Every how many steps the parameters of the network should be\n                serialized to disk (ignored if logger is provided)\n            write_every: Every how many steps the json data should be flushed to disk (ignored if\n                logger is provided)\n            step_size: Every how many steps should observables be logged to disk (default=1)\n            show_progress: If true displays a progress bar (default=True)\n            callback: Callable or list of callable callback functions to stop training given a condition\n        \"\"\"\n\n        if not isinstance(n_iter, numbers.Number):\n            raise ValueError(\n                \"n_iter, the first positional argument to `run`, must be a number!\"\n            )\n\n        if obs is None:\n            obs = {}\n\n        if out is None:\n            out = tuple()\n            print(\n                \"No output specified (out=[apath|nk.logging.JsonLogger(...)]).\"\n                \"Running the optimization but not saving the output.\"\n            )\n\n        # Log only non-root nodes\n        if self._mynode == 0:\n            # if out is a path, create an overwriting Json Log for output\n            if isinstance(out, str):\n                loggers = (JsonLog(out, \"w\", save_params_every, write_every),)\n            else:\n                loggers = _to_iterable(out)\n        else:\n            loggers = tuple()\n            show_progress = False\n\n        callbacks = _to_iterable(callback)\n        callback_stop = False\n\n        with tqdm(total=n_iter, disable=not show_progress) as pbar:\n            old_step = self.step_count\n            first_step = True\n\n            for step in self.iter(n_iter, step_size):\n\n                log_data = self.estimate(obs)\n\n                # if the cost-function is defined then report it in the progress bar\n                if self._loss_stats is not None:\n                    pbar.set_postfix_str(self._loss_name + \"=\" + str(self._loss_stats))\n                    log_data[self._loss_name] = self._loss_stats\n\n                # Execute callbacks before loggers because they can append to log_data\n                for callback in callbacks:\n                    if not callback(step, log_data, self):\n                        callback_stop = True\n\n                for logger in loggers:\n                    logger(self.step_count, log_data, self.state)\n\n                if len(callbacks) > 0:\n                    if mpi.mpi_any(callback_stop):\n                        break\n\n                # Reset the timing of tqdm after the first step, to ignore compilation time\n                if first_step:\n                    first_step = False\n                    pbar.unpause()\n\n                # Update the progress bar\n                pbar.update(self.step_count - old_step)\n                old_step = self.step_count\n\n            # Final update so that it shows up filled.\n            pbar.update(self.step_count - old_step)\n\n        # flush at the end of the evolution so that final values are saved to\n        # file\n        for logger in loggers:\n            logger.flush(self.state)\n\n        return loggers\n\n    def estimate(self, observables):\n        \"\"\"\n        Return MCMC statistics for the expectation value of observables in the\n        current state of the driver.\n\n        Args:\n            observables: A pytree of operators for which statistics should be computed.\n\n        Returns:\n            A pytree of the same structure as the input, containing MCMC statistics\n            for the corresponding operators as leaves.\n        \"\"\"\n        return tree_map(self._estimate_stats, observables)\n\n    def update_parameters(self, dp):\n        \"\"\"\n        Updates the parameters of the machine using the optimizer in this driver\n\n        Args:\n            dp: the pytree containing the updates to the parameters\n        \"\"\"\n        self._optimizer_state, self.state.parameters = apply_gradient(\n            self._optimizer.update, self._optimizer_state, dp, self.state.parameters\n        )\n\n\n@partial(jax.jit, static_argnums=0)\ndef apply_gradient(optimizer_fun, optimizer_state, dp, params):\n    import optax\n\n    updates, new_optimizer_state = optimizer_fun(dp, optimizer_state, params)\n\n    new_params = optax.apply_updates(params, updates)\n    return new_optimizer_state, new_params\n", "path": "netket/driver/abstract_variational_driver.py"}]}