.cat() / .stack() distributions (e.g. Normal distribution)

username: tjoseph
cooked: <p>Hey <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>
<p>Is there a way to stack / cat torch.distributions?</p>
<p>Example:</p>
<pre><code class="lang-auto">import torch
from torch.distributions import Normal

mean1= torch.zeros((5), dtype=torch.float)
std1 = torch.ones((5), dtype=torch.float)
d1 = Normal(mean1, std1)

mean2= torch.ones((5), dtype=torch.float) * 10
std2 = torch.ones((5), dtype=torch.float) * 10
d2 = Normal(mean2, std2)

mean_result = torch.stack([mean1, mean2])
std_result = torch.stack([std1, std2])
d_result = Normal(mean_result, std_result)
</code></pre>
<p>Is there a way to get <code>d_result</code> from <code>d1</code> and <code>d2</code> without using/stacking the parameters (in this case mean/std)? I am looking for a way that is agnostic to the type of distributions!</p>

username: tjoseph
cooked: <p>I am still wondering <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

username: tjoseph
cooked: <p>There seems to be still no solution.</p>

username: AlphaBetaGamma96
cooked: <p>Normal distribution only works for 1-dimensional distributions, what you need is MultivariateNormal which is located here <a href="https://pytorch.org/docs/stable/distributions.html?highlight=torch%20distributions%20multivariatenormal#torch.distributions.multivariate_normal.MultivariateNormal" class="inline-onebox" rel="noopener nofollow ugc">Probability distributions - torch.distributions — PyTorch 1.10.0 documentation</a></p>

username: tjoseph
cooked: <p>Sorry, maybe I should have made this more clearer, but the question is not about Normal distributions, but about torch.distributions.Distribution in general.</p>

username: tom
cooked: <p>So you are looking for a “product distribution” (i.e. variables independently sampled from their respective distributions).<br>
I don’t think this is currently implemented, but it should be possible to implement such a distribution without too much effort. Of course, the next person will want the distributions of the components linked by given copulas instead of being independent, but hey.</p>
<p>Best regards</p>
<p>Thomas</p>

username: tjoseph
cooked: <p>Thank you for answering. To be honest, I don’t understand your answer, but I think, maybe I have not been clear enough:</p>
<pre><code class="lang-auto">import torch
from torch.distributions import Normal

mean1= torch.zeros((5), dtype=torch.float)
std1 = torch.ones((5), dtype=torch.float)
d1 = Normal(mean1, std1)

mean2= torch.ones((5), dtype=torch.float) * 10
std2 = torch.ones((5), dtype=torch.float) * 10
d2 = Normal(mean2, std2)

# Instead of this
mean_result = torch.stack([mean1, mean2])
std_result = torch.stack([std1, std2])
d_result = Normal(mean_result, std_result)

# I want something like this
d_result = torch.stack_along_a_batch_dimension([d1, d2])

# CURRENT use case
event1 = torch.tensor(...)
event2 = torch.tensor(...)

# One CUDA kernel call per distribution-event-pair
log_prob_sum = d1.log_prob(event1) + d2.log_prob(event2)

# WISH: How I wished I could do it
event = torch.stack([event1, event2])
# Only a single CUDA kernel call
log_prob_sum = d_result.log_prob(event)


</code></pre>
<p>with <code>d1</code> and <code>d2</code> being arbirtrary <code>torch.distribution.Distribution</code>.</p>

username: tom
cooked: <p>But so what you describe sounds as if it were equivalent to using d1 and d2 separately, right?<br>
So what you are doing mathematically is to define a distribution on the product space of the domains of the two distributions. (So e.g. sample would sample from d1 and d2 and return the concatenated vector. log_prob would split the vector, take the log_prob of d1 and d2 and return the sum.)</p>
<p>Best regards</p>
<p>Thomas</p>

username: googlebot
cooked: <aside class="quote no-group" data-username="tjoseph" data-post="7" data-topic="92567">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/tjoseph/48/11927_2.png" class="avatar"> tjoseph:</div>
<blockquote>
<p><code># Only a single CUDA kernel call</code></p>
</blockquote>
</aside>
<p>That’s funny, as there are like 5-10 cuda calls, because all log_prob implementations are in Python (so one math op is done at a time). Well, unless JIT is used, but it is cumbersome to use on Distribution objects.</p>
<p>More ontopic, I think distributions.Independent exists for this usecase.</p>

username: tjoseph
cooked: <p>Hi,<br>
thank you for asking. What I want to accomplish has nothing to do with distributions as such. Similar as I can for example input a batch of images in a CNN via looping over the images or via just concatenating the images in the batch dimension and inputing the whole batch, I want to concatenate my distributions and calculate the log prob. So coming back to the example I gave</p>
<aside class="quote no-group" data-username="tom" data-post="8" data-topic="92567">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/tom/48/3162_2.png" class="avatar"> tom:</div>
<blockquote>
<p>log_prob would split the vector, take the log_prob of d1 and d2 and return the sum.</p>
</blockquote>
</aside>
<p>This is wrong. I want two log_probs, one for each event, since I am only talking about concatenation in the batch dimension!</p>

username: tjoseph
cooked: <p>Thank you for answering. I think you are talking about something different.</p>
<p>I have modules that return <code>torch.distribution.Distribution</code>. I do not wish to know which distribution, I just know it is always the same (e.g. always Normal, or always Categorical). I have a video stream and each of the <code>N</code> images image in the video stream is input into this module in a for loop, since the architecture is RNN-like. Then after I want calculate the log_prob for each distribution I got, so <code>N</code> times. Usually when I get tensors back instead of distributions, I would just stack the tensor and then do some processing. However, with distributions I do not see a way to stack them in the batch_dimension, so I have to loop over them again.</p>

username: googlebot
cooked: <p>Sounds like you just need to stack parameters and construct a merged distribution. But there is no interface to extract distribution parameters, except in an ad-hoc way:</p>
<p>{name: getattr(distr, name) for name in distr.arg_constraints.keys()}</p>
<p>stack these by key and construct a merged distr.object in a generic way, e.g. <code>distr[0].__class__(**kwargs)</code></p>

username: JosueCom
cooked: <p>Is there a way to merge the distributions without having to create them separately with sample() and then, stacking them? The assumption is that you want to have vectors with values that come from different distribution (not necessarily normal)</p>