pymc-extras 0.2.0__py3-none-any.whl

pymc_extras/model/marginal/distributions.py

@@ -0,0 +1,276 @@
+from collections.abc import Sequence
+
+import numpy as np
+import pytensor.tensor as pt
+
+from pymc.distributions import Bernoulli, Categorical, DiscreteUniform
+from pymc.logprob.abstract import MeasurableOp, _logprob
+from pymc.logprob.basic import conditional_logp, logp
+from pymc.pytensorf import constant_fold
+from pytensor import Variable
+from pytensor.compile.builders import OpFromGraph
+from pytensor.compile.mode import Mode
+from pytensor.graph import Op, vectorize_graph
+from pytensor.graph.replace import clone_replace, graph_replace
+from pytensor.scan import map as scan_map
+from pytensor.scan import scan
+from pytensor.tensor import TensorVariable
+
+from pymc_extras.distributions import DiscreteMarkovChain
+
+
+class MarginalRV(OpFromGraph, MeasurableOp):
+    """Base class for Marginalized RVs"""
+
+    def __init__(self, *args, dims_connections: tuple[tuple[int | None]], **kwargs) -> None:
+        self.dims_connections = dims_connections
+        super().__init__(*args, **kwargs)
+
+    @property
+    def support_axes(self) -> tuple[tuple[int]]:
+        """Dimensions of dependent RVs that belong to the core (non-batched) marginalized variable."""
+        marginalized_ndim_supp = self.inner_outputs[0].owner.op.ndim_supp
+        support_axes_vars = []
+        for dims_connection in self.dims_connections:
+            ndim = len(dims_connection)
+            marginalized_supp_axes = ndim - marginalized_ndim_supp
+            support_axes_vars.append(
+                tuple(
+                    -i
+                    for i, dim in enumerate(reversed(dims_connection), start=1)
+                    if (dim is None or dim > marginalized_supp_axes)
+                )
+            )
+        return tuple(support_axes_vars)
+
+
+class MarginalFiniteDiscreteRV(MarginalRV):
+    """Base class for Marginalized Finite Discrete RVs"""
+
+
+class MarginalDiscreteMarkovChainRV(MarginalRV):
+    """Base class for Marginalized Discrete Markov Chain RVs"""
+
+
+def get_domain_of_finite_discrete_rv(rv: TensorVariable) -> tuple[int, ...]:
+    op = rv.owner.op
+    dist_params = rv.owner.op.dist_params(rv.owner)
+    if isinstance(op, Bernoulli):
+        return (0, 1)
+    elif isinstance(op, Categorical):
+        [p_param] = dist_params
+        [p_param_length] = constant_fold([p_param.shape[-1]])
+        return tuple(range(p_param_length))
+    elif isinstance(op, DiscreteUniform):
+        lower, upper = constant_fold(dist_params)
+        return tuple(np.arange(lower, upper + 1))
+    elif isinstance(op, DiscreteMarkovChain):
+        P, *_ = dist_params
+        return tuple(range(pt.get_vector_length(P[-1])))
+
+    raise NotImplementedError(f"Cannot compute domain for op {op}")
+
+
+def reduce_batch_dependent_logps(
+    dependent_dims_connections: Sequence[tuple[int | None, ...]],
+    dependent_ops: Sequence[Op],
+    dependent_logps: Sequence[TensorVariable],
+) -> TensorVariable:
+    """Combine the logps of dependent RVs and align them with the marginalized logp.
+
+    This requires reducing extra batch dims and transposing when they are not aligned.
+
+       idx = pm.Bernoulli(idx, shape=(3, 2))  # 0, 1
+       pm.Normal("dep1", mu=idx.T[..., None] * 2, shape=(3, 2, 5))
+       pm.Normal("dep2", mu=idx * 2, shape=(7, 2, 3))
+
+       marginalize(idx)
+
+       The marginalized op will have dims_connections = [(1, 0, None), (None, 0, 1)]
+       which tells us we need to reduce the last axis of dep1 logp and the first of dep2 logp,
+       as well as transpose the remaining axis of dep1 logp before adding the two element-wise.
+
+    """
+    from pymc_extras.model.marginal.graph_analysis import get_support_axes
+
+    reduced_logps = []
+    for dependent_op, dependent_logp, dependent_dims_connection in zip(
+        dependent_ops, dependent_logps, dependent_dims_connections
+    ):
+        if dependent_logp.type.ndim > 0:
+            # Find which support axis implied by the MarginalRV need to be reduced
+            # Some may have already been reduced by the logp expression of the dependent RV (e.g., multivariate RVs)
+            dep_supp_axes = get_support_axes(dependent_op)[0]
+
+            # Dependent RV support axes are already collapsed in the logp, so we ignore them
+            supp_axes = [
+                -i
+                for i, dim in enumerate(reversed(dependent_dims_connection), start=1)
+                if (dim is None and -i not in dep_supp_axes)
+            ]
+            dependent_logp = dependent_logp.sum(supp_axes)
+
+            # Finally, we need to align the dependent logp batch dimensions with the marginalized logp
+            dims_alignment = [dim for dim in dependent_dims_connection if dim is not None]
+            dependent_logp = dependent_logp.transpose(*dims_alignment)
+
+        reduced_logps.append(dependent_logp)
+
+    reduced_logp = pt.add(*reduced_logps)
+    return reduced_logp
+
+
+def align_logp_dims(dims: tuple[tuple[int, None]], logp: TensorVariable) -> TensorVariable:
+    """Align the logp with the order specified in dims."""
+    dims_alignment = [dim for dim in dims if dim is not None]
+    return logp.transpose(*dims_alignment)
+
+
+def inline_ofg_outputs(op: OpFromGraph, inputs: Sequence[Variable]) -> tuple[Variable]:
+    """Inline the inner graph (outputs) of an OpFromGraph Op.
+
+    Whereas `OpFromGraph` "wraps" a graph inside a single Op, this function "unwraps"
+    the inner graph.
+    """
+    return clone_replace(
+        op.inner_outputs,
+        replace=tuple(zip(op.inner_inputs, inputs)),
+    )
+
+
+DUMMY_ZERO = pt.constant(0, name="dummy_zero")
+
+
+@_logprob.register(MarginalFiniteDiscreteRV)
+def finite_discrete_marginal_rv_logp(op: MarginalFiniteDiscreteRV, values, *inputs, **kwargs):
+    # Clone the inner RV graph of the Marginalized RV
+    marginalized_rv, *inner_rvs = inline_ofg_outputs(op, inputs)
+
+    # Obtain the joint_logp graph of the inner RV graph
+    inner_rv_values = dict(zip(inner_rvs, values))
+    marginalized_vv = marginalized_rv.clone()
+    rv_values = inner_rv_values | {marginalized_rv: marginalized_vv}
+    logps_dict = conditional_logp(rv_values=rv_values, **kwargs)
+
+    # Reduce logp dimensions corresponding to broadcasted variables
+    marginalized_logp = logps_dict.pop(marginalized_vv)
+    joint_logp = marginalized_logp + reduce_batch_dependent_logps(
+        dependent_dims_connections=op.dims_connections,
+        dependent_ops=[inner_rv.owner.op for inner_rv in inner_rvs],
+        dependent_logps=[logps_dict[value] for value in values],
+    )
+
+    # Compute the joint_logp for all possible n values of the marginalized RV. We assume
+    # each original dimension is independent so that it suffices to evaluate the graph
+    # n times, once with each possible value of the marginalized RV replicated across
+    # batched dimensions of the marginalized RV
+
+    # PyMC does not allow RVs in the logp graph, even if we are just using the shape
+    marginalized_rv_shape = constant_fold(tuple(marginalized_rv.shape), raise_not_constant=False)
+    marginalized_rv_domain = get_domain_of_finite_discrete_rv(marginalized_rv)
+    marginalized_rv_domain_tensor = pt.moveaxis(
+        pt.full(
+            (*marginalized_rv_shape, len(marginalized_rv_domain)),
+            marginalized_rv_domain,
+            dtype=marginalized_rv.dtype,
+        ),
+        -1,
+        0,
+    )
+
+    try:
+        joint_logps = vectorize_graph(
+            joint_logp, replace={marginalized_vv: marginalized_rv_domain_tensor}
+        )
+    except Exception:
+        # Fallback to Scan
+        def logp_fn(marginalized_rv_const, *non_sequences):
+            return graph_replace(joint_logp, replace={marginalized_vv: marginalized_rv_const})
+
+        joint_logps, _ = scan_map(
+            fn=logp_fn,
+            sequences=marginalized_rv_domain_tensor,
+            non_sequences=[*values, *inputs],
+            mode=Mode().including("local_remove_check_parameter"),
+        )
+
+    joint_logp = pt.logsumexp(joint_logps, axis=0)
+
+    # Align logp with non-collapsed batch dimensions of first RV
+    joint_logp = align_logp_dims(dims=op.dims_connections[0], logp=joint_logp)
+
+    # We have to add dummy logps for the remaining value variables, otherwise PyMC will raise
+    dummy_logps = (DUMMY_ZERO,) * (len(values) - 1)
+    return joint_logp, *dummy_logps
+
+
+@_logprob.register(MarginalDiscreteMarkovChainRV)
+def marginal_hmm_logp(op, values, *inputs, **kwargs):
+    chain_rv, *dependent_rvs = inline_ofg_outputs(op, inputs)
+
+    P, n_steps_, init_dist_, rng = chain_rv.owner.inputs
+    domain = pt.arange(P.shape[-1], dtype="int32")
+
+    # Construct logp in two steps
+    # Step 1: Compute the probability of the data ("emissions") under every possible state (vec_logp_emission)
+
+    # First we need to vectorize the conditional logp graph of the data, in case there are batch dimensions floating
+    # around. To do this, we need to break the dependency between chain and the init_dist_ random variable. Otherwise,
+    # PyMC will detect a random variable in the logp graph (init_dist_), that isn't relevant at this step.
+    chain_value = chain_rv.clone()
+    dependent_rvs = clone_replace(dependent_rvs, {chain_rv: chain_value})
+    logp_emissions_dict = conditional_logp(dict(zip(dependent_rvs, values)))
+
+    # Reduce and add the batch dims beyond the chain dimension
+    reduced_logp_emissions = reduce_batch_dependent_logps(
+        dependent_dims_connections=op.dims_connections,
+        dependent_ops=[dependent_rv.owner.op for dependent_rv in dependent_rvs],
+        dependent_logps=[logp_emissions_dict[value] for value in values],
+    )
+
+    # Add a batch dimension for the domain of the chain
+    chain_shape = constant_fold(tuple(chain_rv.shape))
+    batch_chain_value = pt.moveaxis(pt.full((*chain_shape, domain.size), domain), -1, 0)
+    batch_logp_emissions = vectorize_graph(reduced_logp_emissions, {chain_value: batch_chain_value})
+
+    # Step 2: Compute the transition probabilities
+    # This is the "forward algorithm", alpha_t = p(y | s_t) * sum_{s_{t-1}}(p(s_t | s_{t-1}) * alpha_{t-1})
+    # We do it entirely in logs, though.
+
+    # To compute the prior probabilities of each state, we evaluate the logp of the domain (all possible states)
+    # under the initial distribution. This is robust to everything the user can throw at it.
+    init_dist_value = init_dist_.type()
+    logp_init_dist = logp(init_dist_, init_dist_value)
+    # There is a degerate batch dim for lags=1 (the only supported case),
+    # that we have to work around, by expanding the batch value and then squeezing it out of the logp
+    batch_logp_init_dist = vectorize_graph(
+        logp_init_dist, {init_dist_value: batch_chain_value[:, None, ..., 0]}
+    ).squeeze(1)
+    log_alpha_init = batch_logp_init_dist + batch_logp_emissions[..., 0]
+
+    def step_alpha(logp_emission, log_alpha, log_P):
+        step_log_prob = pt.logsumexp(log_alpha[:, None] + log_P, axis=0)
+        return logp_emission + step_log_prob
+
+    P_bcast_dims = (len(chain_shape) - 1) - (P.type.ndim - 2)
+    log_P = pt.shape_padright(pt.log(P), P_bcast_dims)
+    log_alpha_seq, _ = scan(
+        step_alpha,
+        non_sequences=[log_P],
+        outputs_info=[log_alpha_init],
+        # Scan needs the time dimension first, and we already consumed the 1st logp computing the initial value
+        sequences=pt.moveaxis(batch_logp_emissions[..., 1:], -1, 0),
+    )
+    # Final logp is just the sum of the last scan state
+    joint_logp = pt.logsumexp(log_alpha_seq[-1], axis=0)
+
+    # Align logp with non-collapsed batch dimensions of first RV
+    remaining_dims_first_emission = list(op.dims_connections[0])
+    # The last dim of chain_rv was removed when computing the logp
+    remaining_dims_first_emission.remove(chain_rv.type.ndim - 1)
+    joint_logp = align_logp_dims(remaining_dims_first_emission, joint_logp)
+
+    # If there are multiple emission streams, we have to add dummy logps for the remaining value variables. The first
+    # return is the joint probability of everything together, but PyMC still expects one logp for each emission stream.
+    dummy_logps = (DUMMY_ZERO,) * (len(values) - 1)
+    return joint_logp, *dummy_logps

