SearchLibrium 0.0.241__tar.gz → 0.0.243__tar.gz

{searchlibrium-0.0.241/src/SearchLibrium.egg-info → searchlibrium-0.0.243}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: SearchLibrium
-Version: 0.0.241
+Version: 0.0.243
 Summary: Automated specification search for discrete choice models (logit, mixed logit, nested logit, RRM and more) using metaheuristic optimisation
 Author: Mario Fernando Taco Morales, Alexander Paz, Prithvi Beeramole, Robert Burdett
 Author-email: Zeke Ahern <z.ahern@qut.edu.au>

{searchlibrium-0.0.241 → searchlibrium-0.0.243}/pyproject.toml

@@ -129,7 +129,7 @@ Documentation = "https://pypi.org/project/SearchLibrium/"
 realpython = "SearchLibrium.__main__:main"
 
 [tool.bumpver]
-current_version = "0.0.241"
+current_version = "0.0.243"
 version_pattern = "MAJOR.MINOR.PATCH"
 commit_message = "[skip ci] Bump version {old_version} -> {new_version}"
 commit = true

{searchlibrium-0.0.241 → searchlibrium-0.0.243}/src/SearchLibrium/MixedLogit.py

@@ -63,7 +63,7 @@ def _panel_weighted_shares(prob_np, panel_info):
     return ind, np.mean(ind, axis=0)
 
 class MixedLogit(DiscreteChoiceModel):
-    def __init__(self, halton_opts=None, distributions=['n', 'ln', 'nln', 't', 'tn', 'u'], _jax=True):
+    def __init__(self, halton_opts=None, distributions=['n', 'ln', 'nln', 't', 'tn', 'u'], _jax=None):
         super().__init__(_jax)
         self.descr = "MXL"
         self.halton_opts = halton_opts
@@ -713,8 +713,17 @@ class MixedLogit(DiscreteChoiceModel):
 
         # Optional numba engine (first-class; the T4 minimise_func patch is
         # the equivalent for older installs). Resolved here so the JAX fast
-        # path below is skipped when numba is selected.
-        use_numba = (getattr(self, 'engine', None) == 'numba')
+        # path below is skipped when numba is selected. An explicit
+        # model.engine wins; otherwise the global default (SL_ENGINE env var
+        # or numba_engine.set_default_engine) applies.
+        try:
+            try:
+                from numba_engine import effective_engine as _mx_eff_eng
+            except ImportError:
+                from .numba_engine import effective_engine as _mx_eff_eng
+            use_numba = (_mx_eff_eng(self) == 'numba')
+        except Exception:
+            use_numba = (getattr(self, 'engine', None) == 'numba')
 
         # 2x Kftrans - mean and lambda, 3x Krtrans - mean, s.d., lambda
         # Kchol, Kbw - relate to random variables, non-transformed

{searchlibrium-0.0.241 → searchlibrium-0.0.243}/src/SearchLibrium/MixedLogitGSE.py

@@ -116,7 +116,19 @@ class MixedLogitGSE(MixedLogit):
                 rng = np.random.default_rng(None if _seed is None
                                             else int(_seed) + 10007 * self.Kgrad)
                 self.gamma_draws = rng.standard_normal((self.N, self.Kgrad))
-        # Force JAX-only (scipy fallback doesn't handle gradient loadings)
+        # Force JAX-only (scipy fallback doesn't handle gradient loadings,
+        # and there is no numba port for the gradient-latent terms yet).
+        # If numba is the global default, say so instead of silently using JAX.
+        try:
+            try:
+                from numba_engine import numba_as_default as _gse_nad
+            except ImportError:
+                from .numba_engine import numba_as_default as _gse_nad
+            if bool(_gse_nad()):
+                print("[MixedLogitGSE] gradient loadings require JAX; ignoring "
+                      "the global numba default for this fit.")
+        except Exception:
+            pass
         self._jax = True
         super().fit()
 

{searchlibrium-0.0.241 → searchlibrium-0.0.243}/src/SearchLibrium/_choice_model.py

@@ -67,7 +67,12 @@ from time import time
 import numpy as np
 import pandas as pd
 import scipy.stats as ss
-import jax.numpy as jnp
+try:
+    import jax.numpy as jnp
+    JAX_AVAILABLE = True
+except ImportError:  # pragma: no cover - jax is required; guard is import robustness only
+    jnp = None
+    JAX_AVAILABLE = False
 
 
 
