Model criticism API (ed.copy): possible to copy multiple model components at once?

For the model criticism API it could be useful to be able to copy multiple model components at once.

Consider this toy example:

q = Normal(...)

a = Normal(f(q))

b = Normal(g(q))

q are your latent variables and a and b are observed (g and f are some functions). Now say you do inference with a variational model q tied to qz ({q: qz}). Finally you want to criticize the model.

a_post = ed.copy(a, {q: qz})
b_post = ed.copy(b, {q: qz})

It could be nice to allow a list of random variables or tensors as input to ed.copy() so that nodes that are shared in the computational graph say downstream from qz do not have to be copied twice. (e.g. f and g could be large neural networks which share parameters)

That makes sense. It’s an easy change to the internals in ed.copy to handle lists.