Fix issues with discontinuities

Ceyron · Jul 25, 2024 · 44a7931 · 44a7931
1 parent 1812fd8
commit 44a7931
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/exponax/ic/_discontinuities.py b/exponax/ic/_discontinuities.py
@@ -135,8 +135,8 @@ def gen_one_ic_fn(self, *, key: PRNGKeyArray) -> Discontinuity:
             key_1, key_2, key = jr.split(key, 3)
             lim_1 = jr.uniform(key_1, (), minval=0.0, maxval=self.domain_extent)
             lim_2 = jr.uniform(key_2, (), minval=0.0, maxval=self.domain_extent)
-            lower_limits.append(min(lim_1, lim_2))
-            upper_limits.append(max(lim_1, lim_2))
+            lower_limits.append(jnp.minimum(lim_1, lim_2))
+            upper_limits.append(jnp.maximum(lim_1, lim_2))
 
         lower_limits = tuple(lower_limits)
         upper_limits = tuple(upper_limits)