@@ -87,8 +92,65 @@ except ImportError:
 # A dictionary to manage backends
 BACKENDS = {
     "numpy": np,
-    "jax": jnp,
 }
+if JAX_AVAILABLE:
+    BACKENDS["jax"] = jnp
+
+_JAX_MISSING_WARNED = False
+
+
+def _warn_jax_missing():
+    """Warn once that JAX was requested but is not installed."""
+    global _JAX_MISSING_WARNED
+    if not _JAX_MISSING_WARNED:
+        _JAX_MISSING_WARNED = True
+        warnings.warn(
+            "JAX was requested but could not be imported; falling back to "
+            "the numpy/scipy backend. JAX is a required dependency — check "
+            "your installation.",
+            ImportWarning, stacklevel=3,
+        )
+
+
+def _numba_is_global_default():
+    """True when the process-wide default engine is numba (SL_ENGINE/env)."""
+    try:
+        try:
+            from .numba_engine import numba_as_default as _nad
+        except ImportError:
+            try:
+                from numba_engine import numba_as_default as _nad
+            except ImportError:
+                _nad = None
+        if _nad is not None:
+            return bool(_nad())
+    except Exception:
+        pass
+    try:
+        import os as _os
+        return str(_os.environ.get('SL_ENGINE', '')).strip().lower() == 'numba'
+    except Exception:
+        return False
+
+
+def resolve_jax_backend(requested=None):
+    """Resolve the array backend for ``_jax=<requested>``.
+
+    Returns ``(use_jax, backend_module)``. ``None`` (the default) follows
+    the global default engine: ``False``/numpy when numba is the global
+    default (``SL_ENGINE=numba`` or ``numba_engine.set_default_engine``),
+    ``True``/jax.numpy otherwise (legacy JAX-first behaviour). Explicit
+    ``True``/``False`` is honoured verbatim. When JAX is requested but not
+    installed this degrades to ``(False, numpy)`` with a one-time warning
+    instead of raising.
+    """
+    if requested is None:
+        requested = not _numba_is_global_default()
+    if requested and JAX_AVAILABLE:
+        return True, jnp
+    if requested:
+        _warn_jax_missing()
+    return False, np
 
 # Utility function to select the backend
 def get_backend(jax=True):
@@ -159,8 +221,11 @@ class DiscreteChoiceModel(ABC):
     ''' ---------------------------------------------------------- '''
     ''' Function                                                   '''
     ''' ---------------------------------------------------------- '''
-    def __init__(self, jax = True):
+    def __init__(self, jax = None):
     # {
+        # jax=None follows the global default engine (JAX-first unless
+        # SL_ENGINE=numba / set_default_engine('numba')); explicit
+        # True/False is honoured verbatim.
 
         self.reset_attributes()
         self.fit_intercept = False
@@ -202,8 +267,7 @@ class DiscreteChoiceModel(ABC):
         self.X_original, self.y_original = [], []
         self.weights, self.avail = [], []
         self.init_coeff = []
-        self._jax = jax
-        self.backend = get_backend(jax)
+        self._jax, self.backend = resolve_jax_backend(jax)
         self.descr = ""
     # }
 

{searchlibrium-0.0.241 → searchlibrium-0.0.243}/src/SearchLibrium/latent_class.py

@@ -32,6 +32,103 @@ def _sig_stars(pv: float) -> str:
     return ""
 
 
