prompt: 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()
  
  ```
  
golden_diff: 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]:
  
verification_info: {"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"}]}