We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[TEST] Jax Datasets
### What do you want to test?
Jax dataset components

### 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/components/jax/building/datasets.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 """Commonly used dataset components for system builders"""
17 import abc
18 from dataclasses import dataclass
19 from typing import Any, Callable, Optional
20 
21 import reverb
22 from acme import datasets
23 
24 from mava.components.jax import Component
25 from mava.core_jax import SystemBuilder
26 
27 Transform = Callable[[reverb.ReplaySample], reverb.ReplaySample]
28 
29 
30 class TrainerDataset(Component):
31     @abc.abstractmethod
32     def __init__(
33         self,
34         config: Any,
35     ):
36         """_summary_
37 
38         Args:
39             config : _description_.
40         """
41         self.config = config
42 
43     @abc.abstractmethod
44     def on_building_trainer_dataset(self, builder: SystemBuilder) -> None:
45         """_summary_
46 
47         Args:
48             builder : _description_
49         """
50         pass
51 
52     @staticmethod
53     def name() -> str:
54         """_summary_
55 
56         Returns:
57             _description_
58         """
59         return "trainer_dataset"
60 
61 
62 @dataclass
63 class TransitionDatasetConfig:
64     sample_batch_size: int = 256
65     prefetch_size: Optional[int] = None
66     num_parallel_calls: int = 12
67     max_in_flight_samples_per_worker: Optional[int] = None
68     postprocess: Optional[Transform] = None
69     # dataset_name: str = "transition_dataset"
70 
71 
72 class TransitionDataset(TrainerDataset):
73     def __init__(
74         self,
75         config: TransitionDatasetConfig = TransitionDatasetConfig(),
76     ):
77         """_summary_
78 
79         Args:
80             config : _description_.
81         """
82         self.config = config
83 
84     def on_building_trainer_dataset(self, builder: SystemBuilder) -> None:
85         """_summary_
86 
87         Args:
88             builder : _description_
89         """
90         max_in_flight_samples_per_worker = self.config.max_in_flight_samples_per_worker
91         dataset = datasets.make_reverb_dataset(
92             table=builder.store.trainer_id,
93             server_address=builder.store.data_server_client.server_address,
94             batch_size=self.config.sample_batch_size,
95             prefetch_size=self.config.prefetch_size,
96             num_parallel_calls=self.config.num_parallel_calls,
97             max_in_flight_samples_per_worker=max_in_flight_samples_per_worker,
98             postprocess=self.config.postprocess,
99         )
100 
101         builder.store.dataset = iter(dataset)
102 
103     @staticmethod
104     def config_class() -> Optional[Callable]:
105         """Config class used for component.
106 
107         Returns:
108             config class/dataclass for component.
109         """
110         return TransitionDatasetConfig
111 
112 
113 @dataclass
114 class TrajectoryDatasetConfig:
115     sample_batch_size: int = 256
116     max_in_flight_samples_per_worker: int = 512
117     num_workers_per_iterator: int = -1
118     max_samples_per_stream: int = -1
119     rate_limiter_timeout_ms: int = -1
120     get_signature_timeout_secs: Optional[int] = None
121     # max_samples: int = -1
122     # dataset_name: str = "trajectory_dataset"
123 
124 
125 class TrajectoryDataset(TrainerDataset):
126     def __init__(
127         self,
128         config: TrajectoryDatasetConfig = TrajectoryDatasetConfig(),
129     ):
130         """_summary_
131 
132         Args:
133             config : _description_.
134         """
135         self.config = config
136 
137     def on_building_trainer_dataset(self, builder: SystemBuilder) -> None:
138         """_summary_
139 
140         Args:
141             builder : _description_
142         """
143         dataset = reverb.TrajectoryDataset.from_table_signature(
144             server_address=builder.store.data_server_client.server_address,
145             table=builder.store.trainer_id,
146             max_in_flight_samples_per_worker=2 * self.config.sample_batch_size,
147             num_workers_per_iterator=self.config.num_workers_per_iterator,
148             max_samples_per_stream=self.config.max_samples_per_stream,
149             rate_limiter_timeout_ms=self.config.rate_limiter_timeout_ms,
150             get_signature_timeout_secs=self.config.get_signature_timeout_secs,
151             # max_samples=self.config.max_samples,
152         )
153 
154         # Add batch dimension.
155         dataset = dataset.batch(self.config.sample_batch_size, drop_remainder=True)
156         builder.store.sample_batch_size = self.config.sample_batch_size
157 
158         builder.store.dataset_iterator = dataset.as_numpy_iterator()
159 
160     @staticmethod
161     def config_class() -> Optional[Callable]:
162         """Config class used for component.
163 
164         Returns:
165             config class/dataclass for component.
166         """
167         return TrajectoryDatasetConfig
168
```
--- 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/components/jax/building/datasets.py b/mava/components/jax/building/datasets.py
--- a/mava/components/jax/building/datasets.py
+++ b/mava/components/jax/building/datasets.py
@@ -98,7 +98,7 @@
             postprocess=self.config.postprocess,
         )
 
-        builder.store.dataset = iter(dataset)
+        builder.store.dataset_iterator = iter(dataset)
 
     @staticmethod
     def config_class() -> Optional[Callable]:

{"golden_diff": "diff --git a/mava/components/jax/building/datasets.py b/mava/components/jax/building/datasets.py\n--- a/mava/components/jax/building/datasets.py\n+++ b/mava/components/jax/building/datasets.py\n@@ -98,7 +98,7 @@\n             postprocess=self.config.postprocess,\n         )\n \n-        builder.store.dataset = iter(dataset)\n+        builder.store.dataset_iterator = iter(dataset)\n \n     @staticmethod\n     def config_class() -> Optional[Callable]:\n", "issue": "[TEST] Jax Datasets\n### What do you want to test?\r\nJax dataset components\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\"\"\"Commonly used dataset components for system builders\"\"\"\nimport abc\nfrom dataclasses import dataclass\nfrom typing import Any, Callable, Optional\n\nimport reverb\nfrom acme import datasets\n\nfrom mava.components.jax import Component\nfrom mava.core_jax import SystemBuilder\n\nTransform = Callable[[reverb.ReplaySample], reverb.ReplaySample]\n\n\nclass TrainerDataset(Component):\n    @abc.abstractmethod\n    def __init__(\n        self,\n        config: Any,\n    ):\n        \"\"\"_summary_\n\n        Args:\n            config : _description_.\n        \"\"\"\n        self.config = config\n\n    @abc.abstractmethod\n    def on_building_trainer_dataset(self, builder: SystemBuilder) -> None:\n        \"\"\"_summary_\n\n        Args:\n            builder : _description_\n        \"\"\"\n        pass\n\n    @staticmethod\n    def name() -> str:\n        \"\"\"_summary_\n\n        Returns:\n            _description_\n        \"\"\"\n        return \"trainer_dataset\"\n\n\n@dataclass\nclass TransitionDatasetConfig:\n    sample_batch_size: int = 256\n    prefetch_size: Optional[int] = None\n    num_parallel_calls: int = 12\n    max_in_flight_samples_per_worker: Optional[int] = None\n    postprocess: Optional[Transform] = None\n    # dataset_name: str = \"transition_dataset\"\n\n\nclass TransitionDataset(TrainerDataset):\n    def __init__(\n        self,\n        config: TransitionDatasetConfig = TransitionDatasetConfig(),\n    ):\n        \"\"\"_summary_\n\n        Args:\n            config : _description_.\n        \"\"\"\n        self.config = config\n\n    def on_building_trainer_dataset(self, builder: SystemBuilder) -> None:\n        \"\"\"_summary_\n\n        Args:\n            builder : _description_\n        \"\"\"\n        max_in_flight_samples_per_worker = self.config.max_in_flight_samples_per_worker\n        dataset = datasets.make_reverb_dataset(\n            table=builder.store.trainer_id,\n            server_address=builder.store.data_server_client.server_address,\n            batch_size=self.config.sample_batch_size,\n            prefetch_size=self.config.prefetch_size,\n            num_parallel_calls=self.config.num_parallel_calls,\n            max_in_flight_samples_per_worker=max_in_flight_samples_per_worker,\n            postprocess=self.config.postprocess,\n        )\n\n        builder.store.dataset = iter(dataset)\n\n    @staticmethod\n    def config_class() -> Optional[Callable]:\n        \"\"\"Config class used for component.\n\n        Returns:\n            config class/dataclass for component.\n        \"\"\"\n        return TransitionDatasetConfig\n\n\n@dataclass\nclass TrajectoryDatasetConfig:\n    sample_batch_size: int = 256\n    max_in_flight_samples_per_worker: int = 512\n    num_workers_per_iterator: int = -1\n    max_samples_per_stream: int = -1\n    rate_limiter_timeout_ms: int = -1\n    get_signature_timeout_secs: Optional[int] = None\n    # max_samples: int = -1\n    # dataset_name: str = \"trajectory_dataset\"\n\n\nclass TrajectoryDataset(TrainerDataset):\n    def __init__(\n        self,\n        config: TrajectoryDatasetConfig = TrajectoryDatasetConfig(),\n    ):\n        \"\"\"_summary_\n\n        Args:\n            config : _description_.\n        \"\"\"\n        self.config = config\n\n    def on_building_trainer_dataset(self, builder: SystemBuilder) -> None:\n        \"\"\"_summary_\n\n        Args:\n            builder : _description_\n        \"\"\"\n        dataset = reverb.TrajectoryDataset.from_table_signature(\n            server_address=builder.store.data_server_client.server_address,\n            table=builder.store.trainer_id,\n            max_in_flight_samples_per_worker=2 * self.config.sample_batch_size,\n            num_workers_per_iterator=self.config.num_workers_per_iterator,\n            max_samples_per_stream=self.config.max_samples_per_stream,\n            rate_limiter_timeout_ms=self.config.rate_limiter_timeout_ms,\n            get_signature_timeout_secs=self.config.get_signature_timeout_secs,\n            # max_samples=self.config.max_samples,\n        )\n\n        # Add batch dimension.\n        dataset = dataset.batch(self.config.sample_batch_size, drop_remainder=True)\n        builder.store.sample_batch_size = self.config.sample_batch_size\n\n        builder.store.dataset_iterator = dataset.as_numpy_iterator()\n\n    @staticmethod\n    def config_class() -> Optional[Callable]:\n        \"\"\"Config class used for component.\n\n        Returns:\n            config class/dataclass for component.\n        \"\"\"\n        return TrajectoryDatasetConfig\n", "path": "mava/components/jax/building/datasets.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\"\"\"Commonly used dataset components for system builders\"\"\"\nimport abc\nfrom dataclasses import dataclass\nfrom typing import Any, Callable, Optional\n\nimport reverb\nfrom acme import datasets\n\nfrom mava.components.jax import Component\nfrom mava.core_jax import SystemBuilder\n\nTransform = Callable[[reverb.ReplaySample], reverb.ReplaySample]\n\n\nclass TrainerDataset(Component):\n    @abc.abstractmethod\n    def __init__(\n        self,\n        config: Any,\n    ):\n        \"\"\"_summary_\n\n        Args:\n            config : _description_.\n        \"\"\"\n        self.config = config\n\n    @abc.abstractmethod\n    def on_building_trainer_dataset(self, builder: SystemBuilder) -> None:\n        \"\"\"_summary_\n\n        Args:\n            builder : _description_\n        \"\"\"\n        pass\n\n    @staticmethod\n    def name() -> str:\n        \"\"\"_summary_\n\n        Returns:\n            _description_\n        \"\"\"\n        return \"trainer_dataset\"\n\n\n@dataclass\nclass TransitionDatasetConfig:\n    sample_batch_size: int = 256\n    prefetch_size: Optional[int] = None\n    num_parallel_calls: int = 12\n    max_in_flight_samples_per_worker: Optional[int] = None\n    postprocess: Optional[Transform] = None\n    # dataset_name: str = \"transition_dataset\"\n\n\nclass TransitionDataset(TrainerDataset):\n    def __init__(\n        self,\n        config: TransitionDatasetConfig = TransitionDatasetConfig(),\n    ):\n        \"\"\"_summary_\n\n        Args:\n            config : _description_.\n        \"\"\"\n        self.config = config\n\n    def on_building_trainer_dataset(self, builder: SystemBuilder) -> None:\n        \"\"\"_summary_\n\n        Args:\n            builder : _description_\n        \"\"\"\n        max_in_flight_samples_per_worker = self.config.max_in_flight_samples_per_worker\n        dataset = datasets.make_reverb_dataset(\n            table=builder.store.trainer_id,\n            server_address=builder.store.data_server_client.server_address,\n            batch_size=self.config.sample_batch_size,\n            prefetch_size=self.config.prefetch_size,\n            num_parallel_calls=self.config.num_parallel_calls,\n            max_in_flight_samples_per_worker=max_in_flight_samples_per_worker,\n            postprocess=self.config.postprocess,\n        )\n\n        builder.store.dataset_iterator = iter(dataset)\n\n    @staticmethod\n    def config_class() -> Optional[Callable]:\n        \"\"\"Config class used for component.\n\n        Returns:\n            config class/dataclass for component.\n        \"\"\"\n        return TransitionDatasetConfig\n\n\n@dataclass\nclass TrajectoryDatasetConfig:\n    sample_batch_size: int = 256\n    max_in_flight_samples_per_worker: int = 512\n    num_workers_per_iterator: int = -1\n    max_samples_per_stream: int = -1\n    rate_limiter_timeout_ms: int = -1\n    get_signature_timeout_secs: Optional[int] = None\n    # max_samples: int = -1\n    # dataset_name: str = \"trajectory_dataset\"\n\n\nclass TrajectoryDataset(TrainerDataset):\n    def __init__(\n        self,\n        config: TrajectoryDatasetConfig = TrajectoryDatasetConfig(),\n    ):\n        \"\"\"_summary_\n\n        Args:\n            config : _description_.\n        \"\"\"\n        self.config = config\n\n    def on_building_trainer_dataset(self, builder: SystemBuilder) -> None:\n        \"\"\"_summary_\n\n        Args:\n            builder : _description_\n        \"\"\"\n        dataset = reverb.TrajectoryDataset.from_table_signature(\n            server_address=builder.store.data_server_client.server_address,\n            table=builder.store.trainer_id,\n            max_in_flight_samples_per_worker=2 * self.config.sample_batch_size,\n            num_workers_per_iterator=self.config.num_workers_per_iterator,\n            max_samples_per_stream=self.config.max_samples_per_stream,\n            rate_limiter_timeout_ms=self.config.rate_limiter_timeout_ms,\n            get_signature_timeout_secs=self.config.get_signature_timeout_secs,\n            # max_samples=self.config.max_samples,\n        )\n\n        # Add batch dimension.\n        dataset = dataset.batch(self.config.sample_batch_size, drop_remainder=True)\n        builder.store.sample_batch_size = self.config.sample_batch_size\n\n        builder.store.dataset_iterator = dataset.as_numpy_iterator()\n\n    @staticmethod\n    def config_class() -> Optional[Callable]:\n        \"\"\"Config class used for component.\n\n        Returns:\n            config class/dataclass for component.\n        \"\"\"\n        return TrajectoryDatasetConfig\n", "path": "mava/components/jax/building/datasets.py"}]}