//  [python]
def forward(
        self, x1, x2, 
        graph_nodes, graph_edges, graph_edges_weights, graph_clusters,
        pixel_anchors, pixel_weights, 
        num_nodes_vec, intrinsics, 
        evaluate=False, split="train"
    ):
        batch_size = x1.shape[0]

        image_width = x1.shape[3]
        image_height = x1.shape[2]

        convergence_info = []
        for i in range(batch_size):
            convergence_info.append({
                "total": [],
                "arap": [],
                "data": [],
                "condition_numbers": [],
                "valid": 0,
                "errors": []
            })

        ########################################################################
        # Compute dense flow from source to target.
        ########################################################################
        flow2, flow3, flow4, flow5, flow6, features2 = self.flow_net.forward(x1[:,:3,:,:], x2[:,:3,:,:])
        
        assert torch.isfinite(flow2).all()
        assert torch.isfinite(features2).all()
        
        flow = 20.0 * torch.nn.functional.interpolate(input=flow2, size=(image_height, image_width), mode='bilinear', align_corners=False)

        ########################################################################
        # Initialize graph data.
        ########################################################################
        num_nodes_total = graph_nodes.shape[1]
        num_neighbors = graph_edges.shape[2]

        # We assume we always use 4 nearest anchors.
        assert pixel_anchors.shape[3] == 4

        ########################################################################
        # Apply dense flow to warp the source points to target frame.
        ########################################################################
        x_coords = torch.arange(image_width, dtype=torch.float32, device=x1.device).unsqueeze(0).expand(image_height, image_width).unsqueeze(0)
        y_coords = torch.arange(image_height, dtype=torch.float32, device=x1.device).unsqueeze(1).expand(image_height, image_width).unsqueeze(0)

        xy_coords = torch.cat([x_coords, y_coords], 0)
        xy_coords = xy_coords.unsqueeze(0).repeat(batch_size, 1, 1, 1) # (bs, 2, 448, 640)

        # Apply the flow to pixel coordinates.
        xy_coords_warped = xy_coords + flow
        xy_pixels_warped = xy_coords_warped.clone()

        # Normalize to be between -1, and 1.
        # Since we use "align_corners=False", the boundaries of corner pixels
        # are -1 and 1, not their centers.
        xy_coords_warped[:,0,:,:] = (xy_coords_warped[:,0,:,:]) / (image_width - 1)
        xy_coords_warped[:,1,:,:] = (xy_coords_warped[:,1,:,:]) / (image_height - 1)
        xy_coords_warped = xy_coords_warped * 2 - 1

        # Permute the warped coordinates to fit the grid_sample format.
        xy_coords_warped = xy_coords_warped.permute(0, 2, 3, 1)       

        ########################################################################
        # Construct point-to-point correspondences between source <-> target points.
        ########################################################################
        # Mask out invalid source points.
        source_points = x1[:, 3:, :, :].clone()
        source_anchor_validity = torch.all(pixel_anchors >= 0.0, dim=3)          

        # Sample target points at computed pixel locations.
        target_points = x2[:, 3:, :, :].clone()
        target_matches = torch.nn.functional.grid_sample(
            target_points, xy_coords_warped, mode=opt.gn_depth_sampling_mode, padding_mode='zeros', align_corners=False
        )

        # We filter out any boundary matches, where any of the 4 pixels is invalid.
        target_validity = ((target_points > 0.0) & (target_points <= opt.gn_max_depth)).type(torch.float32)
        target_matches_validity = torch.nn.functional.grid_sample(
            target_validity, xy_coords_warped, mode="bilinear", padding_mode='zeros', align_corners=False
        )
        target_matches_validity = target_matches_validity[:, 2, :, :] >= 0.999

        # Prepare masks for both valid source points and target matches
        valid_source_points  = (source_points[:, 2, :, :] > 0.0)  & (source_points[:, 2, :, :] <= opt.gn_max_depth)  & source_anchor_validity
        valid_target_matches = (target_matches[:, 2, :, :] > 0.0) & (target_matches[:, 2, :, :] <= opt.gn_max_depth) & target_matches_validity

        # Prepare the input of both the MaskNet and AttentionNet, if we actually use either of them
        if opt.use_mask:
            target_rgb = x2[:, :3, :, :].clone()
            target_rgb_warped = torch.nn.functional.grid_sample(target_rgb, xy_coords_warped, padding_mode='zeros', align_corners=False)

            mask_input = torch.cat([x1, target_rgb_warped, target_matches], 1)

        ########################################################################
        # MaskNet
        ########################################################################
        # We predict correspondence weights [0, 1], if we use mask network.
        mask_pred = None
        if opt.use_mask:
            mask_pred = self.mask_net(features2, mask_input).view(batch_size, image_height, image_width)

        # Compute mask of valid correspondences
        valid_correspondences = valid_source_points & valid_target_matches

        # Initialize correspondence weights with 1's. We might overwrite them with MaskNet-predicted weights next
        correspondence_weights = torch.ones(valid_target_matches.shape, dtype=torch.float32, device=valid_target_matches.device)

        # We invalidate target matches later, when we assign a weight to each match.
        if opt.use_mask:
            correspondence_weights = mask_pred
            
            if evaluate:
                # Hard threshold
                if opt.threshold_mask_predictions:
                    correspondence_weights = torch.where(mask_pred < opt.threshold, torch.zeros_like(mask_pred), mask_pred)

                # Patch-wise threshold
                elif opt.patchwise_threshold_mask_predictions:
                    pooled = torch.nn.functional.max_pool2d(input=mask_pred, kernel_size=opt.patch_size, stride=opt.patch_size)
                    pooled = torch.nn.functional.interpolate(input=pooled.unsqueeze(1), size=(opt.image_height, opt.image_width), mode='nearest').squeeze(1)
                    selected = (torch.abs(mask_pred - pooled) <= 1e-8).type(torch.float32) 

                    correspondence_weights = mask_pred * selected #* opt.patch_size**2

        num_valid_correspondences = torch.sum(valid_correspondences, dim=(1, 2))

        ########################################################################
        # Initialize graph data.
        ########################################################################
        num_nodes_total = graph_nodes.shape[1]
        num_neighbors = graph_edges.shape[2]

        # We assume we always use 4 nearest anchors.
        assert pixel_anchors.shape[3] == 4

        node_rotations = torch.eye(3, dtype=x1.dtype, device=x1.device).view(1, 1, 3, 3).repeat(batch_size, num_nodes_total, 1, 1)
        node_translations = torch.zeros((batch_size, num_nodes_total, 3), dtype=x1.dtype, device=x1.device) 
        deformations_validity = torch.zeros((batch_size, num_nodes_total), dtype=x1.dtype, device=x1.device) 
        valid_solve = torch.zeros((batch_size), dtype=torch.uint8, device=x1.device) 
        deformed_points_pred = torch.zeros((batch_size, opt.gn_max_warped_points, 3), dtype=x1.dtype, device=x1.device) 
        deformed_points_idxs = torch.zeros((batch_size, opt.gn_max_warped_points), dtype=torch.int64, device=x1.device) 
        deformed_points_subsampled = torch.zeros((batch_size), dtype=torch.uint8, device=x1.device) 

        # Skip the solver
        if not evaluate and opt.skip_solver:
            return {
                "flow_data": [flow2, flow3, flow4, flow5, flow6], 
                "node_rotations": node_rotations,
                "node_translations": node_translations,
                "deformations_validity": deformations_validity,
                "deformed_points_pred": deformed_points_pred, 
                "valid_solve": valid_solve, 
                "mask_pred": mask_pred,
                "correspondence_info": [
                    xy_coords_warped, 
                    source_points, valid_source_points, 
                    target_matches, valid_target_matches,
                    None, deformed_points_idxs, deformed_points_subsampled
                ], 
                "convergence_info": convergence_info,
                "weight_info": {
                    "total_corres_num": 0,
                    "total_corres_weight": 0.0
                }
            }

        ########################################################################
        # Estimate node deformations using differentiable Gauss-Newton.
        ########################################################################
        correspondences_exist = num_valid_correspondences > 0

        total_num_matches_per_batch = 0

        weight_info = {
            "total_corres_num": 0,
            "total_corres_weight": 0.0
        }

        self.vec_to_skew_mat.to(x1.device)        

        for i in range(batch_size):
            if opt.gn_debug:
                print()
                print("--Sample", i, "in batch--")

            num_nodes_i = num_nodes_vec[i]
            if num_nodes_i > opt.gn_max_nodes or num_nodes_i < opt.gn_min_nodes: 
                print("\tSolver failed: Invalid number of nodes: {0}".format(num_nodes_i))
                convergence_info[i]["errors"].append("Solver failed: Invalid number of nodes: {0}".format(num_nodes_i))
                continue

            if not correspondences_exist[i]: 
                print("\tSolver failed: No valid correspondences before filtering")
                convergence_info[i]["errors"].append("Solver failed: No valid correspondences before filtering")
                continue

            timer_start = timer()

            graph_nodes_i         = graph_nodes[i, :num_nodes_i, :]             # (num_nodes_i, 3)
            graph_edges_i         = graph_edges[i, :num_nodes_i, :]             # (num_nodes_i, 8)
            graph_edges_weights_i = graph_edges_weights[i, :num_nodes_i, :]     # (num_nodes_i, 8)
            graph_clusters_i      = graph_clusters[i, :num_nodes_i, :]          # (num_nodes_i, 1)

            fx = intrinsics[i, 0]
            fy = intrinsics[i, 1]
            cx = intrinsics[i, 2]
            cy = intrinsics[i, 3]

            ###############################################################################################################
            # Filter invalid matches.
            ###############################################################################################################
            valid_correspondences_idxs = torch.where(valid_correspondences[i])

            source_points_filtered = source_points[i].permute(1, 2, 0)
            source_points_filtered = source_points_filtered[valid_correspondences_idxs[0], valid_correspondences_idxs[1], :].view(-1, 3, 1)
            
            target_matches_filtered = target_matches[i].permute(1, 2, 0)
            target_matches_filtered = target_matches_filtered[valid_correspondences_idxs[0], valid_correspondences_idxs[1], :].view(-1, 3, 1)

            xy_pixels_warped_filtered = xy_pixels_warped[i].permute(1, 2, 0) # (height, width, 2)
            xy_pixels_warped_filtered = xy_pixels_warped_filtered[valid_correspondences_idxs[0], valid_correspondences_idxs[1], :].view(-1, 2, 1)

            correspondence_weights_filtered = correspondence_weights[i, valid_correspondences_idxs[0], valid_correspondences_idxs[1]].view(-1) # (num_matches)

            source_anchors = pixel_anchors[i, valid_correspondences_idxs[0], valid_correspondences_idxs[1], :] # (num_matches, 4)
            source_weights = pixel_weights[i, valid_correspondences_idxs[0], valid_correspondences_idxs[1], :] # (num_matches, 4)

            num_matches = source_points_filtered.shape[0]

            ###############################################################################################################
            # Generate weight info to estimate average weight.
            ###############################################################################################################
            weight_info = {
                "total_corres_num": correspondence_weights_filtered.shape[0],
                "total_corres_weight": float(correspondence_weights_filtered.sum())
            }

            ###############################################################################################################
            # Randomly subsample matches, if necessary.
            ###############################################################################################################
            if split == "val" or split == "test":
                max_num_matches = opt.gn_max_matches_eval
            elif split == "train":
                max_num_matches = opt.gn_max_matches_train
            else:
                raise Exception("Split {} is not defined".format(split))

            if num_matches > max_num_matches:
                sampled_idxs = torch.randperm(num_matches)[:max_num_matches]

                source_points_filtered          = source_points_filtered[sampled_idxs]
                target_matches_filtered         = target_matches_filtered[sampled_idxs]
                xy_pixels_warped_filtered       = xy_pixels_warped_filtered[sampled_idxs]
                correspondence_weights_filtered = correspondence_weights_filtered[sampled_idxs]
                source_anchors                  = source_anchors[sampled_idxs]
                source_weights                  = source_weights[sampled_idxs]

                num_matches = max_num_matches

            ###############################################################################################################
            # Remove nodes if their corresponding clusters don't have enough correspondences
            # (Correspondences that have "bad" nodes as anchors will also be removed)
            ###############################################################################################################
            map_opt_nodes_to_complete_nodes_i = list(range(0, num_nodes_i))
            opt_num_nodes_i = num_nodes_i

            if opt.gn_remove_clusters_with_few_matches:
                source_anchors_numpy         = source_anchors.clone().cpu().numpy()
                source_weights_numpy         = source_weights.clone().cpu().numpy()

                # Compute number of correspondences (or matches) per node in the form of the
                # match weight sum
                match_weights_per_node = np.zeros(num_nodes_i)

                # This method adds weight contribution of each match to the corresponding node,
                # allowing duplicate node ids in the flattened array.
                np.add.at(match_weights_per_node, source_anchors_numpy.flatten(), source_weights_numpy.flatten())
                
                total_match_weights = 0.0
                match_weights_per_cluster = {}
                for node_id in range(num_nodes_i):
                    # Get sum of weights for current node.
                    match_weights = match_weights_per_node[node_id]

                    # Get cluster id for current node
                    cluster_id = graph_clusters_i[node_id].item()
                    
                    if cluster_id in match_weights_per_cluster:
                        match_weights_per_cluster[cluster_id] += match_weights
                    else:
                        match_weights_per_cluster[cluster_id] = match_weights

                    total_match_weights += match_weights

                # we'll build a mask that stores which nodes will survive
                valid_nodes_mask_i = torch.ones((num_nodes_i), dtype=torch.bool, device=x1.device)
                
                # if not enough matches in a cluster, mark all cluster's nodes for removal
                node_ids_for_removal = []
                for cluster_id, cluster_match_weights in match_weights_per_cluster.items():
                    if opt.gn_debug:
                        print('cluster_id', cluster_id, cluster_match_weights)

                    if cluster_match_weights < opt.gn_min_num_correspondences_per_cluster:
                        x = torch.where(graph_clusters_i == cluster_id)[0].tolist()
                        node_ids_for_removal += x

                if opt.gn_debug:
                    print("node_ids_for_removal", node_ids_for_removal)

                if len(node_ids_for_removal) > 0:
                    # Mark invalid nodes
                    valid_nodes_mask_i[node_ids_for_removal] = False

                    # Kepp only nodes and edges for valid nodes
                    graph_nodes_i         = graph_nodes_i[valid_nodes_mask_i.squeeze()]
                    graph_edges_i         = graph_edges_i[valid_nodes_mask_i.squeeze()] 
                    graph_edges_weights_i = graph_edges_weights_i[valid_nodes_mask_i.squeeze()] 

                    # Update number of nodes
                    opt_num_nodes_i = graph_nodes_i.shape[0]

                    # Get mask of correspondences for which any one of their anchors is an invalid node
                    valid_corresp_mask = torch.ones((num_matches), dtype=torch.bool, device=x1.device)
                    for node_id_for_removal in node_ids_for_removal:
                        valid_corresp_mask = valid_corresp_mask & torch.all(source_anchors != node_id_for_removal, axis=1)

                    source_points_filtered           = source_points_filtered[valid_corresp_mask]
                    target_matches_filtered          = target_matches_filtered[valid_corresp_mask]
                    xy_pixels_warped_filtered        = xy_pixels_warped_filtered[valid_corresp_mask]
                    correspondence_weights_filtered  = correspondence_weights_filtered[valid_corresp_mask]
                    source_anchors                   = source_anchors[valid_corresp_mask]
                    source_weights                   = source_weights[valid_corresp_mask]

                    num_matches = source_points_filtered.shape[0]

                    # Update node_ids in edges and anchors by mapping old indices to new indices
                    map_opt_nodes_to_complete_nodes_i = []
                    node_count = 0
                    for node_id, is_node_valid in enumerate(valid_nodes_mask_i):
                        if is_node_valid:
                            graph_edges_i[graph_edges_i == node_id]     = node_count
                            source_anchors[source_anchors == node_id]   = node_count
                            map_opt_nodes_to_complete_nodes_i.append(node_id)
                            node_count += 1
                    
            if num_matches == 0: 
                if opt.gn_debug:
                    print("\tSolver failed: No valid correspondences")
                convergence_info[i]["errors"].append("Solver failed: No valid correspondences after filtering")
                continue

            if opt_num_nodes_i > opt.gn_max_nodes or opt_num_nodes_i < opt.gn_min_nodes:
                if opt.gn_debug:
                    print("\tSolver failed: Invalid number of nodes: {0}".format(opt_num_nodes_i))
                convergence_info[i]["errors"].append("Solver failed: Invalid number of nodes: {0}".format(opt_num_nodes_i))
                continue

            # Since source anchor ids need to be long in order to be used as indices,
            # we need to convert them.
            assert torch.all(source_anchors >= 0)
            source_anchors = source_anchors.type(torch.int64)

            ###############################################################################################################
            # Filter invalid edges.
            ###############################################################################################################
            node_ids = torch.arange(opt_num_nodes_i, dtype=torch.int32, device=x1.device).view(-1, 1).repeat(1, num_neighbors) # (opt_num_nodes_i, num_neighbors)
            graph_edge_pairs = torch.cat([node_ids.view(-1, num_neighbors, 1), graph_edges_i.view(-1, num_neighbors, 1)], 2) # (opt_num_nodes_i, num_neighbors, 2)

            valid_edges = graph_edges_i >= 0
            valid_edge_idxs = torch.where(valid_edges)
            graph_edge_pairs_filtered = graph_edge_pairs[valid_edge_idxs[0], valid_edge_idxs[1], :].type(torch.int64)
            graph_edge_weights_pairs  = graph_edges_weights_i[valid_edge_idxs[0], valid_edge_idxs[1]]
            
            num_edges_i = graph_edge_pairs_filtered.shape[0]

            ###############################################################################################################
            # Execute Gauss-Newton solver.
            ###############################################################################################################
            num_gn_iter = self.gn_num_iter
            lambda_data_flow = math.sqrt(self.gn_data_flow)
            lambda_data_depth = math.sqrt(self.gn_data_depth)
            lambda_arap = math.sqrt(self.gn_arap)
            lm_factor = self.gn_lm_factor
            
            # The parameters in GN solver are 3 parameters for rotation and 3 parameters for
            # translation for every node. All node rotation parameters are listed first, and
            # then all node translation parameters are listed.
            #                        x = [w_current_all, t_current_all]
            R_current = torch.eye(3, dtype=x1.dtype, device=x1.device).view(1, 3, 3).repeat(opt_num_nodes_i, 1, 1)
            t_current = torch.zeros((opt_num_nodes_i, 3, 1), dtype=x1.dtype, device=x1.device) 

            if opt.gn_debug:
                print("\tNum. matches: {0} || Num. nodes: {1} || Num. edges: {2}".format(num_matches, opt_num_nodes_i, num_edges_i))

            # Helper structures.
            data_increment_vec_0_3 = torch.arange(0, num_matches * 3, 3, out=torch.cuda.LongTensor(), device=x1.device) # (num_matches)
            data_increment_vec_1_3 = torch.arange(1, num_matches * 3, 3, out=torch.cuda.LongTensor(), device=x1.device) # (num_matches)
            data_increment_vec_2_3 = torch.arange(2, num_matches * 3, 3, out=torch.cuda.LongTensor(), device=x1.device) # (num_matches)

            if num_edges_i > 0:
                arap_increment_vec_0_3 = torch.arange(0, num_edges_i * 3, 3, out=torch.cuda.LongTensor(), device=x1.device) # (num_edges_i)
                arap_increment_vec_1_3 = torch.arange(1, num_edges_i * 3, 3, out=torch.cuda.LongTensor(), device=x1.device) # (num_edges_i)
                arap_increment_vec_2_3 = torch.arange(2, num_edges_i * 3, 3, out=torch.cuda.LongTensor(), device=x1.device) # (num_edges_i)
                arap_one_vec = torch.ones((num_edges_i), dtype=x1.dtype, device=x1.device)

            ill_posed_system = False

            for gn_i in range(num_gn_iter):
                timer_data_start = timer()

                ##########################################
                # Compute data residual and jacobian.
                ##########################################
                jacobian_data = torch.zeros((num_matches * 3, opt_num_nodes_i * 6), dtype=x1.dtype, device=x1.device) # (num_matches*3, opt_num_nodes_i*6)
                deformed_points = torch.zeros((num_matches, 3, 1), dtype=x1.dtype, device=x1.device) 

                for k in range(4): # Our data uses 4 anchors for every point
                    node_idxs_k = source_anchors[:, k] # (num_matches)
                    nodes_k = graph_nodes_i[node_idxs_k].view(num_matches, 3, 1) # (num_matches, 3, 1)

                    # Compute deformed point contribution.                    
                    rotated_points_k = torch.matmul(R_current[node_idxs_k], source_points_filtered - nodes_k) # (num_matches, 3, 1) = (num_matches, 3, 3) * (num_matches, 3, 1)
                    deformed_points_k = rotated_points_k + nodes_k + t_current[node_idxs_k]
                    deformed_points += source_weights[:, k].view(num_matches, 1, 1).repeat(1, 3, 1) * deformed_points_k # (num_matches, 3, 1)

                # Get necessary components of deformed points.
                eps = 1e-7 # Just as good practice, although matches should all have valid depth at this stage

                deformed_x = deformed_points[:, 0, :].view(num_matches) # (num_matches)
                deformed_y = deformed_points[:, 1, :].view(num_matches) # (num_matches)
                deformed_z_inverse = torch.div(1.0, deformed_points[:, 2, :].view(num_matches) + eps) # (num_matches)
                fx_mul_x = fx * deformed_x # (num_matches)
                fy_mul_y = fy * deformed_y # (num_matches)
                fx_div_z = fx * deformed_z_inverse # (num_matches)
                fy_div_z = fy * deformed_z_inverse # (num_matches)
                fx_mul_x_div_z = fx_mul_x * deformed_z_inverse # (num_matches)
                fy_mul_y_div_z = fy_mul_y * deformed_z_inverse # (num_matches)
                minus_fx_mul_x_div_z_2 = -fx_mul_x_div_z * deformed_z_inverse # (num_matches)
                minus_fy_mul_y_div_z_2 = -fy_mul_y_div_z * deformed_z_inverse # (num_matches)

                for k in range(4): # Our data uses 4 anchors for every point
                    node_idxs_k = source_anchors[:, k] # (num_matches)
                    nodes_k = graph_nodes_i[node_idxs_k].view(num_matches, 3, 1) # (num_matches, 3, 1)

                    weights_k = source_weights[:, k] * correspondence_weights_filtered # (num_matches)

                    # Compute skew symetric part.                
                    rotated_points_k = torch.matmul(R_current[node_idxs_k], source_points_filtered - nodes_k) # (num_matches, 3, 1) = (num_matches, 3, 3) * (num_matches, 3, 1)
                    weighted_rotated_points_k = weights_k.view(num_matches, 1, 1).repeat(1, 3, 1) * rotated_points_k # (num_matches, 3, 1)
                    skew_symetric_mat_data = -torch.matmul(self.vec_to_skew_mat, weighted_rotated_points_k).view(num_matches, 3, 3) # (num_matches, 3, 3)

                    # Compute jacobian wrt. TRANSLATION.
                    # FLOW PART
                    jacobian_data[data_increment_vec_0_3, 3 * opt_num_nodes_i + 3 * node_idxs_k + 0] += lambda_data_flow * weights_k * fx_div_z # (num_matches)
                    jacobian_data[data_increment_vec_0_3, 3 * opt_num_nodes_i + 3 * node_idxs_k + 2] += lambda_data_flow * weights_k * minus_fx_mul_x_div_z_2 # (num_matches)
                    jacobian_data[data_increment_vec_1_3, 3 * opt_num_nodes_i + 3 * node_idxs_k + 1] += lambda_data_flow * weights_k * fy_div_z # (num_matches)
                    jacobian_data[data_increment_vec_1_3, 3 * opt_num_nodes_i + 3 * node_idxs_k + 2] += lambda_data_flow * weights_k * minus_fy_mul_y_div_z_2 # (num_matches)
                    
                    # DEPTH PART
                    jacobian_data[data_increment_vec_2_3, 3 * opt_num_nodes_i + 3 * node_idxs_k + 2] += lambda_data_depth * weights_k # (num_matches)

                    # Compute jacobian wrt. ROTATION.
                    # FLOW PART
                    jacobian_data[data_increment_vec_0_3,                   3 * node_idxs_k + 0] += lambda_data_flow * fx_div_z * skew_symetric_mat_data[:, 0, 0] + minus_fx_mul_x_div_z_2 * skew_symetric_mat_data[:, 2, 0]
                    jacobian_data[data_increment_vec_0_3,                   3 * node_idxs_k + 1] += lambda_data_flow * fx_div_z * skew_symetric_mat_data[:, 0, 1] + minus_fx_mul_x_div_z_2 * skew_symetric_mat_data[:, 2, 1]
                    jacobian_data[data_increment_vec_0_3,                   3 * node_idxs_k + 2] += lambda_data_flow * fx_div_z * skew_symetric_mat_data[:, 0, 2] + minus_fx_mul_x_div_z_2 * skew_symetric_mat_data[:, 2, 2]
                    jacobian_data[data_increment_vec_1_3,                   3 * node_idxs_k + 0] += lambda_data_flow * fy_div_z * skew_symetric_mat_data[:, 1, 0] + minus_fy_mul_y_div_z_2 * skew_symetric_mat_data[:, 2, 0]
                    jacobian_data[data_increment_vec_1_3,                   3 * node_idxs_k + 1] += lambda_data_flow * fy_div_z * skew_symetric_mat_data[:, 1, 1] + minus_fy_mul_y_div_z_2 * skew_symetric_mat_data[:, 2, 1]
                    jacobian_data[data_increment_vec_1_3,                   3 * node_idxs_k + 2] += lambda_data_flow * fy_div_z * skew_symetric_mat_data[:, 1, 2] + minus_fy_mul_y_div_z_2 * skew_symetric_mat_data[:, 2, 2]
                    
                    # DEPTH PART
                    jacobian_data[data_increment_vec_2_3,                   3 * node_idxs_k + 0] += lambda_data_depth * skew_symetric_mat_data[:, 2, 0]
                    jacobian_data[data_increment_vec_2_3,                   3 * node_idxs_k + 1] += lambda_data_depth * skew_symetric_mat_data[:, 2, 1]
                    jacobian_data[data_increment_vec_2_3,                   3 * node_idxs_k + 2] += lambda_data_depth * skew_symetric_mat_data[:, 2, 2]

                    assert torch.isfinite(jacobian_data).all(), jacobian_data

                res_data = torch.zeros((num_matches * 3, 1), dtype=x1.dtype, device=x1.device)

                # FLOW PART
                res_data[data_increment_vec_0_3,