def solve_data(u, data, sol, naxes, is_rank_zero):
        s = [0]*u.ndim
        s[naxes] = slice(None)
        paxes = np.setxor1d(range(u.ndim), naxes)
        if u.ndim == 2:
            for i in range(u.shape[paxes[0]]):
                if i == 0 and is_rank_zero:
                    continue
                s[paxes[0]] = i
                s0 = tuple(s)
                sol(u[s0], data[i])

        elif u.ndim == 3:
            for i in range(u.shape[paxes[0]]):
                s[paxes[0]] = i
                for j in range(u.shape[paxes[1]]):
                    if i == 0 and j == 0 and is_rank_zero:
                        continue
                    s[paxes[1]] = j
                    s0 = tuple(s)
                    sol(u[s0], data[i, j])
        return u