How to Create Upper Triangular Matrix with Diagonals of Increasing Power?

RylanSchaeffer: <p>I’d like to create a square upper triangular matrix that is defined as follows for some float <code>c</code> and some dimension <code>N</code>:</p>
<pre><code class="lang-auto">[[1 , c , c^2, ... c^N],
 [0,  1 ,   c, ... c^{N-1}],
 [0,  0 ,   1, ... c^{N-2}],
 .
 .
 .
 [0,  0 ,   0, ....    1]]
</code></pre>
<p>For concreteness, if <code>N=2</code>, then the matrix should be</p>
<pre><code class="lang-auto">[[1, c],
 [0, 1]]
</code></pre>
<p>And if <code>N=3</code>, then the matrix should be:</p>
<pre><code class="lang-auto">
[[1, c, c^2],
 [0, 1,   c],
 [0, 0,   1]]
</code></pre>
<p>How can I do this?</p>

RylanSchaeffer: <p>Right now I’m doing the following. Is there a faster way?</p>
<pre><code class="lang-auto">            discount_matrix = torch.triu(torch.ones(
                size=(N, N)))
            for i in range(1, N):
                tril_idx = torch.triu_indices(
                    row=N,
                    col=N,
                    offset=i)
                discount_matrix[tril_idx[0], tril_idx[1]] *= c

</code></pre>

ptrblck: <p>The Python loop wouldn’t scale that well for larger values of <code>N</code>.<br>
You could use an approach of repeating the initial tensor and calling <code>view</code> on it to create an “interleaved” tensor. This approach would be a bit wasteful in memory, but would yield a better compute performance.</p>
<p>Here is the comparison with a quick profiling for CPU and GPU:</p>
<pre data-code-wrap="python"><code class="lang-python">def fun1(N, c, device):
    discount_matrix = torch.triu(torch.ones(size=(N, N), device=device))
    for i in range(1, N):
        tril_idx = torch.triu_indices(
            row=N,
            col=N,
            offset=i)
        discount_matrix[tril_idx[0], tril_idx[1]] *= c
    return discount_matrix
    

def fun2(N, c, device):
    tmp = torch.tensor([float(c)**i for i in range(N)], device=device)
    tmp = torch.cat((tmp, torch.zeros(1, device=device)))
    res = torch.triu(tmp.repeat(N).view(-1, N)[:-1])
    return res


c = 2
nb_iters = 100

f1_times = []
f2_times = []

device = 'cuda'

for N in range(1, 100):
    # correctness check
    print((fun1(N, c, device) == fun2(N, c, device)).all())
    
    # timing
    torch.cuda.synchronize()
    t0 = time.time()
    for _ in range(nb_iters):
        out = fun1(N, c, device)
    torch.cuda.synchronize()
    t1 = time.time()
    f1_time = (t1 - t0)/nb_iters
    print('fun1({}, {}): {}ms'.format(N, c, f1_time))
    f1_times.append(f1_time)
    
    torch.cuda.synchronize()
    t0 = time.time()
    for _ in range(nb_iters): 
        out = fun2(N, c, device)
    torch.cuda.synchronize()
    t1 = time.time()
    f2_time = (t1 - t0)/nb_iters
    print('fun2({}, {}): {}ms'.format(N, c, f2_time))
    f2_times.append(f2_time)


plt.figure()
plt.plot(np.log10(np.array(f1_times)))
plt.plot(np.log10(np.array(f2_times)))
</code></pre>
<p>Results <code>for N in range(1, 100)</code>. Blue curve: your approach, orange curve: my proposal, y-axis: <code>log10</code> of the time in ms, x-axis: N (sorry for being too lazy to properly add the legends into the plot <img src="https://discuss.pytorch.org/images/emoji/apple/wink.png?v=12" title=":wink:" class="emoji" alt=":wink:" loading="lazy" width="20" height="20"> ).</p>
<p>CPU:<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/1/a/1a76f5c051174656c494c4a948b04b74d83601ee.png" alt="image" data-base62-sha1="3M7hEHJEK0PDf1K6wd2dtZIzrQy" width="381" height="248"></p>
<p>GPU:<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/b/5/b5040e048e4d96788c7d1d6c59a4972af5a673ad.png" alt="image" data-base62-sha1="pPl6MDMHEfoASmFKPmkh3jhYw8J" width="381" height="248"></p>

cbitter: <p>I just encountered the same problem and wanted to share my approach.</p>
<pre><code class="lang-auto">def fun3(N, c, device):
    exponents = torch.triu(torch.ones(N, N, device=device), 1).cumsum(-1)
    power = torch.float_power(c, exponents)
    power_mask = torch.tril(torch.ones(N, N, dtype=torch.bool, device=device), diagonal=-1)
    result = power.masked_fill(power_mask, 0)
    return result
</code></pre>
<p>Avoiding the for-loops seems to be a little bit more performant (tested with GPU):</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/b/b/bbec9cdd672e0a37bf36470c99581b47d65db4c6.png" alt="grafik" data-base62-sha1="qOse6m5KpNNTXkYRdng4aQCT4JE" width="320" height="240"></p>
<p>If anyone wants to replicate my results, you’ll have to adjust the correctness check a little:</p>
<pre><code class="lang-auto">assert torch.isclose(fun2(N, c, device).double(), fun3(N, c, device)).all()
</code></pre>