Open
Description
Description
Numpy has np.block
, which is gives a nice shorthand for repeated concatenations. For example, say I want to make a block lower-triangle matrix:
I can do:
import numpy
A, B, C = np.ones((3, 4, 4))
B *= 2
C *= 3
zeros = np.zeros((4, 4))
D = np.block([[A, zeros], [B, C]])
print(D)
# Result:
[[1. 1. 1. 1. 0. 0. 0. 0.]
[1. 1. 1. 1. 0. 0. 0. 0.]
[1. 1. 1. 1. 0. 0. 0. 0.]
[1. 1. 1. 1. 0. 0. 0. 0.]
[2. 2. 2. 2. 3. 3. 3. 3.]
[2. 2. 2. 2. 3. 3. 3. 3.]
[2. 2. 2. 2. 3. 3. 3. 3.]
[2. 2. 2. 2. 3. 3. 3. 3.]]
Obviously you can do this in a million different ways, with np.hstack
, np.vstack
, np.c_
, np.r_r
, np.concat
, np.stack
, etc, etc. But this one is concise and readable.
In the content of pytensor, it's related to rewrites in the same vein as #1044. We could very easily break apart these block matrices and do the matmul in chunks, especially if we see there are zero matrices among the blocks.