Fix warnings

Author: patrick-kidger

diffrax/brownian/path.py

@@ -43,7 +43,9 @@ def __init__(
         key: "jax.random.PRNGKey",
     ):
         self.shape = (
-            jax.ShapeDtypeStruct(shape, None) if is_tuple_of_ints(shape) else shape
+            jax.ShapeDtypeStruct(shape, jax.dtypes.canonicalize_dtype(None))
+            if is_tuple_of_ints(shape)
+            else shape
         )
         self.key = key
         if any(

diffrax/brownian/tree.py

@@ -75,7 +75,9 @@ def __init__(
         self.t1 = t1
         self.tol = tol
         self.shape = (
-            jax.ShapeDtypeStruct(shape, None) if is_tuple_of_ints(shape) else shape
+            jax.ShapeDtypeStruct(shape, jax.dtypes.canonicalize_dtype(None))
+            if is_tuple_of_ints(shape)
+            else shape
         )
         if any(
             not jnp.issubdtype(x.dtype, jnp.inexact)