pymc-extras 0.2.6__py3-none-any.whl → 0.3.1__py3-none-any.whl

pymc_extras/inference/pathfinder/pathfinder.py

@@ -63,7 +63,7 @@ from rich.text import Text
 # TODO: change to typing.Self after Python versions greater than 3.10
 from typing_extensions import Self
 
-from pymc_extras.inference.laplace import add_data_to_inferencedata
+from pymc_extras.inference.laplace_approx.idata import add_data_to_inference_data
 from pymc_extras.inference.pathfinder.importance_sampling import (
     importance_sampling as _importance_sampling,
 )
@@ -1759,6 +1759,6 @@ def fit_pathfinder(
         importance_sampling=importance_sampling,
     )
 
-    idata = add_data_to_inferencedata(idata, progressbar, model, compile_kwargs)
+    idata = add_data_to_inference_data(idata, progressbar, model, compile_kwargs)
 
     return idata

pymc_extras/linearmodel.py

@@ -2,10 +2,12 @@ import numpy as np
 import pandas as pd
 import pymc as pm
 
+from sklearn.base import BaseEstimator
+
 from pymc_extras.model_builder import ModelBuilder
 
 
-class LinearModel(ModelBuilder):
+class LinearModel(ModelBuilder, BaseEstimator):
     def __init__(
         self, model_config: dict | None = None, sampler_config: dict | None = None, nsamples=100
     ):

pymc_extras/model/marginal/graph_analysis.py

@@ -5,6 +5,7 @@ from itertools import zip_longest
 
 from pymc import SymbolicRandomVariable
 from pymc.model.fgraph import ModelVar
+from pymc.variational.minibatch_rv import MinibatchRandomVariable
 from pytensor.graph import Variable, ancestors
 from pytensor.graph.basic import io_toposort
 from pytensor.tensor import TensorType, TensorVariable
@@ -313,6 +314,9 @@ def _subgraph_batch_dim_connection(var_dims: VAR_DIMS, input_vars, output_vars)
 
             var_dims[node.outputs[0]] = output_dims
 
+        elif isinstance(node.op, MinibatchRandomVariable):
+            var_dims[node.outputs[0]] = inputs_dims[0]
+
         else:
             raise NotImplementedError(f"Marginalization through operation {node} not supported.")