How to sample from Multinomial a variable number of times?

drj: <p>I’m trying to use a <code>torch.distributions.multinomial.Multinomial</code> distribution for a reinforcement learning task. Essentially, my network outputs probabilities over actions, which I want to feed into a <code>Multinomial</code> and then sample from.</p>
<p>My network output looks something like <code>(locations x objects)</code>, where each of 5 different <code>object</code> can be placed in 1 of 10 different <code>locations</code>. This is all fine and good.</p>
<p>The problem I have is that in different rollouts, I might have a variable number of each <code>object</code>. So in addition to my <code>(10 x 5)</code> <code>probs</code> matrix, I also have a <code>(5,)</code> <code>objects</code> vector that’s got how many of each object I have in that scenario. It might look like <code>[2, 3, 1, 0, 6]</code> where I have 2 of the first object, 3 of the second, etc. In this particular rollout, I’d want to sample from my first distribution over <code>locations</code> 2 times, since I have 2 of that <code>object</code>.</p>
<p>In the <code>batch_size = 1</code> scenario, this works great. I do <code>torch.unbind(probs, dim=-1)</code>, and then make a list of <code>Multinomials</code>, with an individual <code>total_count</code> equal to the corresponding number of objects in my <code>objects</code> vector, like this:</p>
<p><code>cats = [torch.distributions.multinomial.Multinomial(total_count = obj, probs=prob) for obj, prob in zip(objects, torch.unbind(probs, dim=-1))]</code></p>
<p>Or, I can do something like <code>[cat.sample(n) for n, cat in zip(objects, cats)]</code>, if I instead decide to set <code>total_count</code> equal to 1 for each object distribution.</p>
<p>However I’m completely stumped on how to handle the situation where my <code>batch_size &gt; 1</code>. Like, I get an output from my network that’s say, <code>(32 x 10 x 5)</code>, and I can again make 5 different categories. But since my <code>objects</code> is now <code>(32 x 5)</code>, when I slice to say, <code>obejcts[:,0</code>] I get a slice with a variable number of <code>objects</code> in it, like: <code>[3, 5, 0, ..., 1]</code>. And as far as I can tell, I can’t fit a list or vector into either <code>total_counts</code> or <code>.sample()</code>.</p>
<p>So, how can I sample from a Multinomial a variable number of times? Ideally, I’d like to put in a <code>(32 x 10</code>) matrix for my <code>probs</code>, and then sample a number of times defined by a vector of size <code>(32, )</code>. I guess, even better, I’d like to be able to pass a <code>(32 x 10 x 5)</code> <code>probs</code> matrix, and then pass a sample matrix of size <code>(32 x 1 x 5)</code> that defines my samples, but I’d take either. Any ideas?</p>

AlphaBetaGamma96: <p>Hi <a class="mention" href="/u/drj">@drj</a>,</p>
<p>Could you share a minimal reproducible example of your issue?</p>
<aside class="quote no-group" data-username="drj" data-post="1" data-topic="202953">
<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/drj/48/61407_2.png" class="avatar"> drj:</div>
<blockquote>
<p>cats = [torch.distributions.multinomial.Multinomial(total_count = obj, probs=prob) for obj, prob in zip(objects, torch.unbind(probs, dim=-1))]</p>
</blockquote>
</aside>
<p>If you have a function that can generate a sample for a batch size of 1, you could try using <code>torch.func.vmap</code> to vectorize over this function to generate data with batch size &gt; 1. The docs for <code>torch.func.vmap</code> are <a href="https://pytorch.org/docs/stable/generated/torch.func.vmap.html" rel="noopener nofollow ugc">here</a>.</p>

drj: <p>Hi <a class="mention" href="/u/alphabetagamma96">@AlphaBetaGamma96</a></p>
<p>Thanks for the suggestion. I haven’t looked at <code>torch.func.vmap</code> yet, so maybe that will solve my problem.</p>
<p>Here’s a short script demonstrating what I’m trying to do:</p>
<pre><code class="lang-auto">import torch

# our goal is to place N vehicles in L locations.
# to simplify, lets assume we have 3 vehicle-types: car, buses, bikes
# we have 4 potential locations
# we get a probabily distribution like (Locations x Vehicles):
#   cars
# [[.08, .7, .25], &lt;- location0
#  [.9, .01, .25],
#  [.01, .25, .25],
#  [.01, .04, .25]]

# we also get a vector of vehicle counts like [1, 3, 0]
# so our output (Locations x Vehicles), after sampling, might look like:
# [[0, 2, 0],
#  [1, 0, 0],
#  [0, 1, 0],
#  [0, 0, 0]]

