Fix warnings

patrick-kidger · Nov 15, 2022 · acb8ccd · acb8ccd
1 parent ea1bdc9
commit acb8ccd
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/diffrax/brownian/path.py b/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(

diff --git a/diffrax/brownian/tree.py b/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)