We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[TEST] Jax System Core Parameter Client
### What do you want to test?
Jax system core parameter client (in `mava/systems/jax/`)

### Outline of test structure
* Unit tests
* Test components and hooks

### Definition of done
Passing checks, cover all hooks, edge cases considered

### Mandatory checklist before making a PR
* [ ] The success criteria laid down in “Definition of done” are met.
* [ ] Test code is documented - docstrings for methods and classes, static types for arguments.
* [ ] Documentation is updated - README, CONTRIBUTING, or other documentation.

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `mava/systems/jax/parameter_client.py`
Content:
```
1 # python3
2 # Copyright 2021 InstaDeep Ltd. All rights reserved.
3 #
4 # Licensed under the Apache License, Version 2.0 (the "License");
5 # you may not use this file except in compliance with the License.
6 # You may obtain a copy of the License at
7 #
8 #     http://www.apache.org/licenses/LICENSE-2.0
9 #
10 # Unless required by applicable law or agreed to in writing, software
11 # distributed under the License is distributed on an "AS IS" BASIS,
12 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 # See the License for the specific language governing permissions and
14 # limitations under the License.
15 
16 """Parameter client for Jax system. Adapted from Deepmind's Acme library"""
17 
18 from concurrent import futures
19 from typing import Any, Dict, List, Optional, Union
20 
21 import jax
22 import numpy as np
23 
24 from mava.systems.jax.parameter_server import ParameterServer
25 from mava.utils.sort_utils import sort_str_num
26 
27 
28 class ParameterClient:
29     """A parameter client for updating parameters from a remote server."""
30 
31     def __init__(
32         self,
33         client: ParameterServer,
34         parameters: Dict[str, Any],
35         get_keys: List[str] = None,
36         set_keys: List[str] = None,
37         update_period: int = 1,
38         devices: Dict[str, Optional[Union[str, jax.xla.Device]]] = {},
39     ):
40         """Initialise the parameter server."""
41         self._all_keys = sort_str_num(list(parameters.keys()))
42         # TODO (dries): Is the below change correct?
43         self._get_keys = get_keys if get_keys is not None else []
44         self._set_keys = set_keys if set_keys is not None else []
45         self._parameters: Dict[str, Any] = parameters
46         self._get_call_counter = 0
47         self._set_call_counter = 0
48         self._set_get_call_counter = 0
49         self._update_period = update_period
50         self._client = client
51         self._devices = devices
52 
53         # note below it is assumed that if one device is specified with a string
54         # they all are - need to test this works
55         # TODO: (Dries/Arnu): check this
56         if len(self._devices) and isinstance(list(self._devices.values())[0], str):
57             for key, device in self._devices.items():
58                 self._devices[key] = jax.devices(device)[0]  # type: ignore
59 
60         self._request = lambda: client.get_parameters(self._get_keys)
61         self._request_all = lambda: client.get_parameters(self._all_keys)
62 
63         self._adjust = lambda: client.set_parameters(
64             {key: self._parameters[key] for key in self._set_keys},
65         )
66 
67         self._add = lambda params: client.add_to_parameters(params)
68 
69         # Create a single background thread to fetch parameters without necessarily
70         # blocking the actor.
71         self._executor = futures.ThreadPoolExecutor(max_workers=1)
72         self._async_add_buffer: Dict[str, Any] = {}
73         self._async_request = lambda: self._executor.submit(self._request)
74         self._async_adjust = lambda: self._executor.submit(self._adjust)
75         self._async_adjust_and_request = lambda: self._executor.submit(
76             self._adjust_and_request
77         )
78         self._async_add: Any = lambda params: self._executor.submit(
79             self._add(params)  # type: ignore
80         )
81 
82         # Initialize this client's future to None to indicate to the `update()`
83         # method that there is no pending/running request.
84         self._get_future: Optional[futures.Future] = None
85         self._set_future: Optional[futures.Future] = None
86         self._set_get_future: Optional[futures.Future] = None
87         self._add_future: Optional[futures.Future] = None
88 
89     def _adjust_and_request(self) -> None:
90         self._client.set_parameters(
91             {key: self._parameters[key] for key in self._set_keys},
92         )
93         self._copy(self._client.get_parameters(self._get_keys))
94 
95     def get_async(self) -> None:
96         """Asynchronously updates the parameters with the latest copy from server."""
97 
98         # Track the number of calls (we only update periodically).
99         if self._get_call_counter < self._update_period:
100             self._get_call_counter += 1
101 
102         period_reached: bool = self._get_call_counter >= self._update_period
103         if period_reached and self._get_future is None:
104             # The update period has been reached and no request has been sent yet, so
105             # making an asynchronous request now.
106             self._get_future = self._async_request()
107             self._get_call_counter = 0
108 
109         if self._get_future is not None and self._get_future.done():
110             # The active request is done so copy the result and remove the future.\
111             self._copy(self._get_future.result())
112             self._get_future = None
113 
114     def set_async(self) -> None:
115         """Asynchronously updates server with the set parameters."""
116         # Track the number of calls (we only update periodically).
117         if self._set_call_counter < self._update_period:
118             self._set_call_counter += 1
119 
120         period_reached: bool = self._set_call_counter >= self._update_period
121 
122         if period_reached and self._set_future is None:
123             # The update period has been reached and no request has been sent yet, so
124             # making an asynchronous request now.
125             self._set_future = self._async_adjust()
126             self._set_call_counter = 0
127             return
128         if self._set_future is not None and self._set_future.done():
129             self._set_future = None
130 
131     def set_and_get_async(self) -> None:
132         """Asynchronously updates server and gets from server."""
133         # Track the number of calls (we only update periodically).
134         if self._set_get_call_counter < self._update_period:
135             self._set_get_call_counter += 1
136         period_reached: bool = self._set_get_call_counter >= self._update_period
137 
138         if period_reached and self._set_get_future is None:
139             # The update period has been reached and no request has been sent yet, so
140             # making an asynchronous request now.
141             self._set_get_future = self._async_adjust_and_request()
142             self._set_get_call_counter = 0
143             return
144         if self._set_get_future is not None and self._set_get_future.done():
145             self._set_get_future = None
146 
147     def add_async(self, params: Dict[str, Any]) -> None:
148         """Asynchronously adds to server parameters."""
149         if self._add_future is not None and self._add_future.done():
150             self._add_future = None
151 
152         names = params.keys()
153         if self._add_future is None:
154             # The update period has been reached and no request has been sent yet, so
155             # making an asynchronous request now.
156             if not self._async_add_buffer:
157                 self._add_future = self._async_add(params)
158             else:
159                 for name in names:
160                     self._async_add_buffer[name] += params[name]
161 
162                 self._add_future = self._async_add(self._async_add_buffer)
163                 self._async_add_buffer = {}
164             return
165         else:
166             # The trainers is going to fast to keep up! Adding
167             # all the values up and only writing them when the
168             # process is ready.
169             if self._async_add_buffer:
170                 for name in names:
171                     self._async_add_buffer[name] += params[name]
172             else:
173                 for name in names:
174                     self._async_add_buffer[name] = params[name]
175 
176     def add_and_wait(self, params: Dict[str, Any]) -> None:
177         """Adds the specified parameters to the corresponding parameters in server \
178         and waits for the process to complete before continuing."""
179         self._client.add_to_parameters(params)
180 
181     def get_and_wait(self) -> None:
182         """Updates the get parameters with the latest copy from server \
183         and waits for the process to complete before continuing."""
184         self._copy(self._request())
185 
186     def get_all_and_wait(self) -> None:
187         """Updates all the parameters with the latest copy from server \
188         and waits for the process to complete before continuing."""
189         self._copy(self._request_all())
190 
191     def set_and_wait(self) -> None:
192         """Updates server with the set parameters \
193         and waits for the process to complete before continuing."""
194         self._adjust()
195 
196     # TODO(Dries/Arnu): this needs a bit of a cleanup
197     def _copy(self, new_parameters: Dict[str, Any]) -> None:
198         """Copies the new parameters to the old ones."""
199         for key in new_parameters.keys():
200             if isinstance(new_parameters[key], dict):
201                 for type1_key in new_parameters[key].keys():
202                     for type2_key in self._parameters[key][type1_key].keys():
203                         if self._devices:
204                             # Move variables to a proper device.
205                             # self._parameters[key][type1_key][
206                             #     type2_key
207                             # ] = jax.device_put(  # type: ignore
208                             #     new_parameters[key][type1_key],
209                             #     self._devices[key][type1_key],
210                             # )
211                             raise NotImplementedError(
212                                 "Support for devices"
213                                 + "have not been implemented"
214                                 + "yet in the parameter client."
215                             )
216                         else:
217                             self._parameters[key][type1_key][
218                                 type2_key
219                             ] = new_parameters[key][type1_key][type2_key]
220             elif isinstance(new_parameters[key], np.ndarray):
221                 if self._devices:
222                     self._parameters[key] = jax.device_put(
223                         new_parameters[key], self._devices[key]  # type: ignore
224                     )
225                 else:
226                     # Note (dries): These in-place operators are used instead
227                     # of direct assignment to not lose reference to the numpy
228                     # array.
229 
230                     self._parameters[key] *= 0
231                     # Remove last dim of numpy array if needed
232                     if new_parameters[key].shape != self._parameters[key].shape:
233                         self._parameters[key] += new_parameters[key][0]
234                     else:
235                         self._parameters[key] += new_parameters[key]
236 
237             elif isinstance(new_parameters[key], tuple):
238                 for i in range(len(self._parameters[key])):
239                     if self._devices:
240                         self._parameters[key][i] = jax.device_put(
241                             new_parameters[key][i],
242                             self._devices[key][i],  # type: ignore
243                         )
244                     else:
245                         self._parameters[key][i] = new_parameters[key][i]
246             else:
247                 NotImplementedError(
248                     f"Parameter type of {type(new_parameters[key])} not implemented."
249                 )
250
```
--- 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/mava/systems/jax/parameter_client.py b/mava/systems/jax/parameter_client.py
--- a/mava/systems/jax/parameter_client.py
+++ b/mava/systems/jax/parameter_client.py
@@ -244,6 +244,6 @@
                     else:
                         self._parameters[key][i] = new_parameters[key][i]
             else:
