import pymc as pm
import numpy as np
from pymc_extras.model.marginal.marginal_model import marginalize

data = np.random.normal(size=10_000)

with pm.Model() as model:
    d = pm.Data("data", data)
    batched_data = pm.Minibatch(d, batch_size=100)
    x = pm.Normal("x", 0., 1.)
    b = pm.Bernoulli("b", 0.5, shape=(100,))
    y = pm.Normal("y", b*x, total_size=len(data), observed=batched_data)

model2 = marginalize(model, [b])

NotImplementedError                       Traceback (most recent call last)
File ~/upstream/pymc-extras/pymc_extras/model/marginal/marginal_model.py:561, in replace_finite_discrete_marginal_subgraph(fgraph, rv_to_marginalize, dependent_rvs, input_rvs)
    560 try:
--> 561     dependent_rvs_dim_connections = subgraph_batch_dim_connection(
    562         rv_to_marginalize, dependent_rvs
    563     )
    564 except (ValueError, NotImplementedError) as e:
    565     # For the perspective of the user this is a NotImplementedError

File ~/upstream/pymc-extras/pymc_extras/model/marginal/graph_analysis.py:365, in subgraph_batch_dim_connection(input_var, output_vars)
    364 var_dims = {input_var: tuple(range(input_var.type.ndim))}
--> 365 var_dims = _subgraph_batch_dim_connection(var_dims, [input_var], output_vars)
    366 ret = []

File ~/upstream/pymc-extras/pymc_extras/model/marginal/graph_analysis.py:317, in _subgraph_batch_dim_connection(var_dims, input_vars, output_vars)
    316     else:
--> 317         raise NotImplementedError(f"Marginalization through operation {node} not supported.")
    319 return var_dims

NotImplementedError: Marginalization through operation MinibatchRandomVariable(y, 10000) not supported.

The above exception was the direct cause of the following exception:

NotImplementedError                       Traceback (most recent call last)
Cell In[21], line 2
      1 from pymc_extras.model.marginal.marginal_model import marginalize
----> 2 model2 = marginalize(model, [b])

File ~/upstream/pymc-extras/pymc_extras/model/marginal/marginal_model.py:244, in marginalize(model, rvs_to_marginalize)
    237     other_direct_rv_ancestors = [
    238         rv
    239         for rv in find_conditional_input_rvs(dependent_rvs, all_rvs)
    240         if rv is not rv_to_marginalize
    241     ]
    242     input_rvs = _unique((*marginalized_rv_input_rvs, *other_direct_rv_ancestors))
--> 244     replace_finite_discrete_marginal_subgraph(fg, rv_to_marginalize, dependent_rvs, input_rvs)
    246 return model_from_fgraph(fg, mutate_fgraph=True)

File ~/upstream/pymc-extras/pymc_extras/model/marginal/marginal_model.py:566, in replace_finite_discrete_marginal_subgraph(fgraph, rv_to_marginalize, dependent_rvs, input_rvs)
    561     dependent_rvs_dim_connections = subgraph_batch_dim_connection(
    562         rv_to_marginalize, dependent_rvs
    563     )
    564 except (ValueError, NotImplementedError) as e:
    565     # For the perspective of the user this is a NotImplementedError
--> 566     raise NotImplementedError(
    567         "The graph between the marginalized and dependent RVs cannot be marginalized efficiently. "
    568         "You can try splitting the marginalized RV into separate components and marginalizing them separately."
    569     ) from e
    571 output_rvs = [rv_to_marginalize, *dependent_rvs]
    572 rng_updates = collect_default_updates(output_rvs, inputs=input_rvs, must_be_shared=False)

NotImplementedError: The graph between the marginalized and dependent RVs cannot be marginalized efficiently. You can try splitting the marginalized RV into separate components and marginalizing them separately.