def METHOD_NAME(TSTEPS: dc.int64, A: dc.float64[Nx, Ny], B: dc.float64[Nx, Ny]):
    lA = np.zeros((lNx + 2, lNy + 2), dtype=A.dtype)
    lB = np.zeros((lNx + 2, lNy + 2), dtype=B.dtype)
    tAB = np.empty((lNx, lNy), dtype=A.dtype)
    req = np.empty((8, ), dtype=MPI_Request)
    Av = np.reshape(A, (Px, lNx, Py, lNy))
    A2 = np.transpose(Av, axes=(0, 2, 1, 3))
    Bv = np.reshape(B, (Px, lNx, Py, lNy))
    B2 = np.transpose(Bv, axes=(0, 2, 1, 3))
    dc.comm.Scatter(A2, tAB)
    lA[1:-1, 1:-1] = tAB
    dc.comm.Scatter(B2, tAB)
    lB[1:-1, 1:-1] = tAB
    for t in range(1, TSTEPS):
        dc.comm.Isend(lA[1, 1:-1], nn, 0, req[0])
        dc.comm.Isend(lA[-2, 1:-1], ns, 1, req[1])
        dc.comm.Isend(lA[1:-1, 1], nw, 2, req[2])
        dc.comm.Isend(lA[1:-1, -2], ne, 3, req[3])
        dc.comm.Irecv(lA[0, 1:-1], nn, 1, req[4])
        dc.comm.Irecv(lA[-1, 1:-1], ns, 0, req[5])
        dc.comm.Irecv(lA[1:-1, 0], nw, 3, req[6])
        dc.comm.Irecv(lA[1:-1, -1], ne, 2, req[7])
        dc.comm.Waitall(req)
        lB[1 + noff:-1 - soff, 1 + woff:-1 -
           eoff] = 0.2 * (lA[1 + noff:-1 - soff, 1 + woff:-1 - eoff] + lA[1 + noff:-1 - soff, woff:-2 - eoff] +
                          lA[1 + noff:-1 - soff, 2 + woff:-eoff] + lA[2 + noff:-soff, 1 + woff:-1 - eoff] +
                          lA[noff:-2 - soff, 1 + woff:-1 - eoff])
        dc.comm.Isend(lB[1, 1:-1], nn, 0, req[0])
        dc.comm.Isend(lB[-2, 1:-1], ns, 1, req[1])
        dc.comm.Isend(lB[1:-1, 1], nw, 2, req[2])
        dc.comm.Isend(lB[1:-1, -2], ne, 3, req[3])
        dc.comm.Irecv(lB[0, 1:-1], nn, 1, req[4])
        dc.comm.Irecv(lB[-1, 1:-1], ns, 0, req[5])
        dc.comm.Irecv(lB[1:-1, 0], nw, 3, req[6])
        dc.comm.Irecv(lB[1:-1, -1], ne, 2, req[7])
        dc.comm.Waitall(req)
        lA[1 + noff:-1 - soff, 1 + woff:-1 -
           eoff] = 0.2 * (lB[1 + noff:-1 - soff, 1 + woff:-1 - eoff] + lB[1 + noff:-1 - soff, woff:-2 - eoff] +
                          lB[1 + noff:-1 - soff, 2 + woff:-eoff] + lB[2 + noff:-soff, 1 + woff:-1 - eoff] +
                          lB[noff:-2 - soff, 1 + woff:-1 - eoff])
    tAB[:] = lA[1:-1, 1:-1]
    dc.comm.Gather(tAB, A2)
    tAB[:] = lB[1:-1, 1:-1]
    dc.comm.Gather(tAB, B2)
    A[:] = np.reshape(np.transpose(A2, (0, 2, 1, 3)), (Nx, Ny))
    B[:] = np.reshape(np.transpose(B2, (0, 2, 1, 3)), (Nx, Ny))