-                NotImplementedError(
+                raise NotImplementedError(
                     f"Parameter type of {type(new_parameters[key])} not implemented."
                 )

{"golden_diff": "diff --git a/mava/systems/jax/parameter_client.py b/mava/systems/jax/parameter_client.py\n--- a/mava/systems/jax/parameter_client.py\n+++ b/mava/systems/jax/parameter_client.py\n@@ -244,6 +244,6 @@\n                     else:\n                         self._parameters[key][i] = new_parameters[key][i]\n             else:\n-                NotImplementedError(\n+                raise NotImplementedError(\n                     f\"Parameter type of {type(new_parameters[key])} not implemented.\"\n                 )\n", "issue": "[TEST] Jax System Core Parameter Client\n### What do you want to test?\r\nJax system core parameter client (in `mava/systems/jax/`)\r\n\r\n### Outline of test structure\r\n* Unit tests\r\n* Test components and hooks\r\n\r\n### Definition of done\r\nPassing checks, cover all hooks, edge cases considered\r\n\r\n### Mandatory checklist before making a PR\r\n* [ ] The success criteria laid down in \u201cDefinition of done\u201d are met.\r\n* [ ] Test code is documented - docstrings for methods and classes, static types for arguments.\r\n* [ ] Documentation is updated - README, CONTRIBUTING, or other documentation.\n", "before_files": [{"content": "# python3\n# Copyright 2021 InstaDeep Ltd. 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\n\"\"\"Parameter client for Jax system. Adapted from Deepmind's Acme library\"\"\"\n\nfrom concurrent import futures\nfrom typing import Any, Dict, List, Optional, Union\n\nimport jax\nimport numpy as np\n\nfrom mava.systems.jax.parameter_server import ParameterServer\nfrom mava.utils.sort_utils import sort_str_num\n\n\nclass ParameterClient:\n    \"\"\"A parameter client for updating parameters from a remote server.\"\"\"\n\n    def __init__(\n        self,\n        client: ParameterServer,\n        parameters: Dict[str, Any],\n        get_keys: List[str] = None,\n        set_keys: List[str] = None,\n        update_period: int = 1,\n        devices: Dict[str, Optional[Union[str, jax.xla.Device]]] = {},\n    ):\n        \"\"\"Initialise the parameter server.\"\"\"\n        self._all_keys = sort_str_num(list(parameters.keys()))\n        # TODO (dries): Is the below change correct?\n        self._get_keys = get_keys if get_keys is not None else []\n        self._set_keys = set_keys if set_keys is not None else []\n        self._parameters: Dict[str, Any] = parameters\n        self._get_call_counter = 0\n        self._set_call_counter = 0\n        self._set_get_call_counter = 0\n        self._update_period = update_period\n        self._client = client\n        self._devices = devices\n\n        # note below it is assumed that if one device is specified with a string\n        # they all are - need to test this works\n        # TODO: (Dries/Arnu): check this\n        if len(self._devices) and isinstance(list(self._devices.values())[0], str):\n            for key, device in self._devices.items():\n                self._devices[key] = jax.devices(device)[0]  # type: ignore\n\n        self._request = lambda: client.get_parameters(self._get_keys)\n        self._request_all = lambda: client.get_parameters(self._all_keys)\n\n        self._adjust = lambda: client.set_parameters(\n            {key: self._parameters[key] for key in self._set_keys},\n        )\n\n        self._add = lambda params: client.add_to_parameters(params)\n\n        # Create a single background thread to fetch parameters without necessarily\n        # blocking the actor.\n        self._executor = futures.ThreadPoolExecutor(max_workers=1)\n        self._async_add_buffer: Dict[str, Any] = {}\n        self._async_request = lambda: self._executor.submit(self._request)\n        self._async_adjust = lambda: self._executor.submit(self._adjust)\n        self._async_adjust_and_request = lambda: self._executor.submit(\n            self._adjust_and_request\n        )\n        self._async_add: Any = lambda params: self._executor.submit(\n            self._add(params)  # type: ignore\n        )\n\n        # Initialize this client's future to None to indicate to the `update()`\n        # method that there is no pending/running request.\n        self._get_future: Optional[futures.Future] = None\n        self._set_future: Optional[futures.Future] = None\n        self._set_get_future: Optional[futures.Future] = None\n        self._add_future: Optional[futures.Future] = None\n\n    def _adjust_and_request(self) -> None:\n        self._client.set_parameters(\n            {key: self._parameters[key] for key in self._set_keys},\n        )\n        self._copy(self._client.get_parameters(self._get_keys))\n\n    def get_async(self) -> None:\n        \"\"\"Asynchronously updates the parameters with the latest copy from server.\"\"\"\n\n        # Track the number of calls (we only update periodically).\n        if self._get_call_counter < self._update_period:\n            self._get_call_counter += 1\n\n        period_reached: bool = self._get_call_counter >= self._update_period\n        if period_reached and self._get_future is None:\n            # The update period has been reached and no request has been sent yet, so\n            # making an asynchronous request now.\n            self._get_future = self._async_request()\n            self._get_call_counter = 0\n\n        if self._get_future is not None and self._get_future.done():\n            # The active request is done so copy the result and remove the future.\\\n            self._copy(self._get_future.result())\n            self._get_future = None\n\n    def set_async(self) -> None:\n        \"\"\"Asynchronously updates server with the set parameters.\"\"\"\n        # Track the number of calls (we only update periodically).\n        if self._set_call_counter < self._update_period:\n            self._set_call_counter += 1\n\n        period_reached: bool = self._set_call_counter >= self._update_period\n\n        if period_reached and self._set_future is None:\n            # The update period has been reached and no request has been sent yet, so\n            # making an asynchronous request now.\n            self._set_future = self._async_adjust()\n            self._set_call_counter = 0\n            return\n        if self._set_future is not None and self._set_future.done():\n            self._set_future = None\n\n    def set_and_get_async(self) -> None:\n        \"\"\"Asynchronously updates server and gets from server.\"\"\"\n        # Track the number of calls (we only update periodically).\n        if self._set_get_call_counter < self._update_period:\n            self._set_get_call_counter += 1\n        period_reached: bool = self._set_get_call_counter >= self._update_period\n\n        if period_reached and self._set_get_future is None:\n            # The update period has been reached and no request has been sent yet, so\n            # making an asynchronous request now.\n            self._set_get_future = self._async_adjust_and_request()\n            self._set_get_call_counter = 0\n            return\n        if self._set_get_future is not None and self._set_get_future.done():\n            self._set_get_future = None\n\n    def add_async(self, params: Dict[str, Any]) -> None:\n        \"\"\"Asynchronously adds to server parameters.\"\"\"\n        if self._add_future is not None and self._add_future.done():\n            self._add_future = None\n\n        names = params.keys()\n        if self._add_future is None:\n            # The update period has been reached and no request has been sent yet, so\n            # making an asynchronous request now.\n            if not self._async_add_buffer:\n                self._add_future = self._async_add(params)\n            else:\n                for name in names:\n                    self._async_add_buffer[name] += params[name]\n\n                self._add_future = self._async_add(self._async_add_buffer)\n                self._async_add_buffer = {}\n            return\n        else:\n            # The trainers is going to fast to keep up! Adding\n            # all the values up and only writing them when the\n            # process is ready.\n            if self._async_add_buffer:\n                for name in names:\n                    self._async_add_buffer[name] += params[name]\n            else:\n                for name in names:\n                    self._async_add_buffer[name] = params[name]\n\n    def add_and_wait(self, params: Dict[str, Any]) -> None:\n        \"\"\"Adds the specified parameters to the corresponding parameters in server \\\n        and waits for the process to complete before continuing.\"\"\"\n        self._client.add_to_parameters(params)\n\n    def get_and_wait(self) -> None:\n        \"\"\"Updates the get parameters with the latest copy from server \\\n        and waits for the process to complete before continuing.\"\"\"\n        self._copy(self._request())\n\n    def get_all_and_wait(self) -> None:\n        \"\"\"Updates all the parameters with the latest copy from server \\\n        and waits for the process to complete before continuing.\"\"\"\n        self._copy(self._request_all())\n\n    def set_and_wait(self) -> None:\n        \"\"\"Updates server with the set parameters \\\n        and waits for the process to complete before continuing.\"\"\"\n        self._adjust()\n\n    # TODO(Dries/Arnu): this needs a bit of a cleanup\n    def _copy(self, new_parameters: Dict[str, Any]) -> None:\n        \"\"\"Copies the new parameters to the old ones.\"\"\"\n        for key in new_parameters.keys():\n            if isinstance(new_parameters[key], dict):\n                for type1_key in new_parameters[key].keys():\n                    for type2_key in self._parameters[key][type1_key].keys():\n                        if self._devices:\n                            # Move variables to a proper device.\n                            # self._parameters[key][type1_key][\n                            #     type2_key\n                            # ] = jax.device_put(  # type: ignore\n                            #     new_parameters[key][type1_key],\n                            #     self._devices[key][type1_key],\n                            # )\n                            raise NotImplementedError(\n                                \"Support for devices\"\n                                + \"have not been implemented\"\n                                + \"yet in the parameter client.\"\n                            )\n                        else:\n                            self._parameters[key][type1_key][\n                                type2_key\n                            ] = new_parameters[key][type1_key][type2_key]\n            elif isinstance(new_parameters[key], np.ndarray):\n                if self._devices:\n                    self._parameters[key] = jax.device_put(\n                        new_parameters[key], self._devices[key]  # type: ignore\n                    )\n                else:\n                    # Note (dries): These in-place operators are used instead\n                    # of direct assignment to not lose reference to the numpy\n                    # array.\n\n                    self._parameters[key] *= 0\n                    # Remove last dim of numpy array if needed\n                    if new_parameters[key].shape != self._parameters[key].shape:\n                        self._parameters[key] += new_parameters[key][0]\n                    else:\n                        self._parameters[key] += new_parameters[key]\n\n            elif isinstance(new_parameters[key], tuple):\n                for i in range(len(self._parameters[key])):\n                    if self._devices:\n                        self._parameters[key][i] = jax.device_put(\n                            new_parameters[key][i],\n                            self._devices[key][i],  # type: ignore\n                        )\n                    else:\n                        self._parameters[key][i] = new_parameters[key][i]\n            else:\n                NotImplementedError(\n                    f\"Parameter type of {type(new_parameters[key])} not implemented.\"\n                )\n", "path": "mava/systems/jax/parameter_client.py"}], "after_files": [{"content": "# python3\n# Copyright 2021 InstaDeep Ltd. 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\n\"\"\"Parameter client for Jax system. Adapted from Deepmind's Acme library\"\"\"\n\nfrom concurrent import futures\nfrom typing import Any, Dict, List, Optional, Union\n\nimport jax\nimport numpy as np\n\nfrom mava.systems.jax.parameter_server import ParameterServer\nfrom mava.utils.sort_utils import sort_str_num\n\n\nclass ParameterClient:\n    \"\"\"A parameter client for updating parameters from a remote server.\"\"\"\n\n    def __init__(\n        self,\n        client: ParameterServer,\n        parameters: Dict[str, Any],\n        get_keys: List[str] = None,\n        set_keys: List[str] = None,\n        update_period: int = 1,\n        devices: Dict[str, Optional[Union[str, jax.xla.Device]]] = {},\n    ):\n        \"\"\"Initialise the parameter server.\"\"\"\n        self._all_keys = sort_str_num(list(parameters.keys()))\n        # TODO (dries): Is the below change correct?\n        self._get_keys = get_keys if get_keys is not None else []\n        self._set_keys = set_keys if set_keys is not None else []\n        self._parameters: Dict[str, Any] = parameters\n        self._get_call_counter = 0\n        self._set_call_counter = 0\n        self._set_get_call_counter = 0\n        self._update_period = update_period\n        self._client = client\n        self._devices = devices\n\n        # note below it is assumed that if one device is specified with a string\n        # they all are - need to test this works\n        # TODO: (Dries/Arnu): check this\n        if len(self._devices) and isinstance(list(self._devices.values())[0], str):\n            for key, device in self._devices.items():\n                self._devices[key] = jax.devices(device)[0]  # type: ignore\n\n        self._request = lambda: client.get_parameters(self._get_keys)\n        self._request_all = lambda: client.get_parameters(self._all_keys)\n\n        self._adjust = lambda: client.set_parameters(\n            {key: self._parameters[key] for key in self._set_keys},\n        )\n\n        self._add = lambda params: client.add_to_parameters(params)\n\n        # Create a single background thread to fetch parameters without necessarily\n        # blocking the actor.\n        self._executor = futures.ThreadPoolExecutor(max_workers=1)\n        self._async_add_buffer: Dict[str, Any] = {}\n        self._async_request = lambda: self._executor.submit(self._request)\n        self._async_adjust = lambda: self._executor.submit(self._adjust)\n        self._async_adjust_and_request = lambda: self._executor.submit(\n            self._adjust_and_request\n        )\n        self._async_add: Any = lambda params: self._executor.submit(\n            self._add(params)  # type: ignore\n        )\n\n        # Initialize this client's future to None to indicate to the `update()`\n        # method that there is no pending/running request.\n        self._get_future: Optional[futures.Future] = None\n        self._set_future: Optional[futures.Future] = None\n        self._set_get_future: Optional[futures.Future] = None\n        self._add_future: Optional[futures.Future] = None\n\n    def _adjust_and_request(self) -> None:\n        self._client.set_parameters(\n            {key: self._parameters[key] for key in self._set_keys},\n        )\n        self._copy(self._client.get_parameters(self._get_keys))\n\n    def get_async(self) -> None:\n        \"\"\"Asynchronously updates the parameters with the latest copy from server.\"\"\"\n\n        # Track the number of calls (we only update periodically).\n        if self._get_call_counter < self._update_period:\n            self._get_call_counter += 1\n\n        period_reached: bool = self._get_call_counter >= self._update_period\n        if period_reached and self._get_future is None:\n            # The update period has been reached and no request has been sent yet, so\n            # making an asynchronous request now.\n            self._get_future = self._async_request()\n            self._get_call_counter = 0\n\n        if self._get_future is not None and self._get_future.done():\n            # The active request is done so copy the result and remove the future.\\\n            self._copy(self._get_future.result())\n            self._get_future = None\n\n    def set_async(self) -> None:\n        \"\"\"Asynchronously updates server with the set parameters.\"\"\"\n        # Track the number of calls (we only update periodically).\n        if self._set_call_counter < self._update_period:\n            self._set_call_counter += 1\n\n        period_reached: bool = self._set_call_counter >= self._update_period\n\n        if period_reached and self._set_future is None:\n            # The update period has been reached and no request has been sent yet, so\n            # making an asynchronous request now.\n            self._set_future = self._async_adjust()\n            self._set_call_counter = 0\n            return\n        if self._set_future is not None and self._set_future.done():\n            self._set_future = None\n\n    def set_and_get_async(self) -> None:\n        \"\"\"Asynchronously updates server and gets from server.\"\"\"\n        # Track the number of calls (we only update periodically).\n        if self._set_get_call_counter < self._update_period:\n            self._set_get_call_counter += 1\n        period_reached: bool = self._set_get_call_counter >= self._update_period\n\n        if period_reached and self._set_get_future is None:\n            # The update period has been reached and no request has been sent yet, so\n            # making an asynchronous request now.\n            self._set_get_future = self._async_adjust_and_request()\n            self._set_get_call_counter = 0\n            return\n        if self._set_get_future is not None and self._set_get_future.done():\n            self._set_get_future = None\n\n    def add_async(self, params: Dict[str, Any]) -> None:\n        \"\"\"Asynchronously adds to server parameters.\"\"\"\n        if self._add_future is not None and self._add_future.done():\n            self._add_future = None\n\n        names = params.keys()\n        if self._add_future is None:\n            # The update period has been reached and no request has been sent yet, so\n            # making an asynchronous request now.\n            if not self._async_add_buffer:\n                self._add_future = self._async_add(params)\n            else:\n                for name in names:\n                    self._async_add_buffer[name] += params[name]\n\n                self._add_future = self._async_add(self._async_add_buffer)\n                self._async_add_buffer = {}\n            return\n        else:\n            # The trainers is going to fast to keep up! Adding\n            # all the values up and only writing them when the\n            # process is ready.\n            if self._async_add_buffer:\n                for name in names:\n                    self._async_add_buffer[name] += params[name]\n            else:\n                for name in names:\n                    self._async_add_buffer[name] = params[name]\n\n    def add_and_wait(self, params: Dict[str, Any]) -> None:\n        \"\"\"Adds the specified parameters to the corresponding parameters in server \\\n        and waits for the process to complete before continuing.\"\"\"\n        self._client.add_to_parameters(params)\n\n    def get_and_wait(self) -> None:\n        \"\"\"Updates the get parameters with the latest copy from server \\\n        and waits for the process to complete before continuing.\"\"\"\n        self._copy(self._request())\n\n    def get_all_and_wait(self) -> None:\n        \"\"\"Updates all the parameters with the latest copy from server \\\n        and waits for the process to complete before continuing.\"\"\"\n        self._copy(self._request_all())\n\n    def set_and_wait(self) -> None:\n        \"\"\"Updates server with the set parameters \\\n        and waits for the process to complete before continuing.\"\"\"\n        self._adjust()\n\n    # TODO(Dries/Arnu): this needs a bit of a cleanup\n    def _copy(self, new_parameters: Dict[str, Any]) -> None:\n        \"\"\"Copies the new parameters to the old ones.\"\"\"\n        for key in new_parameters.keys():\n            if isinstance(new_parameters[key], dict):\n                for type1_key in new_parameters[key].keys():\n                    for type2_key in self._parameters[key][type1_key].keys():\n                        if self._devices:\n                            # Move variables to a proper device.\n                            # self._parameters[key][type1_key][\n                            #     type2_key\n                            # ] = jax.device_put(  # type: ignore\n                            #     new_parameters[key][type1_key],\n                            #     self._devices[key][type1_key],\n                            # )\n                            raise NotImplementedError(\n                                \"Support for devices\"\n                                + \"have not been implemented\"\n                                + \"yet in the parameter client.\"\n                            )\n                        else:\n                            self._parameters[key][type1_key][\n                                type2_key\n                            ] = new_parameters[key][type1_key][type2_key]\n            elif isinstance(new_parameters[key], np.ndarray):\n                if self._devices:\n                    self._parameters[key] = jax.device_put(\n                        new_parameters[key], self._devices[key]  # type: ignore\n                    )\n                else:\n                    # Note (dries): These in-place operators are used instead\n                    # of direct assignment to not lose reference to the numpy\n                    # array.\n\n                    self._parameters[key] *= 0\n                    # Remove last dim of numpy array if needed\n                    if new_parameters[key].shape != self._parameters[key].shape:\n                        self._parameters[key] += new_parameters[key][0]\n                    else:\n                        self._parameters[key] += new_parameters[key]\n\n            elif isinstance(new_parameters[key], tuple):\n                for i in range(len(self._parameters[key])):\n                    if self._devices:\n                        self._parameters[key][i] = jax.device_put(\n                            new_parameters[key][i],\n                            self._devices[key][i],  # type: ignore\n                        )\n                    else:\n                        self._parameters[key][i] = new_parameters[key][i]\n            else:\n                raise NotImplementedError(\n                    f\"Parameter type of {type(new_parameters[key])} not implemented.\"\n                )\n", "path": "mava/systems/jax/parameter_client.py"}]}