We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
TensorBoardLog with MPI
Not sure this can be solved in netket, but when using TensorBoardLog with MPI, I'm getting errors.

MWE (mpitest.py):
```python
import netket as nk
logger = nk.logging.TensorBoardLog("test_mpi_log")
```

When you run the above with
```console
mpirun -np 4 python ./mpitest.py
```

You'll see messages like:
```console
KeyError: 'test_mpi_log'

FileExistsError: [Errno 17] File exists: 'test_mpi_log'
Traceback (most recent call last):
  File "/usr/local/anaconda3/envs/localjw/lib/python3.8/site-packages/tensorboardX/record_writer.py", line 47, in directory_check
    factory = REGISTERED_FACTORIES[prefix]
```





--- 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/logging/tensorboard.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 from numbers import Number
16 
17 from netket.utils import deprecated
18 
19 
20 def tree_log(tree, root, data):
21     """
22     Maps all elements in tree, recursively calling tree_log with a new root string,
23     and when it reaches leaves pushes (string, leave) tuples to data.
24 
25     Args:
26         tree: a pytree where the leaf nodes contain data
27         root: the root of the tags used to log to tensorboard
28         data: a container modified in place
29 
30     """
31     if tree is None:
32         return
33     elif isinstance(tree, list):
34         for (i, val) in enumerate(tree):
35             tree_log(val, root + f"/{i}", data)
36 
37     elif isinstance(tree, list) and hasattr(tree, "_fields"):
38         for key in tree._fields:
39             tree_log(getattr(tree, key), root + f"/{key}", data)
40 
41     elif isinstance(tree, tuple):
42         for (i, val) in enumerate(tree):
43             tree_log(val, root + f"/{i}", data)
44 
45     elif isinstance(tree, dict):
46         for key, value in tree.items():
47             tree_log(value, root + f"/{key}", data)  # noqa: F722
48 
49     elif hasattr(tree, "to_compound"):
50         tree_log(tree.to_compound()[1], root, data)  # noqa: F722
51 
52     elif hasattr(tree, "to_dict"):
53         tree_log(tree.to_dict(), root, data)  # noqa: F722
54 
55     elif isinstance(tree, complex):
56         tree_log(tree.real, root + "/re", data)  # noqa: F722
57         tree_log(tree.imag, root + "/im", data)  # noqa: F722
58 
59     else:
60         data.append((root, tree))
61 
62 
63 class TensorBoardLog:
64     """
65     Creates a tensorboard logger using tensorboardX's summarywriter.
66     Refer to its documentation for further details
67 
68     https://tensorboardx.readthedocs.io/en/latest/tensorboard.html
69 
70     TensorBoardX must be installed.
71 
72     Args:
73         logdir (string): Save directory location. Default is
74           runs/**CURRENT_DATETIME_HOSTNAME**, which changes after each run.
75           Use hierarchical folder structure to compare
76           between runs easily. e.g. pass in 'runs/exp1', 'runs/exp2', etc.
77           for each new experiment to compare across them.
78         comment (string): Comment logdir suffix appended to the default
79           ``logdir``. If ``logdir`` is assigned, this argument has no effect.
80         purge_step (int):
81           When logging crashes at step :math:`T+X` and restarts at step :math:`T`,
82           any events whose global_step larger or equal to :math:`T` will be
83           purged and hidden from TensorBoard.
84           Note that crashed and resumed experiments should have the same ``logdir``.
85         max_queue (int): Size of the queue for pending events and
86           summaries before one of the 'add' calls forces a flush to disk.
87           Default is ten items.
88         flush_secs (int): How often, in seconds, to flush the
89           pending events and summaries to disk. Default is every two minutes.
90         filename_suffix (string): Suffix added to all event filenames in
91           the logdir directory. More details on filename construction in
92           tensorboard.summary.writer.event_file_writer.EventFileWriter.
93         write_to_disk (boolean):
94           If pass `False`, TensorBoardLog will not write to disk.
95     Examples:
96         Logging optimisation to tensorboard.
97 
98         >>> import pytest; pytest.skip("skip automated test of this docstring")
99         >>>
100         >>> import netket as nk
101         >>> # create a summary writer with automatically generated folder name.
102         >>> writer = nk.logging.TensorBoardLog()
103         >>> # folder location: runs/May04_22-14-54_s-MacBook-Pro.local/
104         >>> # create a summary writer using the specified folder name.
105         >>> writer = nk.logging.TensorBoardLog("my_experiment")
106         >>> # folder location: my_experiment
107         >>> # create a summary writer with comment appended.
108         >>> writer = nk.logging.TensorBoardLog(comment="LR_0.1_BATCH_16")
109         >>> # folder location: runs/May04_22-14-54_s-MacBook-Pro.localLR_0.1_BATCH_16/
110     """
111 
112     def __init__(
113         self,
114         *args,
115         **kwargs,
116     ):
117         from tensorboardX import SummaryWriter
118 
119         self._writer = SummaryWriter(*args, **kwargs)
120 
121         self._old_step = 0
122 
123     def __call__(self, step, item, machine):
124 
125         data = []
126         tree_log(item, "", data)
127 
128         for key, val in data:
129             if isinstance(val, Number):
130                 self._writer.add_scalar(key[1:], val, step)
131 
132         self._writer.flush()
133         self._old_step = step
134 
135     def _flush_log(self):
136         self._writer.flush()
137 
138     def _flush_params(self, _):
139         return None
140 
141     def flush(self, machine=None):
142         """
143         Writes to file the content of this logger.
144 
145         :param machine: optionally also writes the parameters of the machine.
146         """
147         self._flush_log()
148 
149         if machine is not None:
150             self._flush_params(machine)
151 
152 
153 # TODO: deprecate in 3.1
154 @deprecated(
155     "TBLog has been renamed to `TensorBoardLog` and will be removed in the next"
156     "minor release. Please update your usages."
157 )
158 def TBLog(*args, **kwargs):
159     return TensorBoardLog(*args, **kwargs)
160
```
--- 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/logging/tensorboard.py b/netket/logging/tensorboard.py
--- a/netket/logging/tensorboard.py
+++ b/netket/logging/tensorboard.py
@@ -114,13 +114,24 @@
         *args,
         **kwargs,
     ):
-        from tensorboardX import SummaryWriter
+        self._init_args = args
+        """Store the args for the lazily initialized SummaryWriter's constructor."""
+        self._init_kwargs = kwargs
+        """Store the kwargs for the lazily initialized SummaryWriter's constructor."""
 
-        self._writer = SummaryWriter(*args, **kwargs)
+        self._writer = None
+        """Lazily initialized summarywriter constructor"""
 
         self._old_step = 0
 
+    def _init_tensoboard(self):
+        from tensorboardX import SummaryWriter
+
+        self._writer = SummaryWriter(*self._init_args, **self._init_kwargs)
+
     def __call__(self, step, item, machine):
+        if self._writer is None:
+            self._init_tensoboard()
 
         data = []
         tree_log(item, "", data)
@@ -133,7 +144,8 @@
         self._old_step = step
 
     def _flush_log(self):
-        self._writer.flush()
+        if self._writer is not None:
+            self._writer.flush()
 
     def _flush_params(self, _):
         return None

{"golden_diff": "diff --git a/netket/logging/tensorboard.py b/netket/logging/tensorboard.py\n--- a/netket/logging/tensorboard.py\n+++ b/netket/logging/tensorboard.py\n@@ -114,13 +114,24 @@\n         *args,\n         **kwargs,\n     ):\n-        from tensorboardX import SummaryWriter\n+        self._init_args = args\n+        \"\"\"Store the args for the lazily initialized SummaryWriter's constructor.\"\"\"\n+        self._init_kwargs = kwargs\n+        \"\"\"Store the kwargs for the lazily initialized SummaryWriter's constructor.\"\"\"\n \n-        self._writer = SummaryWriter(*args, **kwargs)\n+        self._writer = None\n+        \"\"\"Lazily initialized summarywriter constructor\"\"\"\n \n         self._old_step = 0\n \n+    def _init_tensoboard(self):\n+        from tensorboardX import SummaryWriter\n+\n+        self._writer = SummaryWriter(*self._init_args, **self._init_kwargs)\n+\n     def __call__(self, step, item, machine):\n+        if self._writer is None:\n+            self._init_tensoboard()\n \n         data = []\n         tree_log(item, \"\", data)\n@@ -133,7 +144,8 @@\n         self._old_step = step\n \n     def _flush_log(self):\n-        self._writer.flush()\n+        if self._writer is not None:\n+            self._writer.flush()\n \n     def _flush_params(self, _):\n         return None\n", "issue": "TensorBoardLog with MPI\nNot sure this can be solved in netket, but when using TensorBoardLog with MPI, I'm getting errors.\r\n\r\nMWE (mpitest.py):\r\n```python\r\nimport netket as nk\r\nlogger = nk.logging.TensorBoardLog(\"test_mpi_log\")\r\n```\r\n\r\nWhen you run the above with\r\n```console\r\nmpirun -np 4 python ./mpitest.py\r\n```\r\n\r\nYou'll see messages like:\r\n```console\r\nKeyError: 'test_mpi_log'\r\n\r\nFileExistsError: [Errno 17] File exists: 'test_mpi_log'\r\nTraceback (most recent call last):\r\n  File \"/usr/local/anaconda3/envs/localjw/lib/python3.8/site-packages/tensorboardX/record_writer.py\", line 47, in directory_check\r\n    factory = REGISTERED_FACTORIES[prefix]\r\n```\r\n\r\n\r\n\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\nfrom numbers import Number\n\nfrom netket.utils import deprecated\n\n\ndef tree_log(tree, root, data):\n    \"\"\"\n    Maps all elements in tree, recursively calling tree_log with a new root string,\n    and when it reaches leaves pushes (string, leave) tuples to data.\n\n    Args:\n        tree: a pytree where the leaf nodes contain data\n        root: the root of the tags used to log to tensorboard\n        data: a container modified in place\n\n    \"\"\"\n    if tree is None:\n        return\n    elif isinstance(tree, list):\n        for (i, val) in enumerate(tree):\n            tree_log(val, root + f\"/{i}\", data)\n\n    elif isinstance(tree, list) and hasattr(tree, \"_fields\"):\n        for key in tree._fields:\n            tree_log(getattr(tree, key), root + f\"/{key}\", data)\n\n    elif isinstance(tree, tuple):\n        for (i, val) in enumerate(tree):\n            tree_log(val, root + f\"/{i}\", data)\n\n    elif isinstance(tree, dict):\n        for key, value in tree.items():\n            tree_log(value, root + f\"/{key}\", data)  # noqa: F722\n\n    elif hasattr(tree, \"to_compound\"):\n        tree_log(tree.to_compound()[1], root, data)  # noqa: F722\n\n    elif hasattr(tree, \"to_dict\"):\n        tree_log(tree.to_dict(), root, data)  # noqa: F722\n\n    elif isinstance(tree, complex):\n        tree_log(tree.real, root + \"/re\", data)  # noqa: F722\n        tree_log(tree.imag, root + \"/im\", data)  # noqa: F722\n\n    else:\n        data.append((root, tree))\n\n\nclass TensorBoardLog:\n    \"\"\"\n    Creates a tensorboard logger using tensorboardX's summarywriter.\n    Refer to its documentation for further details\n\n    https://tensorboardx.readthedocs.io/en/latest/tensorboard.html\n\n    TensorBoardX must be installed.\n\n    Args:\n        logdir (string): Save directory location. Default is\n          runs/**CURRENT_DATETIME_HOSTNAME**, which changes after each run.\n          Use hierarchical folder structure to compare\n          between runs easily. e.g. pass in 'runs/exp1', 'runs/exp2', etc.\n          for each new experiment to compare across them.\n        comment (string): Comment logdir suffix appended to the default\n          ``logdir``. If ``logdir`` is assigned, this argument has no effect.\n        purge_step (int):\n          When logging crashes at step :math:`T+X` and restarts at step :math:`T`,\n          any events whose global_step larger or equal to :math:`T` will be\n          purged and hidden from TensorBoard.\n          Note that crashed and resumed experiments should have the same ``logdir``.\n        max_queue (int): Size of the queue for pending events and\n          summaries before one of the 'add' calls forces a flush to disk.\n          Default is ten items.\n        flush_secs (int): How often, in seconds, to flush the\n          pending events and summaries to disk. Default is every two minutes.\n        filename_suffix (string): Suffix added to all event filenames in\n          the logdir directory. More details on filename construction in\n          tensorboard.summary.writer.event_file_writer.EventFileWriter.\n        write_to_disk (boolean):\n          If pass `False`, TensorBoardLog will not write to disk.\n    Examples:\n        Logging optimisation to tensorboard.\n\n        >>> import pytest; pytest.skip(\"skip automated test of this docstring\")\n        >>>\n        >>> import netket as nk\n        >>> # create a summary writer with automatically generated folder name.\n        >>> writer = nk.logging.TensorBoardLog()\n        >>> # folder location: runs/May04_22-14-54_s-MacBook-Pro.local/\n        >>> # create a summary writer using the specified folder name.\n        >>> writer = nk.logging.TensorBoardLog(\"my_experiment\")\n        >>> # folder location: my_experiment\n        >>> # create a summary writer with comment appended.\n        >>> writer = nk.logging.TensorBoardLog(comment=\"LR_0.1_BATCH_16\")\n        >>> # folder location: runs/May04_22-14-54_s-MacBook-Pro.localLR_0.1_BATCH_16/\n    \"\"\"\n\n    def __init__(\n        self,\n        *args,\n        **kwargs,\n    ):\n        from tensorboardX import SummaryWriter\n\n        self._writer = SummaryWriter(*args, **kwargs)\n\n        self._old_step = 0\n\n    def __call__(self, step, item, machine):\n\n        data = []\n        tree_log(item, \"\", data)\n\n        for key, val in data:\n            if isinstance(val, Number):\n                self._writer.add_scalar(key[1:], val, step)\n\n        self._writer.flush()\n        self._old_step = step\n\n    def _flush_log(self):\n        self._writer.flush()\n\n    def _flush_params(self, _):\n        return None\n\n    def flush(self, machine=None):\n        \"\"\"\n        Writes to file the content of this logger.\n\n        :param machine: optionally also writes the parameters of the machine.\n        \"\"\"\n        self._flush_log()\n\n        if machine is not None:\n            self._flush_params(machine)\n\n\n# TODO: deprecate in 3.1\n@deprecated(\n    \"TBLog has been renamed to `TensorBoardLog` and will be removed in the next\"\n    \"minor release. Please update your usages.\"\n)\ndef TBLog(*args, **kwargs):\n    return TensorBoardLog(*args, **kwargs)\n", "path": "netket/logging/tensorboard.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\nfrom numbers import Number\n\nfrom netket.utils import deprecated\n\n\ndef tree_log(tree, root, data):\n    \"\"\"\n    Maps all elements in tree, recursively calling tree_log with a new root string,\n    and when it reaches leaves pushes (string, leave) tuples to data.\n\n    Args:\n        tree: a pytree where the leaf nodes contain data\n        root: the root of the tags used to log to tensorboard\n        data: a container modified in place\n\n    \"\"\"\n    if tree is None:\n        return\n    elif isinstance(tree, list):\n        for (i, val) in enumerate(tree):\n            tree_log(val, root + f\"/{i}\", data)\n\n    elif isinstance(tree, list) and hasattr(tree, \"_fields\"):\n        for key in tree._fields:\n            tree_log(getattr(tree, key), root + f\"/{key}\", data)\n\n    elif isinstance(tree, tuple):\n        for (i, val) in enumerate(tree):\n            tree_log(val, root + f\"/{i}\", data)\n\n    elif isinstance(tree, dict):\n        for key, value in tree.items():\n            tree_log(value, root + f\"/{key}\", data)  # noqa: F722\n\n    elif hasattr(tree, \"to_compound\"):\n        tree_log(tree.to_compound()[1], root, data)  # noqa: F722\n\n    elif hasattr(tree, \"to_dict\"):\n        tree_log(tree.to_dict(), root, data)  # noqa: F722\n\n    elif isinstance(tree, complex):\n        tree_log(tree.real, root + \"/re\", data)  # noqa: F722\n        tree_log(tree.imag, root + \"/im\", data)  # noqa: F722\n\n    else:\n        data.append((root, tree))\n\n\nclass TensorBoardLog:\n    \"\"\"\n    Creates a tensorboard logger using tensorboardX's summarywriter.\n    Refer to its documentation for further details\n\n    https://tensorboardx.readthedocs.io/en/latest/tensorboard.html\n\n    TensorBoardX must be installed.\n\n    Args:\n        logdir (string): Save directory location. Default is\n          runs/**CURRENT_DATETIME_HOSTNAME**, which changes after each run.\n          Use hierarchical folder structure to compare\n          between runs easily. e.g. pass in 'runs/exp1', 'runs/exp2', etc.\n          for each new experiment to compare across them.\n        comment (string): Comment logdir suffix appended to the default\n          ``logdir``. If ``logdir`` is assigned, this argument has no effect.\n        purge_step (int):\n          When logging crashes at step :math:`T+X` and restarts at step :math:`T`,\n          any events whose global_step larger or equal to :math:`T` will be\n          purged and hidden from TensorBoard.\n          Note that crashed and resumed experiments should have the same ``logdir``.\n        max_queue (int): Size of the queue for pending events and\n          summaries before one of the 'add' calls forces a flush to disk.\n          Default is ten items.\n        flush_secs (int): How often, in seconds, to flush the\n          pending events and summaries to disk. Default is every two minutes.\n        filename_suffix (string): Suffix added to all event filenames in\n          the logdir directory. More details on filename construction in\n          tensorboard.summary.writer.event_file_writer.EventFileWriter.\n        write_to_disk (boolean):\n          If pass `False`, TensorBoardLog will not write to disk.\n    Examples:\n        Logging optimisation to tensorboard.\n\n        >>> import pytest; pytest.skip(\"skip automated test of this docstring\")\n        >>>\n        >>> import netket as nk\n        >>> # create a summary writer with automatically generated folder name.\n        >>> writer = nk.logging.TensorBoardLog()\n        >>> # folder location: runs/May04_22-14-54_s-MacBook-Pro.local/\n        >>> # create a summary writer using the specified folder name.\n        >>> writer = nk.logging.TensorBoardLog(\"my_experiment\")\n        >>> # folder location: my_experiment\n        >>> # create a summary writer with comment appended.\n        >>> writer = nk.logging.TensorBoardLog(comment=\"LR_0.1_BATCH_16\")\n        >>> # folder location: runs/May04_22-14-54_s-MacBook-Pro.localLR_0.1_BATCH_16/\n    \"\"\"\n\n    def __init__(\n        self,\n        *args,\n        **kwargs,\n    ):\n        self._init_args = args\n        \"\"\"Store the args for the lazily initialized SummaryWriter's constructor.\"\"\"\n        self._init_kwargs = kwargs\n        \"\"\"Store the kwargs for the lazily initialized SummaryWriter's constructor.\"\"\"\n\n        self._writer = None\n        \"\"\"Lazily initialized summarywriter constructor\"\"\"\n\n        self._old_step = 0\n\n    def _init_tensoboard(self):\n        from tensorboardX import SummaryWriter\n\n        self._writer = SummaryWriter(*self._init_args, **self._init_kwargs)\n\n    def __call__(self, step, item, machine):\n        if self._writer is None:\n            self._init_tensoboard()\n\n        data = []\n        tree_log(item, \"\", data)\n\n        for key, val in data:\n            if isinstance(val, Number):\n                self._writer.add_scalar(key[1:], val, step)\n\n        self._writer.flush()\n        self._old_step = step\n\n    def _flush_log(self):\n        if self._writer is not None:\n            self._writer.flush()\n\n    def _flush_params(self, _):\n        return None\n\n    def flush(self, machine=None):\n        \"\"\"\n        Writes to file the content of this logger.\n\n        :param machine: optionally also writes the parameters of the machine.\n        \"\"\"\n        self._flush_log()\n\n        if machine is not None:\n            self._flush_params(machine)\n\n\n# TODO: deprecate in 3.1\n@deprecated(\n    \"TBLog has been renamed to `TensorBoardLog` and will be removed in the next\"\n    \"minor release. Please update your usages.\"\n)\ndef TBLog(*args, **kwargs):\n    return TensorBoardLog(*args, **kwargs)\n", "path": "netket/logging/tensorboard.py"}]}