chore: fix matplotlib and jax typing issues

jumanji/specs_test.py

@@ -184,13 +184,13 @@ def test_shape_element_type_error(self) -> None:
 
     def test_dtype_type_error(self) -> None:
         with pytest.raises(TypeError):
-            specs.Array((1, 2, 3), "32")
+            specs.Array((1, 2, 3), "32")  # type: ignore
 
     def test_scalar_shape(self) -> None:
         specs.Array((), jnp.int32)
 
     def test_string_dtype_error(self) -> None:
-        specs.Array((1, 2, 3), "int32")
+        specs.Array((1, 2, 3), "int32")  # type: ignore
 
     def test_dtype(self) -> None:
         specs.Array((1, 2, 3), int)

jumanji/types.py

@@ -233,4 +233,4 @@ def get_valid_dtype(dtype: Union[jnp.dtype, type]) -> jnp.dtype:
     Returns:
         dtype converted to the correct type precision.
     """
-    return jnp.empty((), dtype).dtype
+    return jnp.empty((), dtype).dtype  # type: ignore

requirements/requirements.txt

@@ -2,7 +2,7 @@ chex>=0.1.3
 dm-env>=1.5
 gym>=0.22.0
 jax>=0.2.26
-matplotlib>=3.3.4
+matplotlib~=3.7.4
 numpy>=1.19.5
 Pillow>=9.0.0
 typing-extensions>=4.0.0