Merge pull request #1020 from rhayes777/feature/jax_assert_within_limits

Update `assert_within_limits` on the `Prior` class

autofit/mapper/prior/abstract.py

@@ -116,17 +116,29 @@ def factor(self):
 
     def assert_within_limits(self, value):
 
-        if jax_wrapper.use_jax:
-            return
-
-        if not (self.lower_limit <= value <= self.upper_limit):
+        def exception_message():
             raise exc.PriorLimitException(
                 "The physical value {} for a prior "
                 "was not within its limits {}, {}".format(
                     value, self.lower_limit, self.upper_limit
                 )
             )
 
+        if jax_wrapper.use_jax:
+            import jax
+            jax.lax.cond(
+                jax.numpy.logical_or(
+                    value < self.lower_limit,
+                    value > self.upper_limit
+                ),
+                lambda _: jax.debug.callback(exception_message),
+                lambda _: None,
+                None
+            )
+
+        elif not (self.lower_limit <= value <= self.upper_limit):
+            exception_message()
+
     @staticmethod
     def for_class_and_attribute_name(cls, attribute_name):
         prior_dict = conf.instance.prior_config.for_class_and_suffix_path(