How to sample N-dimensional discrete data with K classes per dimension?

AlphaBetaGamma96: <p>How can I create a <code>torch.distribution</code> class that samples an N-dimensional discrete space of <code>K</code> classes each?</p>
<p>I initially thought of using <code>multinomial</code> (docs: <a href="https://pytorch.org/docs/stable/distributions.html#multinomial" rel="noopener nofollow ugc">here</a>), but that would require initializing an array, which scales exponentially with respect to the number of dimensions (i.e. the sampling array would be <code>K^N</code> in size).</p>
<p>If I fix this distribution to be equal in probability per class per dimension, is there an easy way to generate samples (as well as the probability, which should just be <code>(1/K)^N</code> for uniformly samples discrete values of <code>K</code> classes in <code>N</code> dimensions)</p>
<p>Any help is welcome!</p>

KFrank: <p>Hi Alpha!</p>
<aside class="quote no-group" data-username="AlphaBetaGamma96" data-post="1" data-topic="204172" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/a/3da27b/48.png" class="avatar"> AlphaBetaGamma96:</div>
<blockquote>
<p>If I fix this distribution to be equal in probability per class per dimension, is there an easy way to generate samples</p>
</blockquote>
</aside>
<p>If I understand what you are looking for, I think <code>multinomial()</code> will<br>
work just fine for you.  As I understand it, you are not looking for any<br>
kind of correlation between class values in different dimensions, so<br>
if this is the case, you should be able to “factor” the random sampling<br>
across dimensions.</p>
<p>Here is an example:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; print (torch.__version__)
2.3.0
&gt;&gt;&gt;
&gt;&gt;&gt; _ = torch.manual_seed (2024)
&gt;&gt;&gt;
&gt;&gt;&gt; k = 5
&gt;&gt;&gt; n = 10
&gt;&gt;&gt;
&gt;&gt;&gt; probs = (torch.ones (1, k) / k).expand (n, k)   # uniform class probabilities for all dimensions
&gt;&gt;&gt;
&gt;&gt;&gt; num_samples = 2
&gt;&gt;&gt; torch.multinomial (probs, num_samples = num_samples, replacement = True)
tensor([[1, 3],
        [1, 3],
        [4, 1],
        [1, 2],
        [2, 0],
        [0, 1],
        [4, 3],
        [3, 1],
        [4, 2],
        [2, 4]])
</code></pre>
<p>Each column in the result tensor is a sample – in this example, there<br>
are two samples.  As you note, there are <code>k**n</code> distinct (vector-valued)<br>
values for a sample and when feeding uniform probabilities to<br>
<code>multinomial()</code>, each such value will be sampled with probability<br>
<code>1 / k**n</code>.</p>
<p>Note that this illustration does not depend on the probabilities being<br>
uniform across classes or dimensions.  For arbitrary probabilities, you<br>
do need to materialize the full <code>[n, k]</code> probability matrix, but that scales<br>
as <code>n*k</code>, rather than exponentially.</p>
<p>If this is not what you are looking for, could you post some concrete<br>
code that generates your desired result (even if it doesn’t scale well)?</p>
<p>Best.</p>
<p>K. Frank</p>

AlphaBetaGamma96: <p>Hi <a class="mention" href="/u/kfrank">@KFrank</a>,</p>
<p>This looks like what I was looking for. My idea case would be to sample an arbitrary N-dim discrete data distribution (via some kind of Normalizing Flow perhaps), but I need a base distribution like the one you’ve provided.</p>
<p>But the only issue is that Normalizing Flows are designed for continuous distributions (rather than discrete distributions, so I’m not sure if it’s feasible to use this as a base distribution for a flow).</p>