Update dependencies

bambi/backend/terms.py

@@ -371,8 +371,7 @@ def build_response_distribution(self, kwargs, pymc_backend):
             dist_rv = weighted_dist(self.name, weights, **kwargs, observed=observed, dims=dims)
         # All of the other response kinds are "not special" and thus are handled the same way
         else:
-            # dist_rv = distribution(self.name, **kwargs)
-            dist_rv = None
+            dist_rv = distribution(self.name, **kwargs)
 
         return dist_rv
 

pyproject.toml

@@ -23,7 +23,7 @@ dependencies = [
     "formulae>=0.5.3",
     "graphviz",
     "pandas>=1.0.0",
-    "pymc>=5.14.0",
+    "pymc@git+https://github.com/pymc-devs/pymc",
 ]
 
 [project.optional-dependencies]
@@ -37,12 +37,9 @@ dev = [
     "quartodoc==0.6.1",
     "seaborn>=0.9.0",
 ]
+
 jax = [
-    "bayeux-ml>=0.1.9",
-    "blackjax==1.1.1", # to make CI pass
-    "jax>=0.3.1",
-    "jaxlib>=0.3.1",
-    "numpyro>=0.9.0",
+    "bayeux-ml>=0.1.12",
 ]
 
 [project.urls]

tests/test_alternative_samplers.py

@@ -5,6 +5,7 @@
 
 import pytest
 
+#pytestmark = pytest.mark.skip("JAX DEPS ARE BROKEN!")
 
 MCMC_METHODS = [getattr(bx.mcmc, k).name for k in bx.mcmc.__all__]
 MCMC_METHODS_FILTERED = [
@@ -39,7 +40,7 @@ def test_inference_method_names_and_kwargs():
     assert "mcmc" in names["pymc"].keys()
     assert "vi" in names["pymc"].keys()
 
-    # Check bayeu inference method family. Currently, only MCMC methods are supported
+    # Check bayeux inference method family. Currently, only MCMC methods are supported
     assert "mcmc" in names["bayeux"].keys()
 
     # Ensure get_kwargs method raises an error if a non-supported method name is passed