pymc_extras/model/marginal/graph_analysis.py

@@ -0,0 +1,372 @@
+import itertools
+
+from collections.abc import Sequence
+from itertools import zip_longest
+
+from pymc import SymbolicRandomVariable
+from pytensor.compile import SharedVariable
+from pytensor.graph import Constant, Variable, ancestors
+from pytensor.graph.basic import io_toposort
+from pytensor.tensor import TensorType, TensorVariable
+from pytensor.tensor.blockwise import Blockwise
+from pytensor.tensor.elemwise import CAReduce, DimShuffle, Elemwise
+from pytensor.tensor.random.op import RandomVariable
+from pytensor.tensor.rewriting.subtensor import is_full_slice
+from pytensor.tensor.shape import Shape
+from pytensor.tensor.subtensor import AdvancedSubtensor, Subtensor, get_idx_list
+from pytensor.tensor.type_other import NoneTypeT
+
+from pymc_extras.model.marginal.distributions import MarginalRV
+
+
+def static_shape_ancestors(vars):
+    """Identify ancestors Shape Ops of static shapes (therefore constant in a valid graph)."""
+    return [
+        var
+        for var in ancestors(vars)
+        if (
+            var.owner
+            and isinstance(var.owner.op, Shape)
+            # All static dims lengths of Shape input are known
+            and None not in var.owner.inputs[0].type.shape
+        )
+    ]
+
+
+def find_conditional_input_rvs(output_rvs, all_rvs):
+    """Find conditionally indepedent input RVs."""
+    blockers = [other_rv for other_rv in all_rvs if other_rv not in output_rvs]
+    blockers += static_shape_ancestors(tuple(all_rvs) + tuple(output_rvs))
+    return [
+        var
+        for var in ancestors(output_rvs, blockers=blockers)
+        if var in blockers or (var.owner is None and not isinstance(var, Constant | SharedVariable))
+    ]
+
+
+def is_conditional_dependent(
+    dependent_rv: TensorVariable, dependable_rv: TensorVariable, all_rvs
+) -> bool:
+    """Check if dependent_rv is conditionall dependent on dependable_rv,
+    given all conditionally independent all_rvs"""
+
+    return dependable_rv in find_conditional_input_rvs((dependent_rv,), all_rvs)
+
+
+def find_conditional_dependent_rvs(dependable_rv, all_rvs):
+    """Find rvs than depend on dependable"""
+    return [
+        rv
+        for rv in all_rvs
+        if (rv is not dependable_rv and is_conditional_dependent(rv, dependable_rv, all_rvs))
+    ]
+
+
+def get_support_axes(op) -> tuple[tuple[int, ...], ...]:
+    if isinstance(op, MarginalRV):
+        return op.support_axes
+    else:
+        # For vanilla RVs, the support axes are the last ndim_supp
+        return (tuple(range(-op.ndim_supp, 0)),)
+
+
+def _advanced_indexing_axis_and_ndim(idxs) -> tuple[int, int]:
+    """Find the output axis and dimensionality of the advanced indexing group (i.e., array indexing).
+
+    There is a special case: when there are non-consecutive advanced indexing groups, the advanced indexing
+    group is always moved to the front.
+
+    See: https://numpy.org/doc/stable/user/basics.indexing.html#combining-advanced-and-basic-indexing
+    """
+    adv_group_axis = None
+    simple_group_after_adv = False
+    for axis, idx in enumerate(idxs):
+        if isinstance(idx.type, TensorType):
+            if simple_group_after_adv:
+                # Special non-consecutive case
+                adv_group_axis = 0
+                break
+            elif adv_group_axis is None:
+                adv_group_axis = axis
+        elif adv_group_axis is not None:
+            # Special non-consecutive case
+            simple_group_after_adv = True
+
+    adv_group_ndim = max(idx.type.ndim for idx in idxs if isinstance(idx.type, TensorType))
+    return adv_group_axis, adv_group_ndim
+
+
+DIMS = tuple[int | None, ...]
+VAR_DIMS = dict[Variable, DIMS]
+
+
+def _broadcast_dims(
+    inputs_dims: Sequence[DIMS],
+) -> DIMS:
+    output_ndim = max((len(input_dim) for input_dim in inputs_dims), default=0)
+
+    # Add missing dims
+    inputs_dims = [
+        (None,) * (output_ndim - len(input_dim)) + input_dim for input_dim in inputs_dims
+    ]
+
+    # Find which known dims show in the output, while checking no mixing
+    output_dims = []
+    for inputs_dim in zip(*inputs_dims):
+        output_dim = None
+        for input_dim in inputs_dim:
+            if input_dim is None:
+                continue
+            if output_dim is not None and output_dim != input_dim:
+                raise ValueError("Different known dimensions mixed via broadcasting")
+            output_dim = input_dim
+        output_dims.append(output_dim)
+
+    # Check for duplicates
+    known_dims = [dim for dim in output_dims if dim is not None]
+    if len(known_dims) > len(set(known_dims)):
+        raise ValueError("Same known dimension used in different axis after broadcasting")
+
+    return tuple(output_dims)
+
+
+def _subgraph_batch_dim_connection(var_dims: VAR_DIMS, input_vars, output_vars) -> VAR_DIMS:
+    for node in io_toposort(input_vars, output_vars):
+        inputs_dims = [
+            var_dims.get(inp, ((None,) * inp.type.ndim) if hasattr(inp.type, "ndim") else ())
+            for inp in node.inputs
+        ]
+
+        if all(dim is None for input_dims in inputs_dims for dim in input_dims):
+            # None of the inputs are related to the batch_axes of the input_vars
+            continue
+
+        elif isinstance(node.op, DimShuffle):
+            [input_dims] = inputs_dims
+            output_dims = tuple(None if i == "x" else input_dims[i] for i in node.op.new_order)
+            var_dims[node.outputs[0]] = output_dims
+
+        elif isinstance(node.op, MarginalRV) or (
+            isinstance(node.op, SymbolicRandomVariable) and node.op.extended_signature is None
+        ):
+            # MarginalRV and SymbolicRandomVariables without signature are a wild-card,
+            # so we need to introspect the inner graph.
+            op = node.op
+            inner_inputs = op.inner_inputs
+            inner_outputs = op.inner_outputs
+
+            inner_var_dims = _subgraph_batch_dim_connection(
+                dict(zip(inner_inputs, inputs_dims)), inner_inputs, inner_outputs
+            )
+
+            support_axes = iter(get_support_axes(op))
+            if isinstance(op, MarginalRV):
+                # The first output is the marginalized variable for which we don't compute support axes
+                support_axes = itertools.chain(((),), support_axes)
+            for i, (out, inner_out) in enumerate(zip(node.outputs, inner_outputs)):
+                if not isinstance(out.type, TensorType):
+                    continue
+                support_axes_out = next(support_axes)
+
+                if inner_out in inner_var_dims:
+                    out_dims = inner_var_dims[inner_out]
+                    if any(
+                        dim is not None for dim in (out_dims[axis] for axis in support_axes_out)
+                    ):
+                        raise ValueError(f"Known dim corresponds to core dimension of {node.op}")
+                    var_dims[out] = out_dims
+
+        elif isinstance(node.op, Elemwise | Blockwise | RandomVariable | SymbolicRandomVariable):
+            # NOTE: User-provided CustomDist may not respect core dimensions on the left.
+
+            if isinstance(node.op, Elemwise):
+                op_batch_ndim = node.outputs[0].type.ndim
+            else:
+                op_batch_ndim = node.op.batch_ndim(node)
+
+            if isinstance(node.op, SymbolicRandomVariable):
+                # SymbolicRandomVariable don't have explicit expand_dims unlike the other Ops considered in this
+                [_, _, param_idxs], _ = node.op.get_input_output_type_idxs(
+                    node.op.extended_signature
+                )
+                for param_idx, param_core_ndim in zip(param_idxs, node.op.ndims_params):
+                    param_dims = inputs_dims[param_idx]
+                    missing_ndim = op_batch_ndim - (len(param_dims) - param_core_ndim)
+                    inputs_dims[param_idx] = (None,) * missing_ndim + param_dims
+
+            if any(
+                dim is not None for input_dim in inputs_dims for dim in input_dim[op_batch_ndim:]
+            ):
+                raise ValueError(
+                    f"Use of known dimensions as core dimensions of op {node.op} not supported."
+                )
+
+            batch_dims = _broadcast_dims(
+                tuple(input_dims[:op_batch_ndim] for input_dims in inputs_dims)
+            )
+            for out in node.outputs:
+                if isinstance(out.type, TensorType):
+                    core_ndim = out.type.ndim - op_batch_ndim
+                    output_dims = batch_dims + (None,) * core_ndim
+                    var_dims[out] = output_dims
+
+        elif isinstance(node.op, CAReduce):
+            [input_dims] = inputs_dims
+
+            axes = node.op.axis
+            if isinstance(axes, int):
+                axes = (axes,)
+            elif axes is None:
+                axes = tuple(range(node.inputs[0].type.ndim))
+
+            if any(input_dims[axis] for axis in axes):
+                raise ValueError(
+                    f"Use of known dimensions as reduced dimensions of op {node.op} not supported."
+                )
+
+            output_dims = [dims for i, dims in enumerate(input_dims) if i not in axes]
+            var_dims[node.outputs[0]] = tuple(output_dims)
+
+        elif isinstance(node.op, Subtensor):
+            value_dims, *keys_dims = inputs_dims
+            # Dims in basic indexing must belong to the value variable, since indexing keys are always scalar
+            assert not any(dim is None for dim in keys_dims)
+            keys = get_idx_list(node.inputs, node.op.idx_list)
+
+            output_dims = []
+            for value_dim, idx in zip_longest(value_dims, keys, fillvalue=slice(None)):
+                if idx == slice(None):
+                    # Dim is kept
+                    output_dims.append(value_dim)
+                elif value_dim is not None:
+                    raise ValueError(
+                        "Partial slicing or indexing of known dimensions not supported."
+                    )
+                elif isinstance(idx, slice):
+                    # Unknown dimensions kept by partial slice.
+                    output_dims.append(None)
+
+            var_dims[node.outputs[0]] = tuple(output_dims)
+
+        elif isinstance(node.op, AdvancedSubtensor):
+            # AdvancedSubtensor dimensions can show up as both the indexed variable and indexing variables
+            value, *keys = node.inputs
+            value_dims, *keys_dims = inputs_dims
+
+            # Just to stay sane, we forbid any boolean indexing...
+            if any(isinstance(idx.type, TensorType) and idx.type.dtype == "bool" for idx in keys):
+                raise NotImplementedError(
+                    f"Array indexing with boolean variables in node {node} not supported."
+                )
+
+            if any(dim is not None for dim in value_dims) and any(
+                dim is not None for key_dims in keys_dims for dim in key_dims
+            ):
+                # Both indexed variable and indexing variables have known dimensions
+                # I am to lazy to think through these, so we raise for now.
+                raise NotImplementedError(
+                    f"Simultaneous use of known dimensions in indexed and indexing variables in node {node} not supported."
+                )
+
+            adv_group_axis, adv_group_ndim = _advanced_indexing_axis_and_ndim(keys)
+
+            if any(dim is not None for dim in value_dims):
+                # Indexed variable has known dimensions
+
+                if any(isinstance(idx.type, NoneTypeT) for idx in keys):
+                    # Corresponds to an expand_dims, for now not supported
+                    raise NotImplementedError(
+                        f"Advanced indexing in node {node} which introduces new axis is not supported."
+                    )
+
+                non_adv_dims = []
+                for value_dim, idx in zip_longest(value_dims, keys, fillvalue=slice(None)):
+                    if is_full_slice(idx):
+                        non_adv_dims.append(value_dim)
+                    elif value_dim is not None:
+                        # We are trying to partially slice or index a known dimension
+                        raise ValueError(
+                            "Partial slicing or advanced integer indexing of known dimensions not supported."
+                        )
+                    elif isinstance(idx, slice):
+                        # Unknown dimensions kept by partial slice.
+                        non_adv_dims.append(None)
+
+                # Insert unknown dimensions corresponding to advanced indexing
+                output_dims = tuple(
+                    non_adv_dims[:adv_group_axis]
+                    + [None] * adv_group_ndim
+                    + non_adv_dims[adv_group_axis:]
+                )
+
+            else:
+                # Indexing keys have known dimensions.
+                # Only array indices can have dimensions, the rest are just slices or newaxis
+
+                # Advanced indexing variables broadcast together, so we apply same rules as in Elemwise
+                adv_dims = _broadcast_dims(keys_dims)
+
+                start_non_adv_dims = (None,) * adv_group_axis
+                end_non_adv_dims = (None,) * (
+                    node.outputs[0].type.ndim - adv_group_axis - adv_group_ndim
+                )
+                output_dims = start_non_adv_dims + adv_dims + end_non_adv_dims
+
+            var_dims[node.outputs[0]] = output_dims
+
+        else:
+            raise NotImplementedError(f"Marginalization through operation {node} not supported.")
+
+    return var_dims
+
+
+def subgraph_batch_dim_connection(input_var, output_vars) -> list[DIMS]:
+    """Identify how the batch dims of input map to the batch dimensions of the output_rvs.
+
+    Example:
+    -------
+    In the example below `idx` has two batch dimensions (indexed 0, 1 from left to right).
+    The two uncommented dependent variables each have 2 batch dimensions where each entry
+    results from a mapping of a single entry from one of these batch dimensions.
+
+    This mapping is transposed in the case of the first dependent variable, and shows up in
+    the same order for the second dependent variable. Each of the variables as a further
+    batch dimension encoded as `None`.
+
+    The commented out third dependent variable combines information from the batch dimensions
+    of `idx` via the `sum` operation. A `ValueError` would be raised if we requested the
+    connection of batch dims.
+
+    .. code-block:: python
+        import pymc as pm
+
+        idx = pm.Bernoulli.dist(shape=(3, 2))
+        dep1 = pm.Normal.dist(mu=idx.T[..., None] * 2, shape=(3, 2, 5))
+        dep2 = pm.Normal.dist(mu=idx * 2, shape=(7, 2, 3))
+        # dep3 = pm.Normal.dist(mu=idx.sum())  # Would raise if requested
+
+        print(subgraph_batch_dim_connection(idx, [], [dep1, dep2]))
+        # [(1, 0, None), (None, 0, 1)]
+
+    Returns:
+    -------
+    list of tuples
+        Each tuple corresponds to the batch dimensions of the output_rv in the order they are found in the output.
+        None is used to indicate a batch dimension that is not mapped from the input.
+
+    Raises:
+    ------
+    ValueError
+        If input batch dimensions are mixed in the graph leading to output_vars.
+
+    NotImplementedError
+        If variable related to marginalized batch_dims is used in an operation that is not yet supported
+    """
+    var_dims = {input_var: tuple(range(input_var.type.ndim))}
+    var_dims = _subgraph_batch_dim_connection(var_dims, [input_var], output_vars)
+    ret = []
+    for output_var in output_vars:
+        output_dims = var_dims.get(output_var, (None,) * output_var.type.ndim)
+        assert len(output_dims) == output_var.type.ndim
+        ret.append(output_dims)
+    return ret