WIP Allow rebuilding graph in toposort_replace

Author: ricardoV94

pymc/model/core.py

@@ -1355,7 +1355,7 @@ def make_obs_var(
             elif not isinstance(data, Variable):
                 data = pt.as_tensor_variable(data, name=name)
 
-            if total_size:
+            if total_size is not None:
                 from pymc.variational.minibatch_rv import create_minibatch_rv
 
                 rv_var = create_minibatch_rv(rv_var, total_size)

pymc/model/transform/minibatch.py

@@ -16,12 +16,12 @@
 from pytensor import Variable
 from pytensor.graph import FunctionGraph, ancestors
 
-from build.lib.pymc.variational.minibatch_rv import MinibatchRandomVariable
 from pymc import Minibatch, Model
 from pymc.data import MinibatchOp
 from pymc.model.fgraph import ModelObservedRV, fgraph_from_model, model_from_fgraph
 from pymc.model.transform.basic import parse_vars
 from pymc.pytensorf import toposort_replace
+from pymc.variational.minibatch_rv import MinibatchRandomVariable
 
 
 def minibatch_model(
@@ -36,14 +36,16 @@ def minibatch_model(
 
     .. warning:: This transformation acts on the leading dimension of the specified data variables and dependent observed RVs. If a dimension other than the first is linked to the minibatched data variables, the resulting model will be invalid.
 
+    .. warning:: When minibatch_vars are not specified, all non-scalar data variables will be minibatch. This can be incorrect!
+
     Parameters
     ----------
     model : Model
         The original model to transform.
     batch_size : int
         The minibatch size to use.
     minibatch_vars : Sequence of Variable or string, optional
-        Data variables to convert to minibatch. If None, all data variables with a leading dimension of size None will be minibatched.
+        Data variables to convert to minibatch. If None, all non scalar data variables will be minibatched.
 
     Returns
     -------
@@ -78,9 +80,7 @@ def minibatch_model(
 
     if minibatch_vars is None:
         original_minibatch_vars = [
-            variable
-            for variable in model.data_vars
-            if (variable.type.ndim > 0) and (variable.type.shape[0] is None)
+            variable for variable in model.data_vars if variable.type.ndim > 0
         ]
     else:
         original_minibatch_vars = parse_vars(model, minibatch_vars)
@@ -89,11 +89,6 @@ def minibatch_model(
                 raise ValueError(
                     f"Cannot minibatch {variable.name} because it is a scalar variable."
                 )
-            if variable.type.shape[0] is not None:
-                raise ValueError(
-                    f"Cannot minibatch {variable.name} because its first dimension is static "
-                    f"(size={variable.type.shape[0]})."
-                )
 
     # TODO: Validate that this graph is actually valid to minibatch. Example: linear regression with sigma fixed
     #  shape, but mu from data --> y cannot be minibatched because of sigma.
@@ -124,10 +119,10 @@ def minibatch_model(
     minibatch_fgraph._coords = fgraph._coords  # type: ignore[attr-defined]
     minibatch_fgraph._dim_lengths = fgraph._dim_lengths  # type: ignore[attr-defined]
     toposort_replace(minibatch_fgraph, pre_minibatch_var_to_dummy)
-    toposort_replace(minibatch_fgraph, dummy_to_minibatch_var)
+    toposort_replace(minibatch_fgraph, dummy_to_minibatch_var, rebuild=True)
 
     # Then replace all observed RVs that depend on the minibatch variables with MinibatchRVs
-    dependent_replacements = {}
+    dependent_replacements = []
     total_size = (pre_minibatch_vars[0].owner.inputs[0].shape[0], ...)
     vars_to_minibatch_set = set(pre_minibatch_vars)
     for model_var in minibatch_fgraph.outputs:
@@ -141,11 +136,10 @@ def minibatch_model(
         # TODO: If vars_to_minibatch had a leading dim, we should check that the dependent RVs also has that same dim
         # And conversely other variables do not have that dim
         observed_rv = model_var.owner.inputs[0]
-        dependent_replacements[observed_rv] = create_minibatch_rv(
-            observed_rv, total_size=total_size
-        )
+        minibatch_rv = create_minibatch_rv(observed_rv, total_size=total_size)
+        dependent_replacements.append((observed_rv, minibatch_rv))
 
-    toposort_replace(minibatch_fgraph, tuple(dependent_replacements.items()))
+    toposort_replace(minibatch_fgraph, dependent_replacements, rebuild=True)
 
     # Finally reintroduce the original data variable outputs
     for pre_minibatch_var in pre_minibatch_vars:
@@ -191,11 +185,11 @@ def remove_minibatch(model: Model) -> Model:
     fgraph, _ = fgraph_from_model(model)
 
     replacements = []
-    for var in fgraph.apply_nodes:
-        if isinstance(var.op, MinibatchOp):
-            replacements.extend(zip(var.inputs, var.outputs))
-        elif isinstance(var.op, MinibatchRandomVariable):
-            replacements.append((var.outputs[0], var.inputs[0]))
+    for node in fgraph.apply_nodes:
+        if isinstance(node.op, MinibatchOp):
+            replacements.extend(zip(node.outputs[:-1], node.inputs[:-1]))
+        elif isinstance(node.op, MinibatchRandomVariable):
+            replacements.append((node.outputs[0], node.inputs[0]))
 
-    toposort_replace(fgraph, replacements)
+    toposort_replace(fgraph, replacements, rebuild=True)
     return model_from_fgraph(fgraph, mutate_fgraph=True)

pymc/pytensorf.py

@@ -13,7 +13,9 @@
 #   limitations under the License.
 import warnings
 
+from collections import deque
 from collections.abc import Iterable, Sequence
+from itertools import chain
 from typing import cast
 
 import numpy as np
@@ -1032,20 +1034,127 @@ def as_symbolic_string(x, **kwargs):
     return StringConstant(stringtype, x)
 
 
+def _replace_rebuild(
+    fgraph: FunctionGraph, replacements: Sequence[tuple[Variable, Variable]], **kwargs
+) -> FunctionGraph:
+    """Replace variables and rebuild dependent graph if needed.
+
+    Rebuilding allows for replacements that change the semantics of the graph
+    (different types), which may not be possible for all Ops.
+    """
+    fg_clients = fgraph.clients
+    fg_variables = fgraph.variables
+
+    def get_client_nodes(vars) -> set[Apply]:
+        # Start with the immediate clients of vars
+        nodes = set()
+        d = list(chain.from_iterable(fg_clients[var] for var in vars if var in fg_variables))
+        while d:
+            node, _ = d.pop()
+            if node in nodes or isinstance(node.op, Output):
+                continue
+            nodes.add(node)
+            # Keep walking to the successor clients
+            d.extend(chain.from_iterable(fg_clients[out] for out in node.outputs))
+        return nodes
+
+    repl_dict = dict(replacements)
+    root_nodes = {var.owner for var in repl_dict.keys()}
+
+    # Build sorted queue with all nodes that depend on replaced variables
+    topo_order = {node: order for order, node in enumerate(fgraph.toposort())}
+    d = deque(
+        sorted(
+            get_client_nodes(repl_dict.keys()),
+            key=lambda node: topo_order[node],
+        )
+    )
+    while d:
+        node = d.popleft()
+        if node in root_nodes:
+            continue
+
+        new_inputs = [repl_dict.get(i, i) for i in node.inputs]
+        if new_inputs == node.inputs:
+            continue
+
+        # We need to remake the node if:
+        #  1. The output type depends on an input value
+        #  2. Any of the input type changed
+        if getattr(node.op, "_output_type_depends_on_input_value", False):
+            remake_node = True
+        else:
+            remake_node = any(
+                inp.type != new_inp.type for inp, new_inp in zip(node.inputs, new_inputs)
+            )
+
+        if remake_node:
+            new_node = node.clone_with_new_inputs(new_inputs, strict=False)
+            fgraph.import_node(new_node, import_missing=True)
+
+            # We are not always allowed to call `fgraph.replace_all` because the output types may be incompatible
+            # We will keep the changes in repl_dict until we can replace a node without remaking it,
+            # or we arrive to the end of the graph, in which case we need to replace the FunctionGraph output
+            for out, new_out in zip(node.outputs, new_node.outputs):
+                new_out.name = out.name
+                repl_dict[out] = new_out
+        else:
+            fgraph.replace_all(tuple(zip(node.inputs, new_inputs)), import_missing=True)
+
+    # If the FunctionGraph outputs themselves were rebuilt we need to handle them
+    for i, (new_output, old_output) in enumerate(
+        zip(
+            (repl_dict.get(out, out) for out in fgraph.outputs),
+            fgraph.outputs,
+        )
+    ):
+        if new_output is old_output:
+            continue
+        fgraph.outputs[i] = new_output
+        fgraph.import_var(new_output, import_missing=True)
+        fgraph.clients[new_output] = [
+            # Output variables have a special Output Op client
+            # We need to transfer it to the new output.
+            # Any other uses of this output variable will already have been substituted in the loop above,
+            # or are part of other outputs we will subsitute next
+            (cl.op.make_node(new_output), idx) if isinstance(cl.op, Output) else (cl, idx)
+            for cl, idx in fgraph.clients[old_output]
+        ]
+    return fgraph
+
+
 def toposort_replace(
     fgraph: FunctionGraph,
     replacements: Sequence[tuple[Variable, Variable]],
     reverse: bool = False,
+    rebuild: bool = False,
 ) -> None:
     """Replace multiple variables in place in topological order."""
+    if rebuild and reverse:
+        raise NotImplementedError("reverse rebuild not yet supported")
+
     fgraph_toposort = {node: i for i, node in enumerate(fgraph.toposort())}
     fgraph_toposort[None] = -1  # Variables without owner are not in the toposort
     sorted_replacements = sorted(
         replacements,
         key=lambda pair: fgraph_toposort[pair[0].owner],
         reverse=reverse,
     )
-    fgraph.replace_all(sorted_replacements, import_missing=True)
+
+    if rebuild:
+        if len(replacements) > 1:
+            # In this case we need to modify the replacements recursively with each other
+            sorted_replacements = [list(pairs) for pairs in sorted_replacements]
+            for i in range(1, len(replacements)):
+                temp_fgraph = FunctionGraph(
+                    outputs=[repl for _, repl in sorted_replacements[i:]],
+                    clone=False,
+                )
+                _replace_rebuild(temp_fgraph, replacements=sorted_replacements[:i])
+                sorted_replacements[i][1] = temp_fgraph.outputs[0]
+        _replace_rebuild(fgraph, sorted_replacements)
+    else:
+        fgraph.replace_all(sorted_replacements, import_missing=True)
 
 
 def normalize_rng_param(rng: None | Variable) -> Variable:

pymc/variational/minibatch_rv.py

@@ -19,6 +19,7 @@
 from pytensor import Variable, config
 from pytensor.graph import Apply, Op
 from pytensor.tensor import NoneConst, TensorVariable, as_tensor_variable
+from pytensor.tensor.type_other import NoneTypeT
 
 from pymc.logprob.abstract import MeasurableOp, _logprob
 from pymc.logprob.basic import logp
@@ -33,7 +34,9 @@ class MinibatchRandomVariable(MeasurableOp, Op):
     def make_node(self, rv, *total_size):
         rv = as_tensor_variable(rv)
         total_size = [
-            as_tensor_variable(t, dtype="int64", ndim=0) if t is not None else NoneConst
+            t
+            if isinstance(t, Variable)
+            else (NoneConst if t is None else as_tensor_variable(t, dtype="int64", ndim=0))
             for t in total_size
         ]
         assert len(total_size) == rv.ndim
@@ -55,45 +58,67 @@ def perform(self, node, inputs, output_storage):
 
 def create_minibatch_rv(
     rv: TensorVariable,
-    total_size: int | None | Sequence[int | EllipsisType | None],
+    total_size: int | TensorVariable | Sequence[int | TensorVariable | EllipsisType | None],
 ) -> TensorVariable:
     """Create variable whose logp is rescaled by total_size."""
+    rv_ndim_supp = rv.owner.op.ndim_supp
+
     if isinstance(total_size, int):
-        if rv.ndim <= 1:
-            total_size = [total_size]
+        total_size = (total_size, *([None] * rv_ndim_supp))
+    elif isinstance(total_size, TensorVariable):
+        if total_size.type.ndim == 0:
+            total_size = (total_size, *([None] * rv_ndim_supp))
+        elif total_size.type.ndim == 1:
+            total_size = tuple(total_size)
         else:
-            missing_ndims = rv.ndim - 1
-            total_size = [total_size] + [None] * missing_ndims
-    elif isinstance(total_size, list | tuple):
-        total_size = list(total_size)
-        if Ellipsis in total_size:
-            # Replace Ellipsis by None
-            if total_size.count(Ellipsis) > 1:
-                raise ValueError("Only one Ellipsis can be present in total_size")
-            sep = total_size.index(Ellipsis)
-            begin = total_size[:sep]
-            end = total_size[sep + 1 :]
-            missing_ndims = max((rv.ndim - len(begin) - len(end), 0))
-            total_size = begin + [None] * missing_ndims + end
-        if len(total_size) > rv.ndim:
-            raise ValueError(f"Length of total_size {total_size} is langer than RV ndim {rv.ndim}")
-    else:
-        raise TypeError(f"Invalid type for total_size: {total_size}")
-
-    return cast(TensorVariable, minibatch_rv(rv, *total_size))
-
-
-def get_scaling(total_size: Sequence[Variable], shape: TensorVariable) -> TensorVariable:
+            raise ValueError(
+                f"Total size must be a 0d or 1d vector got {total_size} with {total_size.type.ndim} dimensions"
+            )
+
+    if not isinstance(total_size, list | tuple):
+        raise ValueError(f"Invalid type for total_size {total_size}: {type(total_size)}")
+
+    if Ellipsis in total_size:
+        # Replace Ellipsis by None
+        if total_size.count(Ellipsis) > 1:
+            raise ValueError("Only one Ellipsis can be present in total_size")
+        sep = total_size.index(Ellipsis)
+        begin = total_size[:sep]
+        end = total_size[sep + 1 :]
+        missing_ndims = max((rv_ndim_supp - len(begin) - len(end), 0))
+        total_size = (*begin, *([None] * missing_ndims), *end)
+
+    if (len(total_size) - rv_ndim_supp) not in (0, 1):
+        raise ValueError(
+            f"Length of total_size {total_size} not compatble with ndim_supp of RV {rv}, "
+            f"got {len(total_size)} but must be {rv_ndim_supp} or {rv_ndim_supp - 1}"
+        )
+
+    out = minibatch_rv(rv, *total_size)
+    assert isinstance(out.owner.op, MinibatchRandomVariable)
+    return cast(TensorVariable, out)
+
+
+def get_scaling(
+    total_size: Sequence[TensorVariable], shape: TensorVariable | Sequence[TensorVariable]
+) -> TensorVariable:
     """Get scaling constant for logp."""
     # mypy doesn't understand we can convert a shape TensorVariable into a tuple
-    shape = tuple(shape)  # type: ignore[assignment]
+    shape = tuple(shape)
+
+    if len(total_size) == (len(shape) - 1):
+        # This happens when RV has no batch dimensions
+        # In that case the total_size corresponds to a dummy shape of 1
+        total_size = (1, *total_size)
+
+    assert len(shape) == len(total_size)
 
-    # Scalar RV
-    if len(shape) == 0:  # type: ignore[arg-type]
-        coef = total_size[0] if not NoneConst.equals(total_size[0]) else 1.0
-    else:
-        coefs = [t / shape[i] for i, t in enumerate(total_size) if not NoneConst.equals(t)]
-        coef = pt.prod(coefs)
+    coefs = [
+        size / dim_length
+        for size, dim_length in zip(total_size, shape)
+        if not isinstance(size.type, NoneTypeT)
+    ]
+    coef = pt.prod(coefs) if len(coefs) > 1 else coefs[0]
 
     return pt.cast(coef, dtype=config.floatX)
 
@@ -102,4 +127,6 @@ def get_scaling(total_size: Sequence[Variable], shape: TensorVariable) -> Tensor
 def minibatch_rv_logprob(op, values, *inputs, **kwargs):
     [value] = values
     rv, *total_size = inputs
-    return logp(rv, value, **kwargs) * get_scaling(total_size, value.shape)
+    raw_logp = logp(rv, value, **kwargs)
+    scaled_logp = raw_logp * get_scaling(total_size, raw_logp.shape)
+    return scaled_logp