float_like fix

geometric-kernels · Aug 15, 2023 · f799178 · f799178
1 parent 00573f7
commit f799178
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/geometric_kernels/lab_extras/jax/extras.py b/geometric_kernels/lab_extras/jax/extras.py
@@ -87,7 +87,7 @@ def float_like(reference: B.JAXNumeric):
     Return the type of the reference if it is a floating point type.
     Otherwise return `double` dtype of a backend based on the reference.
     """
-    reference_dtype = jnp.dtype(reference)
+    reference_dtype = reference.dtype
     if jnp.issubdtype(reference_dtype, jnp.floating):
         return reference_dtype
     else:

diff --git a/geometric_kernels/lab_extras/numpy/extras.py b/geometric_kernels/lab_extras/numpy/extras.py
@@ -74,7 +74,7 @@ def float_like(reference: B.NPNumeric):
     Return the type of the reference if it is a floating point type.
     Otherwise return `double` dtype of a backend based on the reference.
     """
-    reference_dtype = np.dtype(reference)
+    reference_dtype = reference.dtype
     if np.issubdtype(reference_dtype, np.floating):
         return reference_dtype
     else:

diff --git a/geometric_kernels/lab_extras/tensorflow/extras.py b/geometric_kernels/lab_extras/tensorflow/extras.py
@@ -95,7 +95,7 @@ def float_like(reference: B.TFNumeric):
     Return the type of the reference if it is a floating point type.
     Otherwise return `double` dtype of a backend based on the reference.
     """
-    reference_dtype = tf.dtype(reference)
+    reference_dtype = reference.dtype
     if reference_dtype.is_floating:
         return reference_dtype
     else: