core.py

def assemble(self):
    """Assemble an array on this node from a distributed array of object IDs."""
    first_block = ray.get(self.objectids[(0,) * self.ndim])
    dtype = first_block.dtype
    result = np.zeros(self.shape, dtype=dtype)

    indices = np.ndindex(*self.num_blocks)

    # Fetch all the blocks asynchronously
    blocks = ray.get([self.objectids[idx] for idx in indices])

    for idx, block in zip(indices, blocks):
      lower = DistArray.compute_block_lower(idx, self.shape)
      upper = DistArray.compute_block_upper(idx, self.shape)
      result[[slice(l, u) for (l, u) in zip(lower, upper)]] = block

    return result