Data coordinate zero do not match zero of axes coordinate?

I'm drawing a deep network architecture using matplotlib. I can draw a 2d tensor, 3d tensor, but have problem in coect them. The zero of axes seems do not match the zero of data coordinate if the axes is moved or created not by subplot.

The image produced is Failed to coect layers

The codes are following:

def add_tensor_2d(ax, length = 10, feature_num = 5, notification_box = None, x_label = None, y_label = None):


    uniform_data = np.random.randn(feature_num, length)

    sns.heatmap(uniform_data, xticklabels = False, yticklabels = False, square=True,
            linewidths=.5, ax = ax, cbar = False)

    ax.add_patch(Rectangle((0,0), length, feature_num, fill=False, color="black", linewidth=2))

    if not notification_box:
        notification_box = []

    for box in notification_box:
        ax.add_patch(Rectangle(box["box"][:2], box["box"][2], box["box"][3], fill=True, color=box["color"], alpha=box["alpha"], linewidth=5))


def add_tensor_3d(gs, layer_num, length = 10, feature_num = 5, shift = (0.1, 0.1), notification_box = None, x_label = None, y_label = None):

    ax_box = gs.get_position(plt.gcf())

    pos = [ax_box.x0, ax_box.y0, ax_box.width - (layer_num-1)*shift[0], ax_box.height - (layer_num-1)*shift[1]]

    if not notification_box:
        notification_box = []
    else:
        assert len(notification_box) == layer_num, "size of notification box must be equal to the number of layer"

    axes = []
    for i in range(layer_num):

        cur_ax = plt.gcf().add_axes(pos)

        axes.append(cur_ax)
        if notification_box:
            cur_boxes = notification_box[i]
        else:
            cur_boxes = None

        add_tensor_2d(cur_ax, length, feature_num, x_label = x_label, y_label = y_label, notification_box=cur_boxes)

        x_label = None
        y_label = None

        plt.hold(True)

        pos[0] += shift[0]
        pos[1] += shift[1]

    retu axes


def add_coection(source_ax, source_data_coord,  target_ax, target_data_coord):
    con = CoectionPatch(xyA=source_data_coord, xyB=target_data_coord,
                      coordsA='data', coordsB='data',
                      axesA=source_ax, axesB=target_ax,
                      arrowstyle='->', clip_on=False)

    source_ax.add_artist(con)

if __name__ == "__main__":


    text_length = 20
    embedding_dim = 5
    conv_dim = 10

    embedding_layer_pos = [0, 0, 1, 1]

    embedding_ax = plt.subplot(gs[3, :])
    note_boxes = [{"box": [1.5,0, 3, embedding_dim], "color":"blue", "alpha" : 0.6},
              {"box": [7.5,0, 3, embedding_dim], "color":"red", "alpha" : 0.6},
              {"box": [16.5, 0, 3, embedding_dim], "color": "green", "alpha": 0.6}]
    add_tensor_2d(embedding_ax, length=text_length, feature_num=embedding_dim, x_label="words", y_label= "embedding", notification_box=note_boxes)
    plt.hold(True)


    shift = [0.1,0.025]

    note_boxes = [[{"box": [1, 0, 1, conv_dim], "color": "blue", "alpha": 0.6}, {"box": [0, 8, text_length, 1], "color": "cyan", "alpha": 0.6}, ],
              [{"box": [4, 0, 1, conv_dim], "color": "red", "alpha": 0.6}, {"box": [0, 5, text_length, 1], "color": "magenta", "alpha": 0.6},],
              [{"box": [16, 0, 1, conv_dim], "color": "green", "alpha": 0.6}, {"box": [0, 2, text_length, 1], "color": "orange", "alpha": 0.6},]]
    conv_axes = add_tensor_3d(gs[1:3, :], 3, length=text_length, feature_num=conv_dim, shift = shift, notification_box=note_boxes)

    plt.hold(True)


    add_coection(embedding_ax, (3, embedding_dim), conv_axes[0] , (1.5, 0))


    max_pool_ax = plt.subplot(gs[0, :])
    note_boxes = [{"box": [5, 0, 1, 1], "color": "cyan", "alpha": 0.6},
              {"box": [16, 0, 1, 1], "color": "magenta", "alpha": 0.6},
              {"box": [35, 0, 1, 1], "color": "orange", "alpha": 0.6}]
    add_tensor_2d(max_pool_ax, length=text_length*2, feature_num=1, notification_box=note_boxes)
    plt.hold(True)
    plt.gcf().tight_layout()
    plt.show()