Commit e76a412c authored by Gordian Edenhofer's avatar Gordian Edenhofer
Browse files

Gently catch not implemented JAX CF None Zero-Mode

parent acad6d56
Pipeline #119748 passed with stages
in 25 minutes
......@@ -729,11 +729,16 @@ class CorrelatedFieldMaker:
warn(f"unable to add JAX operator for {dofdex=}")
self._jax_cfm = None
if self._jax_cfm is not None:
self._jax_cfm.set_amplitude_total_offset(
offset_mean=offset_mean, offset_std=jax_offset_std
)
if not isinstance(self._azm, float):
self._azm._jax_expr = self._jax_cfm.azm
try:
self._jax_cfm.set_amplitude_total_offset(
offset_mean=offset_mean, offset_std=jax_offset_std
)
if not isinstance(self._azm, float):
self._azm._jax_expr = self._jax_cfm.azm
except TypeError as e:
self._jax_cfm = None
if isinstance(e, TypeError):
warn(f"no JAX operator for this configuration;\n{e}")
def finalize(self, prior_info=100):
"""Finishes model construction process and returns the constructed
......
......@@ -152,7 +152,8 @@ def SimpleCorrelatedField(
op._jax_expr = cf
op.amplitude._jax_expr = cfm.amplitude
op.power_spectrum._jax_expr = cfm.power_spectrum
except ImportError:
pass
except (ImportError, TypeError) as e:
if isinstance(e, TypeError):
warn(f"no JAX operator for this configuration;\n{e}")
return op
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment