Skip to content

Add helper for constructing block matrices #1100

Open
@jessegrabowski

Description

@jessegrabowski

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:

$$ D = \begin{bmatrix} A & 0 \\ B & C \end{bmatrix} $$

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.

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions