Merge pull request #167 from invrs-io/types

fix type errors
invrs-io · Jan 22, 2025 · 87abb59 · 87abb59
2 parents 6d5fd90 + 20a48f1
commit 87abb59
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/src/invrs_gym/utils/transforms.py b/src/invrs_gym/utils/transforms.py
@@ -16,9 +16,9 @@ def rescaled_density_array(
     upper_bound: float,
 ) -> jnp.ndarray:
     """Return a density array for specified lower and upper bounds."""
-    array = density.array - density.lower_bound
-    array /= density.upper_bound - density.lower_bound
-    array *= upper_bound - lower_bound
+    array = jnp.asarray(density.array - density.lower_bound)
+    array /= jnp.asarray(density.upper_bound - density.lower_bound)
+    array *= jnp.asarray(upper_bound - lower_bound)
     return jnp.asarray(array + lower_bound)