# Copyright (c) 2017-2019 Uber Technologies, Inc.
# SPDX-License-Identifier: Apache-2.0

import numbers
import warnings

import torch

from .messenger import Messenger
from .runtime import apply_stack


class DoMessenger(Messenger):
    """
    Given a stochastic function with some sample statements
    and a dictionary of values at names,
    set the return values of those sites equal to the values
    as if they were hard-coded to those values
    and introduce fresh sample sites with the same names
    whose values do not propagate.

    Composes freely with :func:`~pyro.poutine.handlers.condition`
    to represent counterfactual distributions over potential outcomes.
    See Single World Intervention Graphs [1] for additional details and theory.

    Consider the following Pyro program:

        >>> def model(x):
        ...     s = pyro.param("s", torch.tensor(0.5))
        ...     z = pyro.sample("z", dist.Normal(x, s))
        ...     return z ** 2

    To intervene with a value for site `z`, we can write

        >>> intervened_model = pyro.poutine.do(model, data={"z": torch.tensor(1.)})

    This is equivalent to replacing `z = pyro.sample("z", ...)` with
    `z = torch.tensor(1.)`
    and introducing a fresh sample site pyro.sample("z", ...) whose value is not used elsewhere.

    References

    [1] `Single World Intervention Graphs: A Primer`,
        Thomas Richardson, James Robins

    :param fn: a stochastic function (callable containing Pyro primitive calls)
    :param data: a ``dict`` mapping sample site names to interventions
    :returns: stochastic function decorated with a :class:`~pyro.poutine.do_messenger.DoMessenger`
    """
    def __init__(self, data):
        super().__init__()
        self.data = data
        self._intervener_id = str(id(self))

    def _pyro_sample(self, msg):
        if msg.get('_intervener_id', None) != self._intervener_id and \
                self.data.get(msg['name']) is not None:

            if msg.get('_intervener_id', None) is not None:
                warnings.warn(
                    "Attempting to intervene on variable {} multiple times,"
                    "this is almost certainly incorrect behavior".format(msg['name']),
                    RuntimeWarning)

            msg['_intervener_id'] = self._intervener_id

            # split node, avoid reapplying self recursively to new node
            new_msg = msg.copy()
            apply_stack(new_msg)

            # apply intervention
            intervention = self.data[msg['name']]
            msg['name'] = msg['name'] + "__CF"  # mangle old name

            if isinstance(intervention, (numbers.Number, torch.Tensor)):
                msg['value'] = intervention
                msg['is_observed'] = True
                msg['stop'] = True
            else:
                raise NotImplementedError(
                    "Interventions of type {} not implemented (yet)".format(type(intervention)))

        return None

pyro/poutine/do_messenger.py
# Copyright (c) 2017-2019 Uber Technologies, Inc.
# SPDX-License-Identifier: Apache-2.0

import numbers
import warnings

import torch

from .messenger import Messenger
from .runtime import apply_stack


class DoMessenger(Messenger):
    """
    Given a stochastic function with some sample statements
    and a dictionary of values at names,
    set the return values of those sites equal to the values
    as if they were hard-coded to those values
    and introduce fresh sample sites with the same names
    whose values do not propagate.

    Composes freely with :func:`~pyro.poutine.handlers.condition`
    to represent counterfactual distributions over potential outcomes.
    See Single World Intervention Graphs [1] for additional details and theory.

    Consider the following Pyro program:

        >>> def model(x):
        ...     s = pyro.param("s", torch.tensor(0.5))
        ...     z = pyro.sample("z", dist.Normal(x, s))
        ...     return z ** 2

    To intervene with a value for site `z`, we can write

        >>> intervened_model = pyro.poutine.do(model, data={"z": torch.tensor(1.)})

    This is equivalent to replacing `z = pyro.sample("z", ...)` with
    `z = torch.tensor(1.)`
    and introducing a fresh sample site pyro.sample("z", ...) whose value is not used elsewhere.

    References

    [1] `Single World Intervention Graphs: A Primer`,
        Thomas Richardson, James Robins

    :param fn: a stochastic function (callable containing Pyro primitive calls)
    :param data: a ``dict`` mapping sample site names to interventions
    :returns: stochastic function decorated with a :class:`~pyro.poutine.do_messenger.DoMessenger`
    """
    def __init__(self, data):
        super().__init__()
        self.data = data
        self._intervener_id = str(id(self))

    def _pyro_sample(self, msg):
        if msg.get('_intervener_id', None) != self._intervener_id and \
                self.data.get(msg['name']) is not None:

            if msg.get('_intervener_id', None) is not None:
                warnings.warn(
                    "Attempting to intervene on variable {} multiple times,"
                    "this is almost certainly incorrect behavior".format(msg['name']),
                    RuntimeWarning)

            msg['_intervener_id'] = self._intervener_id

            # split node, avoid reapplying self recursively to new node
            new_msg = msg.copy()
            new_msg["cond_indep_stack"] = ()  # avoid entering plates twice
            apply_stack(new_msg)

            # apply intervention
            intervention = self.data[msg['name']]
            msg['name'] = msg['name'] + "__CF"  # mangle old name

            if isinstance(intervention, (numbers.Number, torch.Tensor)):
                msg['value'] = intervention
                msg['is_observed'] = True
                msg['stop'] = True
            else:
                raise NotImplementedError(
                    "Interventions of type {} not implemented (yet)".format(type(intervention)))

        return None