+# ---------------------------------------------------------------------------
+# Numba engine: fused weighted-MNL (value, grad) for the EM M-steps.
+# ---------------------------------------------------------------------------
+# The numpy path allocates several (N, J)/(N, J, K) temporaries per objective
+# evaluation; the numba kernel fuses the whole case loop into one pass with
+# O(J*K) stack temporaries. Numerics mirror the numpy branch exactly
+# (max-subtraction, 1e-300 clipping, L1/L2/sd penalties). Compiled once and
+# cached at module level; all dtypes are float64.
+
+_NB_MNL_OBJ = {"fn": None}
+
+
+def _numba_weighted_mnl():
+    """Return the cached njit weighted-MNL objective, or raise."""
+    if _NB_MNL_OBJ["fn"] is not None:
+        return _NB_MNL_OBJ["fn"]
+    try:
+        import numba as _nb
+    except Exception as exc:
+        raise RuntimeError(
+            "engine='numba' requested but numba is not importable. "
+            "Install numba or use engine='auto'/'jax'/'numpy'.") from exc
+
+    @_nb.njit(cache=True)
+    def _obj(beta, Xc, y, w, avail, l2, l1, sd_pen, sd_k):
+        n_obs = Xc.shape[0]
+        n_alts = Xc.shape[1]
+        n_par = Xc.shape[2]
+        ll = 0.0
+        grad = np.zeros(n_par, dtype=np.float64)
+        U = np.empty(n_alts, dtype=np.float64)
+        E = np.empty(n_alts, dtype=np.float64)
+        for n in range(n_obs):
+            wn = w[n]
+            umax = -1e300
+            for j in range(n_alts):
+                if avail[n, j] > 0.0:
+                    uj = 0.0
+                    for k in range(n_par):
+                        uj += Xc[n, j, k] * beta[k]
+                    U[j] = uj
+                    if uj > umax:
+                        umax = uj
+                else:
+                    U[j] = -1e300
+            s = 0.0
+            for j in range(n_alts):
+                if avail[n, j] > 0.0:
+                    ej = np.exp(U[j] - umax)
+                    E[j] = ej
+                    s += ej
+                else:
+                    E[j] = 0.0
+            if s <= 0.0:
+                continue  # no available alternative: skip (numpy path yields p=0)
+            for j in range(n_alts):
+                if avail[n, j] <= 0.0:
+                    continue
+                p = E[j] / s
+                yij = y[n, j]
+                if yij > 0.0:
+                    cp = p if p > 1e-300 else 1e-300
+                    ll += wn * np.log(cp)
+                d = (yij - p) * wn
+                if d != 0.0:
+                    for k in range(n_par):
+                        grad[k] += d * Xc[n, j, k]
+        s2 = 0.0
+        s1 = 0.0
+        for k in range(n_par):
+            s2 += beta[k] * beta[k]
+            ak = beta[k] if beta[k] >= 0.0 else -beta[k]
+            s1 += ak
+        pen = l2 * s2 + l1 * s1
+        if sd_pen > 0.0 and sd_k > 0 and sd_k <= n_par:
+            sds = 0.0
+            for k in range(n_par - sd_k, n_par):
+                sds += beta[k] * beta[k]
+            pen += sd_pen * sds
+        value = -ll + pen
+        for k in range(n_par):
+            b = beta[k]
+            g = -grad[k] + 2.0 * l2 * b
+            if b > 0.0:
+                g += l1
+            elif b < 0.0:
+                g -= l1
+            grad[k] = g
+        if sd_pen > 0.0 and sd_k > 0 and sd_k <= n_par:
+            for k in range(n_par - sd_k, n_par):
+                grad[k] += 2.0 * sd_pen * beta[k]
+        return value, grad
+
+    _NB_MNL_OBJ["fn"] = _obj
+    return _obj
+
+
 class LatentClassMixedLogit(DiscreteChoiceModel):
     """Fast latent-class discrete choice model with optional JAX acceleration.
 
@@ -47,7 +144,7 @@ class LatentClassMixedLogit(DiscreteChoiceModel):
         class_maxiter=50,
         tol=1e-6,
         random_state=0,
-        _jax=True,
+        _jax=None,
         n_init=1,
         optimise_membership=True,
         membership_maxiter=50,
@@ -60,6 +157,7 @@ class LatentClassMixedLogit(DiscreteChoiceModel):
         gamma_max_abs=None,
         share_floor=0.0,
         membership_standardize=False,
+        engine="auto",
     ):
         self.n_classes = int(n_classes)
         self.maxiter = int(maxiter)
@@ -99,6 +197,17 @@ class LatentClassMixedLogit(DiscreteChoiceModel):
         self.membership_standardize = bool(membership_standardize)
         self._memb_standardize_ = None
         self._memb_floor_shrinks_ = 0
+        # ── compute engine for the EM weighted M-steps ───────────────
+        # "auto" (default): JAX when available, else numpy/scipy.
+        # "jax" / "numba" / "numpy": force that backend (numba falls back
+        # with a warning when not importable). search() forwards **kwargs,
+        # so engine="numba" works there too.
+        eng = str(engine or "auto").strip().lower()
+        if eng not in ("auto", "jax", "numba", "numpy"):
+            warnings.warn(f"Unknown engine={engine!r}; using 'auto'.")
+            eng = "auto"
+        self.engine = eng
+        self._backend = None  # resolved lazily by _resolve_backend()
         # Diagnostics populated by fit()/fit_direct().
         self.starts_ = []
         self.sort_perm_ = None
@@ -124,9 +233,24 @@ class LatentClassMixedLogit(DiscreteChoiceModel):
         self.total_iter = 0
         self.num_params = None
         self.search_results = None
-        self._jax = bool(_jax)
         self._jax_enabled = False
 
+        if _jax is None:
+            # Follow the global default engine (SL_ENGINE=numba ->
+            # numpy/numba paths, no JAX); otherwise legacy JAX-first.
+            try:
+                from ._choice_model import resolve_jax_backend as _rjb
+            except ImportError:
+                try:
+                    from _choice_model import resolve_jax_backend as _rjb
+                except ImportError:
+                    _rjb = None
+            if _rjb is not None:
+                _jax = bool(_rjb(None)[0])
+            else:
+                _jax = True
+        self._jax = bool(_jax)
+
         if self._jax:
             try:
                 import jax
@@ -703,6 +827,63 @@ class LatentClassMixedLogit(DiscreteChoiceModel):
         setattr(self, cache_grad, self.jit(self.value_and_grad(objective)))
         return objective
 
+    def _numba_global_default(self) -> bool:
+        """True when the process-wide default engine is numba."""
+        try:
+            from .numba_engine import numba_as_default as _nad
+        except ImportError:
+            try:
+                from numba_engine import numba_as_default as _nad
+            except ImportError:
+                _nad = None
+        if _nad is not None:
+            try:
+                return bool(_nad())
+            except Exception:
+                return False
+        try:
+            import os as _os
+            return str(_os.environ.get('SL_ENGINE', '')).strip().lower() == 'numba'
+        except Exception:
+            return False
+
+    def _resolve_backend(self) -> str:
+        """Resolve the effective compute backend for EM M-steps.
+
+        Returns one of "jax" | "numba" | "numpy". Result is cached on
+        ``self._backend``; explicit ``engine="jax"/"numba"`` falls back with a
+        warning when the backend is unavailable.
+        """
+        if self._backend in ("jax", "numba", "numpy"):
+            return self._backend
+        eng = str(getattr(self, "engine", "auto") or "auto").lower()
+        if eng == "numba":
+            try:
+                import numba  # noqa: F401
+                self._backend = "numba"
+                return self._backend
+            except Exception:
+                warnings.warn("engine='numba' requested but numba is not "
+                              "importable; falling back.")
+        elif eng == "jax":
+            if self._jax_enabled:
+                self._backend = "jax"
+                return self._backend
+            warnings.warn("engine='jax' requested but JAX is unavailable; "
+                          "falling back.")
+        if eng == "numpy":
+            self._backend = "numpy"
+        else:  # "auto" or any fallback above
+            if self._numba_global_default():
+                try:
+                    import numba  # noqa: F401
+                    self._backend = "numba"
+                    return self._backend
+                except Exception:
+                    pass
+            self._backend = "jax" if self._jax_enabled else "numpy"
+        return self._backend
+
     def _weighted_m_step(self, beta0, weights, class_idx=None):
         """Weighted M-step for a single class's betas."""
         weights = np.asarray(weights, dtype=float)
