Thank you for answering. To be honest, I don’t understand your answer, but I think, maybe I have not been clear enough:
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)



with d1 and d2 being arbirtrary torch.distribution.Distribution.