prompt: You will be provided with a partial code base and an issue statement explaining a problem to resolve.
  <issue>
  [BUG][Jax] Default executor flattens observations
  ### Describe the bug
  The default executor flattens observations which is undesirable for environments which can benefit from spatial knowledge.
  
  ### To Reproduce
  `observation = executor.store.observation.observation.reshape((1, -1))` <- inside `action_selection.py`
  
  ### Expected behavior
  It was expected that observations would be fed to neural networks unedited.
  
  ### Possible Solution
  `observation = utils.add_batch_dim(executor.store.observation.observation)`
  
  
  </issue>
  <code>
  [start of mava/components/jax/executing/action_selection.py]
  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 """Execution components for system builders"""
  17 
  18 from dataclasses import dataclass
  19 
  20 import jax
  21 
  22 from mava.components.jax import Component
  23 from mava.core_jax import SystemExecutor
  24 
  25 
  26 @dataclass
  27 class ExecutorSelectActionProcessConfig:
  28     pass
  29 
  30 
  31 class FeedforwardExecutorSelectAction(Component):
  32     def __init__(
  33         self,
  34         config: ExecutorSelectActionProcessConfig = ExecutorSelectActionProcessConfig(),
  35     ):
  36         """_summary_
  37 
  38         Args:
  39             config : _description_.
  40         """
  41         self.config = config
  42 
  43     # Select actions
  44     def on_execution_select_actions(self, executor: SystemExecutor) -> None:
  45         """Summary"""
  46         executor.store.actions_info = {}
  47         executor.store.policies_info = {}
  48         for agent, observation in executor.store.observations.items():
  49             action_info, policy_info = executor.select_action(agent, observation)
  50             executor.store.actions_info[agent] = action_info
  51             executor.store.policies_info[agent] = policy_info
  52 
  53     # Select action
  54     def on_execution_select_action_compute(self, executor: SystemExecutor) -> None:
  55         """Summary"""
  56 
  57         agent = executor.store.agent
  58         network = executor.store.networks["networks"][
  59             executor.store.agent_net_keys[agent]
  60         ]
  61 
  62         observation = executor.store.observation.observation.reshape((1, -1))
  63         rng_key, executor.store.key = jax.random.split(executor.store.key)
  64 
  65         # TODO (dries): We are currently using jit in the networks per agent.
  66         # We can also try jit over all the agents in a for loop. This would
  67         # allow the jit function to save us even more time.
  68         executor.store.action_info, executor.store.policy_info = network.get_action(
  69             observation, rng_key
  70         )
  71 
  72     @staticmethod
  73     def name() -> str:
  74         """_summary_"""
  75         return "action_selector"
  76
  [end of mava/components/jax/executing/action_selection.py]
  </code>
  
  I need you to solve this issue by generating a single patch file that I can apply directly to this repository using git apply. Please respond with a single patch file in the following format.
  <patch>
  diff --git a/file.py b/file.py
  --- a/file.py
  +++ b/file.py
  @@ -1,27 +1,35 @@
   def euclidean(a, b):
  -    while b:
  -        a, b = b, a % b
  -    return a
  +    if b == 0:
  +        return a
  +    return euclidean(b, a % b)
   
   
   def bresenham(x0, y0, x1, y1):
       points = []
       dx = abs(x1 - x0)
       dy = abs(y1 - y0)
  -    sx = 1 if x0 < x1 else -1
  -    sy = 1 if y0 < y1 else -1
  -    err = dx - dy
  +    x, y = x0, y0
  +    sx = -1 if x0 > x1 else 1
  +    sy = -1 if y0 > y1 else 1
   
  -    while True:
  -        points.append((x0, y0))
  -        if x0 == x1 and y0 == y1:
  -            break
  -        e2 = 2 * err
  -        if e2 > -dy:
  -            err -= dy
  -            x0 += sx
  -        if e2 < dx:
  -            err += dx
  -            y0 += sy
  +    if dx > dy:
  +        err = dx / 2.0
  +        while x != x1:
  +            points.append((x, y))
  +            err -= dy
  +            if err < 0:
  +                y += sy
  +                err += dx
  +            x += sx
  +    else:
  +        err = dy / 2.0
  +        while y != y1:
  +            points.append((x, y))
  +            err -= dx
  +            if err < 0:
  +                x += sx
  +                err += dy
  +            y += sy
  +
  +    points.append((x, y))
       return points
  
  </patch>
