(PETSc4py) Creating Distributed Sparse Block Matrix with `createAIJWithArrays'

Question

I am trying to create a sparse block matrix in PETSc4py using local matrices assembled with SciPy. Below is a simple example code with comments

import numpy as np
import petsc4py
from petsc4py import PETSc
import scipy.sparse as sp

def main():
    # Initialize parallel PETSc4py
    petsc4py.init()
    
    comm = PETSc.COMM_WORLD
    
    comm_rank = comm.getRank()
    comm_size = comm.getSize()
    
    # Create local matrices
    n = 4
    n_local = n
    n_global = int(comm_size * n)
    
    A_local = sp.random(n_local, n_local, density = 0.5, format = 'csr') \
        + sp.identity(n)
    A_local = (A_local + A_local.transpose()) / 2.

    # Get CSRSpec for createAIJWithArrays
    [II, JJ, VV] = sp.find(A_local)
    [_, II] = np.unique(II, return_index = True)
    II = np.array(np.append(II, len(II) - 1), dtype = np.int32)
    
    # Create PETSc sparse matrix
    A = PETSc.Mat()
    A.createAIJWithArrays(size = [n_global, n_global], csr = [II, JJ, VV],
                          bsize = [n_local, n_local], comm = comm)
    
    # Communicate off-rank values and setup internal data structures for
    # performing parallel operations
    A.assemblyBegin()
    A.assemblyEnd()
    
    PETSc.Sys.Print(A)

    
if __name__ == '__main__':
    main()

This code works properly when running a single process, but when running with two processes it returns the error

Traceback (most recent call last):
  File "/mnt/c/Users/jltor/Education and Research/misc/PETSc4py/matmpibaij/main.py", line 43, in <module>
    main()
  File "/mnt/c/Users/jltor/Education and Research/misc/PETSc4py/matmpibaij/main.py", line 31, in main
    A.createAIJWithArrays(size = [n_global, n_global], csr = [II, JJ, VV],
  File "petsc4py/PETSc/Mat.pyx", line 449, in petsc4py.PETSc.Mat.createAIJWithArrays
petsc4py.PETSc.Error: error code 55
[0] MatCreateMPIAIJWithArrays() at /home/conda/feedstock_root/build_artifacts/petsc_1688117131282/work/src/mat/impls/aij/mpi/mpiaij.c:4216
[0] MatMPIAIJSetPreallocationCSR() at /home/conda/feedstock_root/build_artifacts/petsc_1688117131282/work/src/mat/impls/aij/mpi/mpiaij.c:4013
[0] MatMPIAIJSetPreallocationCSR_MPIAIJ() at /home/conda/feedstock_root/build_artifacts/petsc_1688117131282/work/src/mat/impls/aij/mpi/mpiaij.c:3936
[0] MatMPIAIJSetPreallocation() at /home/conda/feedstock_root/build_artifacts/petsc_1688117131282/work/src/mat/impls/aij/mpi/mpiaij.c:4148
[0] MatMPIAIJSetPreallocation_MPIAIJ() at /home/conda/feedstock_root/build_artifacts/petsc_1688117131282/work/src/mat/impls/aij/mpi/mpiaij.c:2946
[0] MatSeqAIJSetPreallocation() at /home/conda/feedstock_root/build_artifacts/petsc_1688117131282/work/src/mat/impls/aij/seq/aij.c:3949
[0] MatSeqAIJSetPreallocation_SeqAIJ() at /home/conda/feedstock_root/build_artifacts/petsc_1688117131282/work/src/mat/impls/aij/seq/aij.c:4015
[0] PetscMallocA() at /home/conda/feedstock_root/build_artifacts/petsc_1688117131282/work/src/sys/memory/mal.c:411
[0] PetscMallocAlign() at /home/conda/feedstock_root/build_artifacts/petsc_1688117131282/work/src/sys/memory/mal.c:53
[0] Out of memory. Allocated: 0, Used by process: 80764928
[0] Memory requested 18446744073709551616
[1]PETSC ERROR: ------------------------------------------------------------------------
[1]PETSC ERROR: Caught signal number 11 SEGV: Segmentation Violation, probably memory access out of range
[1]PETSC ERROR: Try option -start_in_debugger or -on_error_attach_debugger
[1]PETSC ERROR: or see https://petsc.org/release/faq/#valgrind and https://petsc.org/release/faq/
[1]PETSC ERROR: configure using --with-debugging=yes, recompile, link, and run
[1]PETSC ERROR: to get more information on the crash.
[1]PETSC ERROR: Run with -malloc_debug to check if memory corruption is causing the crash.

Error Code 55 seems to correspond to an inability to allocate requested memory.

Why is so much memory being requested, and what do I need to change to this example code to fix this?

I have tried running this code with different numbers of processes and different size block matrices, which yielded the same result (including the same amount of memory requested). I also have looked into PETSc4py Error Code 55, which seemed to only return a result on repeatedly setting up a KSP object.