Allow row_reduce to work on stacks of matrices?

[sliedes]

In plain numpy, operations like np.linalg.matrix_rank can work on stacks of matrices:

In [51]: np.linalg.matrix_rank(np.random.randint(0, 2, size=(10, 8, 8)))
Out[51]: array([7, 8, 8, 8, 7, 8, 8, 8, 7, 8])

In [52]: np.linalg.matrix_rank(np.random.randint(0, 2, size=(2, 5, 8, 8)))
Out[52]: 
array([[8, 8, 8, 7, 7],
       [7, 8, 7, 8, 8]])

However, galois apparently is restricted to operating on a single matrix at a time:

In [53]: GF = galois.GF(2)

In [54]: np.linalg.matrix_rank(GF(np.random.randint(0, 1, size=(2, 5, 8, 8))))
...
ValueError: Only 2-D matrices can be converted to reduced row echelon form, not 4-D.

Would it be feasible to make this work on stacks of matrices (essentially, vectorizing the operation)?

[mhostetter]

I wasn't aware np.linalg.matrix_rank() accepted an arbitrary number of dimensions.

In theory, yes. All of these JIT'd functions would need a loop that iterates over the extra dimensions. Perhaps the array could be reshaped to 3-D and loop over the first dimension. Inside the loop, the same operations would happen. After the loop, the array would be reshaped to the N-D array.

This still just loops the code, not really "vectorized". The benefit is the loop is inside the JIT function, so it's faster than a Python loop. Is that what you had in mind?

[sliedes]

I can imagine that might potentially be faster than the current implementation for a stack of matrices. Hard to tell.

I do think Gaussian elimination should be vectorizable, though, albeit I understand if you decide it's not worth the complexity. The pivot would become a vector (one per matrix in the stack) instead of a scalar.

I can imagine something like this (untested code):

A_rre = a.reshape(-1, *A[-2:])

p = np.zeros(A_rre.shape[-1], dtype=np.int_)

for j in range(ncols):
    # find first nonzero for each matrix
    idxs = np.argmax(A_rre[..., p:, j] != 0, axis=-1)  # gives the first non-zero for non-zero rows, 0 for zero rows
    mask = (A_rre[idxs] != 0)  # mask matrices where this this row is non-zero
    if np.all(~mask):
        continue

    # from now on, only operate on the masked matrices

    i = p[mask] + idxs[mask]

    # swap rows p and i
    A_rre[mask, [p, i], :] = A_rre[mask, [i, p], :]

    # force pivot to be 1
    A_rre[mask, p, :] /= A_rre[mask, p, j]
    # ... (np.outer can probably be replaced with dot with reshape or einsum, and so on)

return A_rre.shape(*A.shape[:-2])

[mhostetter]

I'm open to adding support for this. My time is a bit limited at the moment. If you'd like to submit a PR, I can work with you to get it merged. Otherwise, I'll look into it when available.