Update register (#272)

* update register

* fix test

* update pymc branch

Author: aloctavodia

pymc_bart/__init__.py

@@ -11,7 +11,7 @@
 #   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
-import pymc as pm
+from pymc.sampling import mcmc
 
 from pymc_bart.bart import BART
 from pymc_bart.pgbart import PGBART
@@ -47,4 +47,6 @@
 __version__ = "0.11.0"
 
 
-pm.STEP_METHODS = list(pm.STEP_METHODS) + [PGBART]
+methods = mcmc.STEP_METHODS
+if not any(method is PGBART for method in methods):
+    methods.append(PGBART)

requirements.txt

@@ -1,5 +1,5 @@
 pytensor>=3.0.0
-pymc @ git+https://github.com/pymc-devs/pymc.git@v6
+pymc @ git+https://github.com/pymc-devs/pymc.git
 arviz-stats[xarray]>=1.1.0
 numba
 matplotlib

tests/test_bart.py

@@ -97,7 +97,7 @@ def test_shared_variable(response):
         idata = pm.sample(tune=100, draws=100, chains=2, random_seed=3415)
         ppc = pm.sample_posterior_predictive(idata)
         pm.set_data({"data_X": X[:3]})
-        ppc2 = pm.sample_posterior_predictive(idata)
+        ppc2 = pm.sample_posterior_predictive(idata, sample_vars=["mu", "y"])
 
     assert ppc.posterior_predictive["y"].shape == (2, 100, 50)
     assert ppc2.posterior_predictive["y"].shape == (2, 100, 3)