block_lu_working.py

@ray.remote
def lu_decomp_invert(lu_decomp):
    """Takes the inverse of each of the components in the P, L, U
    decomposition. Needed as a helper function for the block-level LU decomp.
    """
    return tuple(np.linalg.inv(x) for x in lu_decomp)


@ray.remote(num_return_vals=3)
def block_lu(a, block_size=100):
    """Returns the LU decomposition of a square matrix.

    Parameters
    ----------
    a : array_like

    Returns
    -------
    p : array_like
    l : array_like
    u : array_like
    """
    if a.shape[0] <= block_size or a.shape[1] <= block_size:
        return ray.get(ra.linalg.lu.remote(a))

    p, l, u = np.zeros(a.shape), np.zeros(a.shape), np.zeros(a.shape)
    num_blocks = int(np.ceil(float(a.shape[0]) / block_size))

    # Compute all single-node LU decompositions in parallel

    block_decomps_remote = []

    for idx in range(num_blocks):
        block_low = block_size * idx
        block_high = block_low + block_size

        a11 = a[block_low:block_high, block_low:block_high]
        block_decomps_remote.append(ra.linalg.lu.remote(a11))

        a12 = a[block_low:block_high, block_high:]
        a21 = a[block_high:, block_low:block_high]

        # Modify a with the Schur complements as we go along so we don't have
        # to repeat the computations later

        schur_complement = np.dot(a21, np.dot(np.linalg.inv(a11), a12))
        a[block_high:, block_high:] -= schur_complement

    block_decomps = ray.get(block_decomps_remote)

    # Compute the inverses for each of the LU components in parallel

    block_decomp_inverses = ray.get([lu_decomp_invert.remote(decomp)
                                     for decomp in block_decomps_remote])

    for idx, (p11, _, _) in enumerate(block_decomps):
        block_low = block_size * idx
        block_high = block_low + block_size

        p[block_low:block_high, block_low:block_high] = p11

    # Perform coalescing.
    # TODO(kvchen): Modify most of this code to use distributed routines.

    for idx, (plu, plu_inverse) in enumerate(zip(block_decomps,
                                                 block_decomp_inverses)):
        p11, l11, u11 = plu
        p11_inverse, l11_inverse, u11_inverse = plu_inverse

        block_low = block_size * idx
        block_high = block_low + block_size

        if idx < num_blocks - 1:
            a12 = a[block_low:block_high, block_high:]
            a21 = a[block_high:, block_low:block_high]

            # Inverse of the permutation matrix is just its transpose

            p22_inverse = p[block_high:, block_high:].T

            u12, l21 = ray.get([
                ra.dot.remote(l11_inverse, ra.dot.remote(p11_inverse, a12)),
                ra.dot.remote(p22_inverse, ra.dot.remote(a21, u11_inverse)),
            ])

            l[block_high:, block_low:block_high] = l21
            u[block_low:block_high, block_high:] = u12

        l[block_low:block_high, block_low:block_high] = l11
        u[block_low:block_high, block_low:block_high] = u11

    return p, l, u

what s the meaning of ra.dot.remote?

ray.get(ra.linalg.lu.remote(a))?
what is the work?