pyro/poutine/do_messenger.py
diff --git a/pyro/poutine/do_messenger.py b/pyro/poutine/do_messenger.py
index 8ddda1fa81..1b053f0f17 100644
--- a/pyro/poutine/do_messenger.py
+++ b/pyro/poutine/do_messenger.py
@@ -66,6 +66,7 @@ def _pyro_sample(self, msg):
 
             # split node, avoid reapplying self recursively to new node
             new_msg = msg.copy()
+            new_msg["cond_indep_stack"] = ()  # avoid entering plates twice
             apply_stack(new_msg)
 
             # apply intervention
diff --git a/tests/poutine/test_counterfactual.py b/tests/poutine/test_counterfactual.py
index a8f9a06617..286c771eba 100644
--- a/tests/poutine/test_counterfactual.py
+++ b/tests/poutine/test_counterfactual.py
@@ -73,3 +73,27 @@ def model():
                 assert_equal(interventions[name], actual_values[name])
             if interventions[name] != observations[name]:
                 assert_not_equal(interventions[name], tr.nodes[name]['value'])
+
+
+def test_plate_duplication_smoke():
+
+    def model(N):
+
+        with pyro.plate("x_plate", N):
+            z1 = pyro.sample("z1", dist.MultivariateNormal(torch.zeros(2), torch.eye(2)))
+            z2 = pyro.sample("z2", dist.MultivariateNormal(torch.zeros(2), torch.eye(2)))
+            return pyro.sample("x", dist.MultivariateNormal(z1+z2, torch.eye(2)))
+
+    fix_z1 = torch.tensor([[-6.1258, -6.1524],
+                           [-4.1513, -4.3080]])
+
+    obs_x = torch.tensor([[-6.1258, -6.1524],
+                          [-4.1513, -4.3080]])
+
+    do_model = poutine.do(model, data={"z1": fix_z1})
+    do_model = poutine.condition(do_model, data={"x": obs_x})
+    do_auto = pyro.infer.autoguide.AutoMultivariateNormal(do_model)
+    optim = pyro.optim.Adam({"lr": 0.05})
+
+    svi = pyro.infer.SVI(do_model, do_auto, optim, pyro.infer.Trace_ELBO())
+    svi.step(len(obs_x))

[bug] Runtime error during SVI inference when using poutine.do()
### Issue Description

Setting: a simple model with 2 latent Gaussians z1 and z2, giving rise to x ~ N( z1+z2, I).

In this setting p(z2 | x, z1) should be the same as p(z2 | x, do(z1)). 

I wanted to check whether the current Pyro interface reflects this and it seems it does not.

My initial thought is that there is a difference in how .do() and .condition() broadcast the constants across the plate context.

### Environment

 - OS and python version: MacOS 10.14.6, Python: 3.8.6
 - PyTorch version: 1.9.0.dev20210502 (nightly version)
 - Pyro version: 1.6.0.

### Code Snippet

Replication code:
https://pastebin.com/Ki2PYX7z

