From 0e59fc3905ca3126da04fade926949734f2e4ee6 Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Tue, 1 Apr 2025 17:41:53 +0100 Subject: [PATCH] readded lost code --- autofit/mapper/prior/abstract.py | 19 ++++++++++++++++--- 1 file changed, 16 insertions(+), 3 deletions(-) diff --git a/autofit/mapper/prior/abstract.py b/autofit/mapper/prior/abstract.py index 91605ed2c..63053123d 100644 --- a/autofit/mapper/prior/abstract.py +++ b/autofit/mapper/prior/abstract.py @@ -110,10 +110,8 @@ def factor(self): return self.message.factor 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( @@ -121,6 +119,21 @@ def assert_within_limits(self, value): ) ) + 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(