# batch_size = 1 case
probs1 = torch.arange(4 * 3).reshape(4, 3)
vehicles1 = [1, 3, 1]


def smpl(probs, vehicles):
    cats = [
        torch.distributions.multinomial.Multinomial(total_count=n, probs=prob)
        for n, prob in zip(vehicles, torch.unbind(probs, dim=-1))
    ]
    sample_ = torch.stack([cat.sample() for cat in cats], dim=-1)
    return sample_


# can also do
def smpl2(probs, vehicles):
    cats = [
        torch.distributions.multinomial.Multinomial(probs=prob)
        for prob in torch.unbind(probs, dim=-1)
    ]
    sample_ = torch.stack(
        [cat.sample((n,)).sum(dim=0) for n, cat in zip(vehicles, cats)], dim=-1
    )
    return sample_


# batch_size = 5 case
probs5 = torch.arange(5 * 4 * 3).reshape(5, 4, 3)
vehicles5 = torch.arange(5 * 3).reshape(5, 3)


# now the first approach doesn't work  because total_counts can only be an int
def smpl_batch(probs, vehicles):
    cats = [
        torch.distributions.multinomial.Multinomial(total_counts=n, probs=prob)
        for n, prob in zip(torch.unbind(vehicles, dim=-1), torch.unbind(probs, dim=-1))
    ]
    sample_ = torch.stack([cat.sample() for cat in cats], dim=-1)
    return sample_


# ERROR


# and the second approach might work if i could figure out the right slicing on something like:
def smpl2_batch(probs, vehicles):
    cats = [
        torch.distributions.multinomial.Multinomial(probs=prob)
        for prob in torch.unbind(probs, dim=-1)
    ]
    sample_ = torch.stack(
        [cat.sample((max(n),)) for n, cat in zip(torch.unbind(vehicles, dim=-1), cats)],
        dim=-1,
    )
    # do some crazy slicing/summing here on sample_ and vehicles
    return sample_


# but i can't figure out how the slicing/summing would work to get back to a (batch x locations x vehicles) size

if __name__ == "__main__":
    print("Batch Size: 1")
    print(f"Probs: {probs1}")
    print(f"Vehicles: {vehicles1}")
    print(f"Result first approach: {smpl(probs1, vehicles1)}\n")
    print(f"Result second approach: {smpl2(probs1, vehicles1)}\n")
    print("Batch Size: 5")
    print(f"Probs: {probs5}")
    print(f"Vehicles: {vehicles5}")
    print(f"Result: {smpl2_batch(probs5, vehicles5)}")
</code></pre>
<p>At the end of the day, I’d really like to be able to pass a 1D vector to <code>total_count</code> that matches the innermost dimension of <code>probs</code> or <code>logits</code>, representing <code>how many times to sample from that distribution</code>.</p>

drj: <p>For anyone finding this in the future, I figured out the slicing. This works for me:</p>
<pre><code class="lang-auto">def smpl2_batch(probs, vehicles):
    cats = [
        torch.distributions.multinomial.Multinomial(probs=prob)
        for prob in torch.unbind(probs, dim=-1)
    ]
    vehicles = torch.unbind(vehicles, dim=-1)
    arr = [cat.sample((max(v),)) for v, cat in zip(vehicles, cats)
    return torch.stack(
        [torch.sum(
            (torch.arange(arr.size(0)) &lt; v.unsqueeze(-1)).T.unsqueeze(-1)
            * arr,
            dim=1)
        for v, arr in zip(vehicles, arr)
        ],
        dim=-1,
    )
</code></pre>
<p>I’ll admit it’s pretty ugly, but it essentially samples as many times as we have vehicles, for each vehicle-type: <code>arr = [cat.sample((max(v),)) for v, cat in zip(vehicles, cats)</code>.</p>
<p>Then, it builds a mask for each vehicle-type: <code>torch.arange(arr.size(0))</code>, and sets all the samples to zero that are greater then the count of vehicles for that type: <code>v.unsqueeze(-1)</code>.</p>
<p>Then we manipulate that <code>mask</code> so it can broadcast to the <code>(samples x batch x locations)</code> matrix: <code>.T.unsqueeze(-1)</code>, multiply them, and <code>sum()</code> along the <code>samples</code> dimension to collapse it to shape <code>(batch x locations)</code>.</p>
<p>Repeat for each vehicle-type, stack the matrices along the last axis to get <code>(batch x locations x vehicles)</code>.</p>