ci: fix tests in pipeline

jumanji/environments/routing/connector/env_test.py

@@ -41,7 +41,7 @@ def is_target_on_grid(agent: Agent, grid: chex.Array) -> chex.Array:
     return jnp.any(grid[agent.target] == get_target(agent.id))
 
 
-def test_connector__reset(connector: Connector, key: jax.random.KeyArray) -> None:
+def test_connector__reset(connector: Connector, key: jax.random.PRNGKey) -> None:
     """Test that all heads and targets are on the board."""
     state, timestep = connector.reset(key)
 

jumanji/environments/routing/cvrp/reward_test.py

@@ -46,10 +46,10 @@ def test_sparse_reward__compute_tour_length() -> None:
     assert jnp.isclose(tour_length, 6.8649917)
 
     trajectory = jnp.array([0, 7, 8, 9, 10, 0, 1, 2, 3, 4, 5, 0, 6])
-    assert compute_tour_length(coordinates, trajectory) == 6.8649917
+    assert jnp.isclose(compute_tour_length(coordinates, trajectory), 6.8649917)
 
     trajectory = jnp.array([0, 7, 8, 9, 10, 0, 1, 2, 3, 4, 5, 0, 6, 0])
-    assert compute_tour_length(coordinates, trajectory) == 6.8649917
+    assert jnp.isclose(compute_tour_length(coordinates, trajectory), 6.8649917)
 
 
 def test_dense_reward(cvrp_dense_reward: CVRP, dense_reward: DenseReward) -> None:

jumanji/environments/routing/mmst/utils.py

@@ -596,10 +596,10 @@ def multi_random_walk(
     ]
 
     # Get the total number of edges we need to add when merging the graphs.
-    sum_ratio: int = np.sum(np.arange(1, num_agents))
+    sum_ratio = np.arange(1, num_agents).sum()
     frac = np.cumsum(
-        [total_edges_merge_graph * (i) / sum_ratio for i in range(1, num_agents - 1)]
-    )
+        total_edges_merge_graph * np.arange(1, num_agents - 1) / sum_ratio
+    ).astype(np.int32)
     edges_per_merge_graph = jnp.split(jnp.arange(total_edges_merge_graph), frac)
     num_edges_per_merge_graph = [len(edges) for edges in edges_per_merge_graph]
 
@@ -608,6 +608,7 @@ def multi_random_walk(
     total_edges = num_edges_per_sub_graph[0]
     merge_graph_keys = jax.random.split(base_key, num_agents - 1)
 
+    # TODO: could do a scan to speed up compilation
     for i in range(1, num_agents):
         total_edges += num_edges_per_sub_graph[i] + num_edges_per_merge_graph[i - 1]
         graph_i = correct_graph_offset(graphs[i - 1], nodes_offsets[i])

jumanji/environments/routing/tsp/reward_test.py

@@ -69,7 +69,7 @@ def test_sparse_reward(  # noqa: CCR001
                     tour_length = jnp.linalg.norm(
                         sorted_cities - sorted_cities_rolled, axis=-1
                     ).sum()
-                    assert reward == -tour_length
+                    assert jnp.isclose(reward, -tour_length)
                 else:
                     # Check that the reward is 0 for every non-final valid action.
                     assert reward == 0

pyproject.toml

@@ -42,5 +42,8 @@ module = [
     "hydra.*",
     "omegaconf.*",
     "huggingface_hub.*",
+    "requests.*",
+    "pkg_resources.*",
+    "PIL.*",
 ]
 ignore_missing_imports = true