@@ -712,9 +893,35 @@ class LatentClassMixedLogit(DiscreteChoiceModel):
         else:
             X_c = self.X
 
-        use_jax = self._jax_enabled
+        use_jax = self._resolve_backend() == "jax"
+        use_numba = self._resolve_backend() == "numba"
+
+        if use_numba:
+            # Fused njit objective: one pass per evaluation, no (N,J,K)
+            # temporaries. Arrays are made contiguous once per M-step (not per
+            # iteration) since L-BFGS-B calls `objective` hundreds of times.
+            _Xc = np.ascontiguousarray(X_c, dtype=np.float64)
+            _y = np.ascontiguousarray(self.y, dtype=np.float64)
+            _w = np.ascontiguousarray(weights, dtype=np.float64)
+            _a = np.ascontiguousarray(self.avail, dtype=np.float64)
+            _sd_pen, _sd_k = 0.0, 0
+            if hasattr(self, 'sd_penalty') and self.sd_penalty > 0:
+                _kbw = getattr(self, '_Kbw', getattr(self, 'Kr', 0))
+                try:
+                    _kbw = int(_kbw)
+                except Exception:
+                    _kbw = 0
+                if _kbw > 0 and int(np.asarray(beta0).size) >= _kbw:
+                    _sd_pen, _sd_k = float(self.sd_penalty), _kbw
+            _fn = _numba_weighted_mnl()
 