golden_diff: diff --git a/mava/components/jax/executing/action_selection.py b/mava/components/jax/executing/action_selection.py
  --- a/mava/components/jax/executing/action_selection.py
  +++ b/mava/components/jax/executing/action_selection.py
  @@ -18,6 +18,7 @@
   from dataclasses import dataclass
   
   import jax
  +from acme.jax import utils
   
   from mava.components.jax import Component
   from mava.core_jax import SystemExecutor
  @@ -59,7 +60,7 @@
               executor.store.agent_net_keys[agent]
           ]
   
  -        observation = executor.store.observation.observation.reshape((1, -1))
  +        observation = utils.add_batch_dim(executor.store.observation.observation)
           rng_key, executor.store.key = jax.random.split(executor.store.key)
   
           # TODO (dries): We are currently using jit in the networks per agent.
  
verification_info: {"golden_diff": "diff --git a/mava/components/jax/executing/action_selection.py b/mava/components/jax/executing/action_selection.py\n--- a/mava/components/jax/executing/action_selection.py\n+++ b/mava/components/jax/executing/action_selection.py\n@@ -18,6 +18,7 @@\n from dataclasses import dataclass\n \n import jax\n+from acme.jax import utils\n \n from mava.components.jax import Component\n from mava.core_jax import SystemExecutor\n@@ -59,7 +60,7 @@\n             executor.store.agent_net_keys[agent]\n         ]\n \n-        observation = executor.store.observation.observation.reshape((1, -1))\n+        observation = utils.add_batch_dim(executor.store.observation.observation)\n         rng_key, executor.store.key = jax.random.split(executor.store.key)\n \n         # TODO (dries): We are currently using jit in the networks per agent.\n", "issue": "[BUG][Jax] Default executor flattens observations\n### Describe the bug\r\nThe default executor flattens observations which is undesirable for environments which can benefit from spatial knowledge.\r\n\r\n### To Reproduce\r\n`observation = executor.store.observation.observation.reshape((1, -1))` <- inside `action_selection.py`\r\n\r\n### Expected behavior\r\nIt was expected that observations would be fed to neural networks unedited.\r\n\r\n### Possible Solution\r\n`observation = utils.add_batch_dim(executor.store.observation.observation)`\r\n\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\"\"\"Execution components for system builders\"\"\"\n\nfrom dataclasses import dataclass\n\nimport jax\n\nfrom mava.components.jax import Component\nfrom mava.core_jax import SystemExecutor\n\n\n@dataclass\nclass ExecutorSelectActionProcessConfig:\n    pass\n\n\nclass FeedforwardExecutorSelectAction(Component):\n    def __init__(\n        self,\n        config: ExecutorSelectActionProcessConfig = ExecutorSelectActionProcessConfig(),\n    ):\n        \"\"\"_summary_\n\n        Args:\n            config : _description_.\n        \"\"\"\n        self.config = config\n\n    # Select actions\n    def on_execution_select_actions(self, executor: SystemExecutor) -> None:\n        \"\"\"Summary\"\"\"\n        executor.store.actions_info = {}\n        executor.store.policies_info = {}\n        for agent, observation in executor.store.observations.items():\n            action_info, policy_info = executor.select_action(agent, observation)\n            executor.store.actions_info[agent] = action_info\n            executor.store.policies_info[agent] = policy_info\n\n    # Select action\n    def on_execution_select_action_compute(self, executor: SystemExecutor) -> None:\n        \"\"\"Summary\"\"\"\n\n        agent = executor.store.agent\n        network = executor.store.networks[\"networks\"][\n            executor.store.agent_net_keys[agent]\n        ]\n\n        observation = executor.store.observation.observation.reshape((1, -1))\n        rng_key, executor.store.key = jax.random.split(executor.store.key)\n\n        # TODO (dries): We are currently using jit in the networks per agent.\n        # We can also try jit over all the agents in a for loop. This would\n        # allow the jit function to save us even more time.\n        executor.store.action_info, executor.store.policy_info = network.get_action(\n            observation, rng_key\n        )\n\n    @staticmethod\n    def name() -> str:\n        \"\"\"_summary_\"\"\"\n        return \"action_selector\"\n", "path": "mava/components/jax/executing/action_selection.py"}]}