Is it possible to further vectorize the operations of this colored point torch.bucketization?

username: kevindckr
cooked: <p>I have <code>n</code> points each with an associated color</p>
<pre><code class="lang-auto">n = 100
colors = torch.rand((n, 3))
points = torch.randint(0, 10, (n, 2))
x_min, x_max = torch.min(points[:, 0]).data, torch.max(points[:, 0]).data
y_min, y_max = torch.min(points[:, 1]).data, torch.max(points[:, 1]).data
</code></pre>
<p>I want to discretize the point’s colors onto some <code>target</code> grid</p>
<pre><code class="lang-auto">target = torch.zeros((5, 5, 3))
eps = torch.finfo(torch.float32).eps # Bucketizing from the "left"
x_steps = torch.linspace(x_min, x_max + eps, target.shape[0] + 1)
y_steps = torch.linspace(y_min, y_max + eps, target.shape[1] + 1)
</code></pre>
<p>I can determine which points fall within each grid location with <code>bucketize</code></p>
<pre><code class="lang-auto">x_bucket_indices = torch.bucketize(points[:, 0], x_steps)
y_bucket_indices = torch.bucketize(points[:, 1], y_steps)
</code></pre>
<p>Though the issue comes when attempting to update <code>target</code> based on the <code>(0, 1, n)</code> number of colors that may be present in each of its cells/buckets/elements.</p>
<pre><code class="lang-auto">for ix in range(target.shape[0]):
    row_inds = torch.where(x_bucket_indices == ix)[0]
    for iy in range(target.shape[1]):
        col_inds = torch.where(y_bucket_indices == iy)[0]
        cell_inds = intersection(row_inds, col_inds)

        # Collect the color of the points that fall within this cell
        target[ix, iy, :] = pool_color(colors[cell_inds, :])
</code></pre>
<p>(Utilizing these two helper functions)</p>
<pre><code class="lang-auto"># Define how to "pool" the colors in each cell of the grid(/bucket)
def pool_color(colors):
    if len(colors) == 0: # Interpolate? (Later)
        return torch.zeros((3,))
    if len(colors) == 1: # Single points single color
        return colors[0]
    if len(colors) &gt; 1:  # Average of all colors in cell
        return torch.mean(colors, dim=0)

# Define a helper function to find the intersection of unique vectors
def intersection(x, y):
    # x[1, 2, 3, 4], y[3, 4] -&gt; [3, 4]
    combined = torch.cat((x, y))
    unq, cnt = combined.unique(return_counts=True)
    return unq[cnt &gt; 1]
</code></pre>
<hr>
<p><strong>Is there a way to further vectorize the operations within the <code>for</code> loops such that they are performed on the GPU?</strong></p>
<p>I think one of the biggest issues is not having access to something akin to <code>map</code> in PyTorch and the generation of a potentially “very” ragged/jagged tensor depending how the <code>pool_color</code> is applied.</p>
<p>Though the problem itself seems very parallelizable. Provided a set of indices, pull from shared memory the colors, perform <code>pool_color</code>, and write the result to shared memory in target. I’m just not sure how to implement this any further in torch.</p>

username: tom
cooked: <p>If you compute a global bucket index (by combining x and y) you can then use the third-party <a href="https://pytorch-scatter.readthedocs.io/en/latest/functions/scatter.html" rel="noopener nofollow ugc">PyTorch Scatter</a> library.</p>
<p>Best regards</p>
<p>Thomas</p>

username: kevindckr
cooked: <p>Thanks Thomas. I think that may work for my case. I’ve played around with a bit this morning. Just struggling a bit with generalizing it to the 2D case.</p>
<p>I think one of the issues may be, but I am not yet certain, that the <code>src</code> tensor in the <code>scatter</code> operation may only be allowed to be <code>(n,)</code> in shape without any more axes. For example, how would the operation know on what dimension to average over if a tensor of <code>(n, q, r, s)</code> was provided?</p>
<p>I’ll keep looking into it, but just an fyi for others who may have suggestions or input I have not yet solved the problem.</p>