-        if use_jax:
+            def objective(beta):
+                _bb = np.ascontiguousarray(np.asarray(beta, dtype=float))
+                _v, _g = _fn(_bb, _Xc, _y, _w, _a,
+                             float(self.l2_penalty), float(self.l1_penalty),
+                             _sd_pen, _sd_k)
+                return float(_v), np.asarray(_g, dtype=float)
+        elif use_jax:
             self._build_jax_weighted_objective(
                 X_c if class_idx is not None else None,
                 class_idx=class_idx

{searchlibrium-0.0.241 → searchlibrium-0.0.243}/src/SearchLibrium/logistic_regression.py

@@ -41,7 +41,7 @@ from .multinomial_logit import MultinomialLogit
 class LogisticRegression(MultinomialLogit):
     """Binomial logistic regression fitted by SearchLibrium's MNL estimator."""
 
-    def __init__(self, _jax: bool = True, base_alt: int = 0, **kwargs):
+    def __init__(self, _jax=None, base_alt: int = 0, **kwargs):
         super().__init__(_jax=_jax)
         self.base_alt = int(base_alt)          # reference alternative (0 = no event)
         self._varnames: Optional[list] = None

{searchlibrium-0.0.241 → searchlibrium-0.0.243}/src/SearchLibrium/mixed_nested.py

@@ -13,9 +13,8 @@ class MixedNested(MixedLogit, NestedLogit):
     Mixed Nested Logit Model.
     """
 
-    def __init__(self, _jax=True):
+    def __init__(self, _jax=None):
         MixedLogit.__init__(self, _jax=_jax)
-        self._jax = _jax
         self._set_backend(_jax)
         self.descr = "Mixed Nested Logit"
         self.nests = {}
@@ -25,12 +24,25 @@ class MixedNested(MixedLogit, NestedLogit):
         self.base_param_count = 0
 
     def _set_backend(self, use_jax):
+        if use_jax is None:
+            # Follow the global default engine (numba -> numpy backend).
+            try:
+                from ._choice_model import resolve_jax_backend
+            except ImportError:
+                from _choice_model import resolve_jax_backend
+            use_jax, _be = resolve_jax_backend(None)
+        self._jax = bool(use_jax)
         if use_jax:
-            import jax
-
-            jax.config.update("jax_enable_x64", True)
-            import jax.numpy as jnp
-
+            try:
+                import jax
+                jax.config.update("jax_enable_x64", True)
+                import jax.numpy as jnp
+            except ImportError:
+                import numpy as _np_fallback
+                self._jax = False
+                self.np = _np_fallback
+                self.backend = _np_fallback
+                return
             self.np = jnp
             self.backend = jnp
         else:

{searchlibrium-0.0.241 → searchlibrium-0.0.243}/src/SearchLibrium/mixedrrm.py

@@ -413,8 +413,18 @@ class MixedRandomRegret(RandomRegret, MixedLogit):
             # RandomRegret.fit directly (NOT fit_jax: its fallback invokes
             # self.fit(), which would recurse back here via the MRO).
             return RandomRegret.fit(self)
+        # Engine resolution: explicit model.engine wins, otherwise the global
+        # default (SL_ENGINE env var or numba_engine.set_default_engine).
         try:
-            if getattr(self, 'engine', None) == 'numba':
+            try:
+                from numba_engine import effective_engine as _mrrm_eff_eng
+            except ImportError:
+                from .numba_engine import effective_engine as _mrrm_eff_eng
+            _mrrm_use_numba = (_mrrm_eff_eng(self) == 'numba')
+        except Exception:
+            _mrrm_use_numba = (getattr(self, 'engine', None) == 'numba')
+        try:
+            if _mrrm_use_numba:
                 try:
                     from numba_rrm import fit_mrrm_numba
                 except ImportError:

{searchlibrium-0.0.241 → searchlibrium-0.0.243}/src/SearchLibrium/multinomial_logit.py

@@ -184,7 +184,7 @@ class MultinomialLogit(DiscreteChoiceModel):
     ''' --------------------------------- '''
     ''' Function. Constructor             '''
     ''' --------------------------------- '''
-    def __init__(self, _jax = True): # {
+    def __init__(self, _jax = None): # {
         super(MultinomialLogit, self).__init__(_jax)  # Base class initialisations
         self.descr = "MNL"
     # }

{searchlibrium-0.0.241 → searchlibrium-0.0.243}/src/SearchLibrium/multinomial_nested.py

@@ -12,30 +12,39 @@ class NestedLogit(MultinomialLogit):
     Handles nested structure of alternatives.
     """
 
-    def __init__(self, _jax = True):
+    def __init__(self, _jax = None):
         super(NestedLogit, self).__init__(_jax)
         self.descr = "Nested Logit"
         self.robust = False
         self.robust_corr = None
 
-        # Dynamically set the backend once during initialization
-        if _jax:
-            import jax
-            jax.config.update("jax_enable_x64", True)  #
-            import jax.numpy as jnp
-            from jax import grad, jacfwd, jit, lax, vmap
-            from scipy.optimize import  minimize
-            self.np = jnp  # Assign JAX's NumPy-like module
-            self.jaxgrad = grad
-            self.vmap = vmap
-            self.lax = lax
-            self.jacfwd = jacfwd
-            self.jit = jit
-            from jaxopt import ScipyMinimize
-            from scipy.stats import norm
-            self.jaxoptmin = ScipyMinimize
-            self.minimize = minimize
-        else:
+        # Dynamically set the backend once during initialization.
+        # self._jax is already resolved by the base class (None follows the
+        # global default engine: numba -> numpy backend, no JAX import).
+        if self._jax:
+            try:
+                import jax
+                jax.config.update("jax_enable_x64", True)  #
+                import jax.numpy as jnp
+                from jax import grad, jacfwd, jit, lax, vmap
+                from scipy.optimize import  minimize
+                self.np = jnp  # Assign JAX's NumPy-like module
+                self.jaxgrad = grad
+                self.vmap = vmap
+                self.lax = lax
+                self.jacfwd = jacfwd
+                self.jit = jit
+                from jaxopt import ScipyMinimize
+                from scipy.stats import norm
+                self.jaxoptmin = ScipyMinimize
+                self.minimize = minimize
+            except ImportError:
+                import warnings as _warnings
+                _warnings.warn(
+                    "JAX backend requested but jax/jaxopt is not importable; "
+                    "falling back to the numpy backend.")
+                self._jax = False
+        if not self._jax:
             import numpy as np
             from scipy.optimize import minimize
             self.np = np  # Assign standard NumPy
@@ -63,22 +72,25 @@ class NestedLogit(MultinomialLogit):
         self.__dict__.update(state)
         # Re-import the right modules based on _jax flag
         if getattr(self, "_jax", False):
-            import jax
-            jax.config.update("jax_enable_x64", True)
-            import jax.numpy as jnp
-            from jax import grad, jacfwd, jit, lax, vmap
-            from jaxopt import ScipyMinimize
-            from scipy.optimize import minimize
-
-            self.np = jnp
-            self.jaxgrad = grad
-            self.vmap = vmap
-            self.lax = lax
-            self.jacfwd = jacfwd
-            self.jit = jit
-            self.jaxoptmin = ScipyMinimize
-            self.minimize = minimize
-        else:
+            try:
+                import jax
+                jax.config.update("jax_enable_x64", True)
+                import jax.numpy as jnp
+                from jax import grad, jacfwd, jit, lax, vmap
+                from jaxopt import ScipyMinimize
+                from scipy.optimize import minimize
+
+                self.np = jnp
+                self.jaxgrad = grad
+                self.vmap = vmap
+                self.lax = lax
+                self.jacfwd = jacfwd
+                self.jit = jit
+                self.jaxoptmin = ScipyMinimize
+                self.minimize = minimize
+            except ImportError:
+                self._jax = False
+        if not getattr(self, "_jax", False):
             import numpy as np
             from scipy.optimize import minimize
 

{searchlibrium-0.0.241 → searchlibrium-0.0.243}/src/SearchLibrium/multinomial_probit.py

@@ -25,7 +25,7 @@ class MultinomialProbit(MultinomialLogit):
     softmax with a probit-style probability construction based on normal CDFs.
     """
 
-    def __init__(self, _jax=True):
+    def __init__(self, _jax=None):
         super(MultinomialProbit, self).__init__(_jax)
         self.descr = "Multinomial Probit (pairwise normal approximation)"