jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def get_cam_rays(H, W, K):
    grid = create_meshgrid(H, W, normalized_coordinates=False)[0]
    i, j = grid.unbind(-1)
    # the direction here is without +0.5 pixel centering as calibration is not so accurate
    # see https://github.com/bmild/nerf/issues/24
    fx, fy, cx, cy = K[0, 0], K[1, 1], K[0, 2], K[1, 2]
    directions = \
        torch.stack([(i - cx) / fx, -(j - cy) / fy, -torch.ones_like(i)], -1)  # (H, W, 3)
    # 求半径
    directions = directions / torch.norm(directions, dim=-1, keepdim=True)
    return directions.numpy()
//  [python]
def get_cam_rays(H, W, K):
    grid = create_meshgrid(H, W, normalized_coordinates=False)[0]
    i, j = grid.unbind(-1)
    # the direction here is without +0.5 pixel centering as calibration is not so accurate
    # see https://github.com/bmild/nerf/issues/24
    fx, fy, cx, cy = K[0, 0], K[1, 1], K[0, 2], K[1, 2]
    directions = \
        torch.stack([(i - cx) / fx, -(j - cy) / fy, -torch.ones_like(i)], -1)  # (H, W, 3)
    # 求半径
    directions = directions / torch.linalg.norm(directions, dim=-1, keepdim=True)
    return directions.numpy()
//  [python]
def log(self, name, step, delta, res, keys_range=None):
        toc = (time() - self.time)*1000
        norm_d = torch.norm(delta).item()
        norm_f = torch.norm(res).item()
        text = '[{}:{:6.2f}]: step = {:3d}, ||delta|| = {:.4f}, ||res|| = {:.4f}'.format(name, toc, step, norm_d, norm_f)
        print(text)
        self.time = time()
        return norm_d, norm_f
//  [python]
def log(self, name, step, delta, res, keys_range=None):
        toc = (time() - self.time)*1000
        norm_d = torch.linalg.norm(delta).item()
        norm_f = torch.linalg.norm(res).item()
        text = '[{}:{:6.2f}]: step = {:3d}, ||delta|| = {:.4f}, ||res|| = {:.4f}'.format(name, toc, step, norm_d, norm_f)
        print(text)
        self.time = time()
        return norm_d, norm_f
//  [python]
def check_keypoints(self, keypoints3d):
        flag = (keypoints3d[..., -1]>0.3).sum() > 5
        if len(self.records) > 1:
            pre = self.records[-1]
            k_pre = self.get_keypoints3d([pre])
            dist = torch.norm(keypoints3d[..., :3] - k_pre[..., :3], dim=-1)
            conf = torch.sqrt(keypoints3d[..., 3] * k_pre[..., 3])
            dist_mean = (dist * conf).sum()/conf.sum()
            flag = flag and dist_mean < 0.1
        return flag
//  [python]
def check_keypoints(self, keypoints3d):
        flag = (keypoints3d[..., -1]>0.3).sum() > 5
        if len(self.records) > 1:
            pre = self.records[-1]
            k_pre = self.get_keypoints3d([pre])
            dist = torch.linalg.norm(keypoints3d[..., :3] - k_pre[..., :3], dim=-1)
            conf = torch.sqrt(keypoints3d[..., 3] * k_pre[..., 3])
            dist_mean = (dist * conf).sum()/conf.sum()
            flag = flag and dist_mean < 0.1
        return flag
//  [python]
def meshArea3D(vertices, faces):
    # Get edge vectors
    fverts = vertices[faces]
    edge1 = fverts[:,1,:] - fverts[:,0,:]
    edge2 = fverts[:,2,:] - fverts[:,0,:]

    # Cross product
    cross = torch.cross(edge1, edge2, dim = 1)
    area = torch.sum(0.5 * torch.norm(cross, dim = 1))

    return area
//  [python]
def meshArea3D(vertices, faces):
    # Get edge vectors
    fverts = vertices[faces]
    edge1 = fverts[:,1,:] - fverts[:,0,:]
    edge2 = fverts[:,2,:] - fverts[:,0,:]

    # Cross product
    cross = torch.cross(edge1, edge2, dim = 1)
    area = torch.sum(0.5 * torch.linalg.norm(cross, dim = 1))

    return area
//  [python]
def fit(self, x, y, criterion=None):
        use_cuda = self.weight.data.is_cuda
        if criterion is None:
            criterion = nn.BCEWithLogitsLoss()
        reduction = criterion.reduction
        criterion.reduction = 'sum'
        if isinstance(x, np.ndarray) or isinstance(y, np.ndarray):
            x = torch.from_numpy(x)
            y = torch.from_numpy(y)
        if self.reverse_complement:
            n_samples, n_features = x.shape
            n_features = n_features // 2
            x = torch.cat([x[:, :n_features], x[:, n_features:]])
        if use_cuda:
            x = x.cuda()
            y = y.cuda()

        if self.bias is not None:
            scale_bias = (x ** 2).mean(-1).sqrt().mean().item()
            self.scale_bias = scale_bias

        def eval_loss(w):
            w = w.reshape((self.out_features, -1))
            if self.weight.grad is not None:
                self.weight.grad = None
            if self.bias is None:
                self.weight.data.copy_(torch.from_numpy(w))
            else:
                if self.bias.grad is not None:
                    self.bias.grad = None
                self.weight.data.copy_(torch.from_numpy(w[:, :-1]))
                self.bias.data.copy_(torch.from_numpy(w[:, -1]))
            y_pred = self(x).squeeze_(-1)
            loss = criterion(y_pred, y)
            loss.backward()
            if self.alpha != 0.0:
                if self.penalty == "l2":
                    penalty = 0.5 * self.alpha * torch.norm(self.weight)**2
                elif self.penalty == "l1":
                    penalty = self.alpha * torch.norm(self.weight, p=1)
                    penalty.backward()
                loss = loss + penalty
            return loss.item()

        def eval_grad(w):
            dw = self.weight.grad.data
            if self.alpha != 0.0:
                if self.penalty == "l2":
                    dw.add_(self.alpha, self.weight.data)
            if self.bias is not None:
                db = self.bias.grad.data
                dw = torch.cat((dw, db.view(-1, 1)), dim=1)
            return dw.cpu().numpy().ravel().astype("float64")

        w_init = self.weight.data
        if self.bias is not None:
            w_init = torch.cat((w_init, self.bias.data.view(-1, 1)), dim=1)
        w_init = w_init.cpu().numpy().astype("float64")

        w = optimize.fmin_l_bfgs_b(
            eval_loss, w_init, fprime=eval_grad, maxiter=100, disp=0)
        if isinstance(w, tuple):
            w = w[0]

        w = w.reshape((self.out_features, -1))
        self.weight.grad.data.zero_()
        if self.bias is None:
            self.weight.data.copy_(torch.from_numpy(w))
        else:
            self.bias.grad.data.zero_()
            self.weight.data.copy_(torch.from_numpy(w[:, :-1]))
            self.bias.data.copy_(torch.from_numpy(w[:, -1]))
        criterion.reduction = reduction
        if self.bias is not None:
            self.bias.data.mul_(self.scale_bias)
        self.scale_bias = None
//  [python]
def fit(self, x, y, criterion=None):
        use_cuda = self.weight.data.is_cuda
        if criterion is None:
            criterion = nn.BCEWithLogitsLoss()
        reduction = criterion.reduction
        criterion.reduction = 'sum'
        if isinstance(x, np.ndarray) or isinstance(y, np.ndarray):
            x = torch.from_numpy(x)
            y = torch.from_numpy(y)
        if self.reverse_complement:
            n_samples, n_features = x.shape
            n_features = n_features // 2
            x = torch.cat([x[:, :n_features], x[:, n_features:]])
        if use_cuda:
            x = x.cuda()
            y = y.cuda()

        if self.bias is not None:
            scale_bias = (x ** 2).mean(-1).sqrt().mean().item()
            self.scale_bias = scale_bias

        def eval_loss(w):
            w = w.reshape((self.out_features, -1))
            if self.weight.grad is not None:
                self.weight.grad = None
            if self.bias is None:
                self.weight.data.copy_(torch.from_numpy(w))
            else:
                if self.bias.grad is not None:
                    self.bias.grad = None
                self.weight.data.copy_(torch.from_numpy(w[:, :-1]))
                self.bias.data.copy_(torch.from_numpy(w[:, -1]))
            y_pred = self(x).squeeze_(-1)
            loss = criterion(y_pred, y)
            loss.backward()
            if self.alpha != 0.0:
                if self.penalty == "l2":
                    penalty = 0.5 * self.alpha * torch.linalg.norm(self.weight)**2
                elif self.penalty == "l1":
                    penalty = self.alpha * torch.linalg.norm(self.weight, p=1)
                    penalty.backward()
                loss = loss + penalty
            return loss.item()

        def eval_grad(w):
            dw = self.weight.grad.data
            if self.alpha != 0.0:
                if self.penalty == "l2":
                    dw.add_(self.alpha, self.weight.data)
            if self.bias is not None:
                db = self.bias.grad.data
                dw = torch.cat((dw, db.view(-1, 1)), dim=1)
            return dw.cpu().numpy().ravel().astype("float64")

        w_init = self.weight.data
        if self.bias is not None:
            w_init = torch.cat((w_init, self.bias.data.view(-1, 1)), dim=1)
        w_init = w_init.cpu().numpy().astype("float64")

        w = optimize.fmin_l_bfgs_b(
            eval_loss, w_init, fprime=eval_grad, maxiter=100, disp=0)
        if isinstance(w, tuple):
            w = w[0]

        w = w.reshape((self.out_features, -1))
        self.weight.grad.data.zero_()
        if self.bias is None:
            self.weight.data.copy_(torch.from_numpy(w))
        else:
            self.bias.grad.data.zero_()
            self.weight.data.copy_(torch.from_numpy(w[:, :-1]))
            self.bias.data.copy_(torch.from_numpy(w[:, -1]))
        criterion.reduction = reduction
        if self.bias is not None:
            self.bias.data.mul_(self.scale_bias)
        self.scale_bias = None
//  [python]
def create_cost_and_derivates(manifold, matrix, backend):
    euclidean_gradient = euclidean_hessian = None

    if backend == "autograd":

        @pymanopt.function.autograd(manifold)
        def cost(Y):
            return np.linalg.norm(Y @ Y.T - matrix, "fro") ** 2

    elif backend == "jax":

        @pymanopt.function.jax(manifold)
        def cost(Y):
            return jnp.linalg.norm(Y @ Y.T - matrix, "fro") ** 2

    elif backend == "numpy":

        @pymanopt.function.numpy(manifold)
        def cost(Y):
            return np.linalg.norm(Y @ Y.T - matrix, "fro") ** 2

        @pymanopt.function.numpy(manifold)
        def euclidean_gradient(Y):
            return 4 * (Y @ Y.T - matrix) @ Y

        @pymanopt.function.numpy(manifold)
        def euclidean_hessian(Y, U):
            return 4 * ((Y @ U.T + U @ Y.T) @ Y + (Y @ Y.T - matrix) @ U)

    elif backend == "pytorch":
        matrix_ = torch.from_numpy(matrix)

        @pymanopt.function.pytorch(manifold)
        def cost(Y):
            X = Y @ torch.transpose(Y, 1, 0)
            return torch.norm(X - matrix_) ** 2

    elif backend == "tensorflow":

        @pymanopt.function.tensorflow(manifold)
        def cost(Y):
            X = Y @ tf.transpose(Y)
            return tf.norm(X - matrix) ** 2

    else:
        raise ValueError(f"Unsupported backend '{backend}'")

    return cost, euclidean_gradient, euclidean_hessian
//  [python]
def create_cost_and_derivates(manifold, matrix, backend):
    euclidean_gradient = euclidean_hessian = None

    if backend == "autograd":

        @pymanopt.function.autograd(manifold)
        def cost(Y):
            return np.linalg.norm(Y @ Y.T - matrix, "fro") ** 2

    elif backend == "jax":

        @pymanopt.function.jax(manifold)
        def cost(Y):
            return jnp.linalg.norm(Y @ Y.T - matrix, "fro") ** 2

    elif backend == "numpy":

        @pymanopt.function.numpy(manifold)
        def cost(Y):
            return np.linalg.norm(Y @ Y.T - matrix, "fro") ** 2

        @pymanopt.function.numpy(manifold)
        def euclidean_gradient(Y):
            return 4 * (Y @ Y.T - matrix) @ Y

        @pymanopt.function.numpy(manifold)
        def euclidean_hessian(Y, U):
            return 4 * ((Y @ U.T + U @ Y.T) @ Y + (Y @ Y.T - matrix) @ U)

    elif backend == "pytorch":
        matrix_ = torch.from_numpy(matrix)

        @pymanopt.function.pytorch(manifold)
        def cost(Y):
            X = Y @ torch.transpose(Y, 1, 0)
            return torch.linalg.norm(X - matrix_) ** 2

    elif backend == "tensorflow":

        @pymanopt.function.tensorflow(manifold)
        def cost(Y):
            X = Y @ tf.transpose(Y)
            return tf.norm(X - matrix) ** 2

    else:
        raise ValueError(f"Unsupported backend '{backend}'")

    return cost, euclidean_gradient, euclidean_hessian
//  [python]
def clip_gradient(self,cond_g):
		
		norm_grad = tr.norm(cond_g)
		clip_coef = self.clip_value / (norm_grad + 1e-6)
		if clip_coef<1.:
			self.dot_prod = self.dot_prod/norm_grad
			return cond_g/norm_grad
		else:
			return cond_g
//  [python]
def clip_gradient(self,cond_g):
		
		norm_grad = tr.linalg.norm(cond_g)
		clip_coef = self.clip_value / (norm_grad + 1e-6)
		if clip_coef<1.:
			self.dot_prod = self.dot_prod/norm_grad
			return cond_g/norm_grad
		else:
			return cond_g
//  [python]
def get_keypoints3d(self, records):
        THRES_WRIST = 0.05
        keypoints3d = super().get_keypoints3d(records)
        keypoints3d = keypoints3d[:, self.INDEX_HALF]
        handl = super().get_keypoints3d(records, key='handl3d')
        handr = super().get_keypoints3d(records, key='handr3d')
        dist_ll = torch.norm(keypoints3d[:, 7, :3] - handl[:, 0, :3], dim=-1)
        dist_rr = torch.norm(keypoints3d[:, 4, :3] - handr[:, 0, :3], dim=-1)
        handl[dist_ll>THRES_WRIST] = 0.
        handr[dist_rr>THRES_WRIST] = 0.
        keypoints3d = np.hstack([keypoints3d, handl, handr])
        conf = keypoints3d[..., 3:]
        keypoints3d = np.hstack([(keypoints3d[..., :3] * conf).sum(axis=0)/(1e-5 + conf.sum(axis=0)), conf.min(axis=0)])
        keypoints3d = keypoints3d[None]
        # if (keypoints3d.shape[0] == 10):
        return torch.Tensor(keypoints3d)
//  [python]
def get_keypoints3d(self, records):
        THRES_WRIST = 0.05
        keypoints3d = super().get_keypoints3d(records)
        keypoints3d = keypoints3d[:, self.INDEX_HALF]
        handl = super().get_keypoints3d(records, key='handl3d')
        handr = super().get_keypoints3d(records, key='handr3d')
        dist_ll = torch.linalg.norm(keypoints3d[:, 7, :3] - handl[:, 0, :3], dim=-1)
        dist_rr = torch.linalg.norm(keypoints3d[:, 4, :3] - handr[:, 0, :3], dim=-1)
        handl[dist_ll>THRES_WRIST] = 0.
        handr[dist_rr>THRES_WRIST] = 0.
        keypoints3d = np.hstack([keypoints3d, handl, handr])
        conf = keypoints3d[..., 3:]
        keypoints3d = np.hstack([(keypoints3d[..., :3] * conf).sum(axis=0)/(1e-5 + conf.sum(axis=0)), conf.min(axis=0)])
        keypoints3d = keypoints3d[None]
        # if (keypoints3d.shape[0] == 10):
        return torch.Tensor(keypoints3d)
//  [python]
def add_densification_stats(self, viewspace_point_tensor, update_filter):
        self.xyz_gradient_accum[update_filter] += torch.norm(viewspace_point_tensor.grad[update_filter,:2], dim=-1, keepdim=True)
        self.denom[update_filter] += 1
//  [python]
def add_densification_stats(self, viewspace_point_tensor, update_filter):
        self.xyz_gradient_accum[update_filter] += torch.linalg.norm(viewspace_point_tensor.grad[update_filter,:2], dim=-1, keepdim=True)
        self.denom[update_filter] += 1
//  [python]
def densify_and_clone(self, grads, grad_threshold, scene_extent):
        # Extract points that satisfy the gradient condition
        selected_pts_mask = torch.where(torch.norm(grads, dim=-1) >= grad_threshold, True, False)
        selected_pts_mask = torch.logical_and(selected_pts_mask,
                                              torch.max(self.get_scaling, dim=1).values <= self.percent_dense*scene_extent)
        
        new_xyz = self._xyz[selected_pts_mask]
        new_features_dc = self._features_dc[selected_pts_mask]
        new_features_rest = self._features_rest[selected_pts_mask]
        new_opacities = self._opacity[selected_pts_mask]
        new_scaling = self._scaling[selected_pts_mask]
        new_rotation = self._rotation[selected_pts_mask]

        self.densification_postfix(new_xyz, new_features_dc, new_features_rest, new_opacities, new_scaling, new_rotation)
//  [python]
def densify_and_clone(self, grads, grad_threshold, scene_extent):
        # Extract points that satisfy the gradient condition
        selected_pts_mask = torch.where(torch.linalg.norm(grads, dim=-1) >= grad_threshold, True, False)
        selected_pts_mask = torch.logical_and(selected_pts_mask,
                                              torch.max(self.get_scaling, dim=1).values <= self.percent_dense*scene_extent)
        
        new_xyz = self._xyz[selected_pts_mask]
        new_features_dc = self._features_dc[selected_pts_mask]
        new_features_rest = self._features_rest[selected_pts_mask]
        new_opacities = self._opacity[selected_pts_mask]
        new_scaling = self._scaling[selected_pts_mask]
        new_rotation = self._rotation[selected_pts_mask]

        self.densification_postfix(new_xyz, new_features_dc, new_features_rest, new_opacities, new_scaling, new_rotation)
//  [python]
def l21shrink_torch(epsilon, X):
    """
    Args:
        epsilon: the shrinkage parameter
        X: matrix to shrink on
    Ref:
        wiki Regularization: {https://en.wikipedia.org/wiki/Regularization_(mathematics)}
    Returns:
            The shrunk matrix
    """
    output = X.clone()
    norm = torch.norm(X, p='fro', dim=0)
    for i in range(X.shape[1]):
        if norm[i] > epsilon:
            for j in range(X.shape[0]):
                output[j, i] = X[j, i] - epsilon * X[j, i] / norm[i]
        else:
            output[:, i] = 0.
    return output
//  [python]
def l21shrink_torch(epsilon, X):
    """
    Args:
        epsilon: the shrinkage parameter
        X: matrix to shrink on
    Ref:
        wiki Regularization: {https://en.wikipedia.org/wiki/Regularization_(mathematics)}
    Returns:
            The shrunk matrix
    """
    output = X.clone()
    norm = torch.linalg.norm(X, ord='fro', dim=0)
    for i in range(X.shape[1]):
        if norm[i] > epsilon:
            for j in range(X.shape[0]):
                output[j, i] = X[j, i] - epsilon * X[j, i] / norm[i]
        else:
            output[:, i] = 0.
    return output
//  [python]
def batch_rodrigues_jacobi(rvec):
    shape = rvec.shape
    rvec = rvec.view(-1, 3)
    device = rvec.device
    dSkew = torch.zeros(3, 9, device=device)
    dSkew[0, 5] = -1
    dSkew[1, 6] = -1
    dSkew[2, 1] = -1
    dSkew[0, 7] =  1
    dSkew[1, 2] =  1
    dSkew[2, 3] =  1
    dSkew = dSkew[None]
    theta = torch.norm(rvec, dim=-1, keepdim=True) + 1e-5
    c = torch.cos(theta)
    s = torch.sin(theta)
    c1 = 1 - c
    itheta = 1 / theta
    r = rvec / theta
    zeros = torch.zeros_like(r[:, :1])
    rx, ry, rz = torch.split(r, 1, dim=1)
    rrt = torch.matmul(r[:, :, None], r[:, None, :])
    skew = torch.cat([zeros, -rz, ry, rz, zeros, -rx, -ry, rx, zeros], dim=1) \
        .view((r.shape[0], 3, 3))
    I = torch.eye(3, device=rvec.device, dtype=rvec.dtype)[None]
    rot_mat = I + s[:, None] * skew + c1[:, None] * torch.bmm(skew, skew)
    
    drrt = torch.stack([
        rx + rx, ry, rz, ry, zeros, zeros, rz, zeros, zeros,
        zeros, rx, zeros, rx, ry + ry, rz, zeros, rz, zeros,
        zeros, zeros, rx, zeros, zeros, ry, rx, ry, rz + rz
    ], dim=-1).view((r.shape[0], 3, 9))
    jacobi = torch.zeros((r.shape[0], 3, 9), device=rvec.device, dtype=rvec.dtype)
    for i in range(3):
        ri = r[:, i:i+1]
        a0 = -s * ri
        a1 = (s - 2*c1*itheta)*ri
        a2 = c1 * itheta
        a3 = (c-s*itheta)*ri
        a4 = s * itheta
        jaco = a0[:, None] * I + a1[:, None] * rrt + a2[:, None] * drrt[:, i].view(-1, 3, 3) + a3[:, None] * skew + a4[:, None] * dSkew[:, i].view(-1, 3, 3)
        jacobi[:, i] = jaco.view(-1, 9)
    rot_mat = rot_mat.view(*shape[:-1], 3, 3)
    jacobi = jacobi.view(*shape[:-1], 3, 9)
    return rot_mat, jacobi
//  [python]
def batch_rodrigues_jacobi(rvec):
    shape = rvec.shape
    rvec = rvec.view(-1, 3)
    device = rvec.device
    dSkew = torch.zeros(3, 9, device=device)
    dSkew[0, 5] = -1
    dSkew[1, 6] = -1
    dSkew[2, 1] = -1
    dSkew[0, 7] =  1
    dSkew[1, 2] =  1
    dSkew[2, 3] =  1
    dSkew = dSkew[None]
    theta = torch.linalg.norm(rvec, dim=-1, keepdim=True) + 1e-5
    c = torch.cos(theta)
    s = torch.sin(theta)
    c1 = 1 - c
    itheta = 1 / theta
    r = rvec / theta
    zeros = torch.zeros_like(r[:, :1])
    rx, ry, rz = torch.split(r, 1, dim=1)
    rrt = torch.matmul(r[:, :, None], r[:, None, :])
    skew = torch.cat([zeros, -rz, ry, rz, zeros, -rx, -ry, rx, zeros], dim=1) \
        .view((r.shape[0], 3, 3))
    I = torch.eye(3, device=rvec.device, dtype=rvec.dtype)[None]
    rot_mat = I + s[:, None] * skew + c1[:, None] * torch.bmm(skew, skew)
    
    drrt = torch.stack([
        rx + rx, ry, rz, ry, zeros, zeros, rz, zeros, zeros,
        zeros, rx, zeros, rx, ry + ry, rz, zeros, rz, zeros,
        zeros, zeros, rx, zeros, zeros, ry, rx, ry, rz + rz
    ], dim=-1).view((r.shape[0], 3, 9))
    jacobi = torch.zeros((r.shape[0], 3, 9), device=rvec.device, dtype=rvec.dtype)
    for i in range(3):
        ri = r[:, i:i+1]
        a0 = -s * ri
        a1 = (s - 2*c1*itheta)*ri
        a2 = c1 * itheta
        a3 = (c-s*itheta)*ri
        a4 = s * itheta
        jaco = a0[:, None] * I + a1[:, None] * rrt + a2[:, None] * drrt[:, i].view(-1, 3, 3) + a3[:, None] * skew + a4[:, None] * dSkew[:, i].view(-1, 3, 3)
        jacobi[:, i] = jaco.view(-1, 9)
    rot_mat = rot_mat.view(*shape[:-1], 3, 3)
    jacobi = jacobi.view(*shape[:-1], 3, 9)
    return rot_mat, jacobi
//  [python]
def _affine_invariant_metric(
        self, x: torch.Tensor, y: torch.Tensor, keepdim=False
    ) -> torch.Tensor:
        r"""Affine Invariant Metric distance.

        {}

        References
        ----------
        A Riemannian framework for tensor computing. 2006.
        """.format(
            self._dist_doc
        )
        inv_sqrt_x = linalg.sym_inv_sqrtm1(x)
        return torch.norm(
            linalg.sym_logm(inv_sqrt_x @ y @ inv_sqrt_x),
            dim=[-1, -2],
            keepdim=keepdim,
        )
//  [python]
def _affine_invariant_metric(
        self, x: torch.Tensor, y: torch.Tensor, keepdim=False
    ) -> torch.Tensor:
        r"""Affine Invariant Metric distance.

        {}

        References
        ----------
        A Riemannian framework for tensor computing. 2006.
        """.format(
            self._dist_doc
        )
        inv_sqrt_x = linalg.sym_inv_sqrtm1(x)
        return torch.linalg.norm(
            linalg.sym_logm(inv_sqrt_x @ y @ inv_sqrt_x),
            dim=[-1, -2],
            keepdim=keepdim,
        )
//  [python]
def get_geman_mclure_weight(xyz, xyz_target, mu):
    """
    compute the weights defined here for the iterative reweighted least square.
    http://vladlen.info/papers/fast-global-registration.pdf
    """
    norm2 = torch.norm(xyz_target - xyz, dim=1) ** 2
    return (mu / (mu + norm2)).view(-1, 1)
//  [python]
def get_geman_mclure_weight(xyz, xyz_target, mu):
    """
    compute the weights defined here for the iterative reweighted least square.
    http://vladlen.info/papers/fast-global-registration.pdf
    """
    norm2 = torch.linalg.norm(xyz_target - xyz, dim=1) ** 2
    return (mu / (mu + norm2)).view(-1, 1)
//  [python]
def fitShape(self, keypoints3d, params, weight, option):
        kintree = np.array(self.cfg.shape.kintree)
        nShapes = params['shapes'].shape[-1]
        # shapes: (1, 10)
        shapes = params['shapes'].T
        shapes0 = shapes.clone()
        device, dtype = shapes.device, shapes.dtype
        lengths3d_est = torch.norm(keypoints3d[:, kintree[:, 1], :3] - keypoints3d[:, kintree[:, 0], :3], dim=-1)
        conf = torch.sqrt(keypoints3d[:, kintree[:, 1], 3:] * keypoints3d[:, kintree[:, 0], 3:])
        conf = conf.repeat(1, 1, 3).reshape(-1, 1)
        nFrames = keypoints3d.shape[0]
        # jacobian: (nFrames, nLimbs, 3, nShapes)
        jacob_limb_shapes = self.jacobian_limb_shapes[None].repeat(nFrames, 1, 1, 1)
        jacob_limb_shapes = jacob_limb_shapes.reshape(-1, nShapes)
        # 注意：这里乘到雅克比的应该是 sqrt(conf)，这里把两个合并了
        JTJ_limb_shapes = jacob_limb_shapes.t() @ (jacob_limb_shapes * conf)
        lossnorm = 0
        self.time = time()
        for iter_ in range(option.max_iter):
            # perform shape blending
            shape_offset = self.k_shapeBlend @ shapes
            keyShaped = self.k_template + shape_offset[..., 0]
            JTJ = JTJ_limb_shapes
            JTr = torch.zeros((nShapes, 1), device=device, dtype=dtype)
            # 并行添加所有的骨架
            dir = keyShaped[kintree[:, 1]] - keyShaped[kintree[:, 0]]
            dir_normalized = dir / torch.norm(dir, dim=-1, keepdim=True)
            # res: (nFrames, nLimbs, 3)
            res = lengths3d_est[..., None] * dir_normalized[None] - dir[None]
            res = conf * res.reshape(-1, 1)
            JTr += jacob_limb_shapes.t() @ res
            self.add_any_reg(shapes, shapes0, JTJ, JTr, w=weight.init_shape)
            delta = torch.linalg.solve(JTJ, JTr)
            shapes += delta
            norm_d, norm_f = self.log('shape', iter_, delta, res)
            if torch.norm(delta) < option.gtol:
                break
            if iter_ > 0 and abs(norm_f - lossnorm)/norm_f < option.ftol:
                break
            lossnorm = norm_f
        shapes = shapes.t()
        params['shapes'] = shapes
        return params, keyShaped
//  [python]
def fitShape(self, keypoints3d, params, weight, option):
        kintree = np.array(self.cfg.shape.kintree)
        nShapes = params['shapes'].shape[-1]
        # shapes: (1, 10)
        shapes = params['shapes'].T
        shapes0 = shapes.clone()
        device, dtype = shapes.device, shapes.dtype
        lengths3d_est = torch.linalg.norm(keypoints3d[:, kintree[:, 1], :3] - keypoints3d[:, kintree[:, 0], :3], dim=-1)
        conf = torch.sqrt(keypoints3d[:, kintree[:, 1], 3:] * keypoints3d[:, kintree[:, 0], 3:])
        conf = conf.repeat(1, 1, 3).reshape(-1, 1)
        nFrames = keypoints3d.shape[0]
        # jacobian: (nFrames, nLimbs, 3, nShapes)
        jacob_limb_shapes = self.jacobian_limb_shapes[None].repeat(nFrames, 1, 1, 1)
        jacob_limb_shapes = jacob_limb_shapes.reshape(-1, nShapes)
        # 注意：这里乘到雅克比的应该是 sqrt(conf)，这里把两个合并了
        JTJ_limb_shapes = jacob_limb_shapes.t() @ (jacob_limb_shapes * conf)
        lossnorm = 0
        self.time = time()
        for iter_ in range(option.max_iter):
            # perform shape blending
            shape_offset = self.k_shapeBlend @ shapes
            keyShaped = self.k_template + shape_offset[..., 0]
            JTJ = JTJ_limb_shapes
            JTr = torch.zeros((nShapes, 1), device=device, dtype=dtype)
            # 并行添加所有的骨架
            dir = keyShaped[kintree[:, 1]] - keyShaped[kintree[:, 0]]
            dir_normalized = dir / torch.linalg.norm(dir, dim=-1, keepdim=True)
            # res: (nFrames, nLimbs, 3)
            res = lengths3d_est[..., None] * dir_normalized[None] - dir[None]
            res = conf * res.reshape(-1, 1)
            JTr += jacob_limb_shapes.t() @ res
            self.add_any_reg(shapes, shapes0, JTJ, JTr, w=weight.init_shape)
            delta = torch.linalg.solve(JTJ, JTr)
            shapes += delta
            norm_d, norm_f = self.log('shape', iter_, delta, res)
            if torch.linalg.norm(delta) < option.gtol:
                break
            if iter_ > 0 and abs(norm_f - lossnorm)/norm_f < option.ftol:
                break
            lossnorm = norm_f
        shapes = shapes.t()
        params['shapes'] = shapes
        return params, keyShaped
//  [python]
def _log_eucliden_metric(
        self, x: torch.Tensor, y: torch.Tensor, keepdim=False
    ) -> torch.Tensor:
        r"""Log-Eucliden Metric distance.

        {}

        References
        ----------
        Log‐Euclidean metrics for fast and simple calculus on diffusion tensors. 2006.
        """.format(
            self._dist_doc
        )
        return torch.norm(
            linalg.sym_logm(x) - linalg.sym_logm(y),
            dim=[-1, -2],
            keepdim=keepdim,
        )
//  [python]
def _log_eucliden_metric(
        self, x: torch.Tensor, y: torch.Tensor, keepdim=False
    ) -> torch.Tensor:
        r"""Log-Eucliden Metric distance.

        {}

        References
        ----------
        Log‐Euclidean metrics for fast and simple calculus on diffusion tensors. 2006.
        """.format(
            self._dist_doc
        )
        return torch.linalg.norm(
            linalg.sym_logm(x) - linalg.sym_logm(y),
            dim=[-1, -2],
            keepdim=keepdim,
        )
//  [python]
def _get_val1(self, val): return torch.norm(val).pow(2)/val.numel()
//  [python]
def _get_val1(self, val): return torch.linalg.norm(val).pow(2)/val.numel()
//  [python]
def orthogonal_error(self, weight):
        size = weight.shape[0]
        return torch.norm(torch.norm(torch.eye(size).to(weight.device) -
                              torch.mm(weight, torch.t(weight)), 2) +
                           torch.norm(torch.eye(size).to(weight.device) -
                                      torch.mm(torch.t(weight), weight), 2), 2)
//  [python]
def orthogonal_error(self, weight):
        size = weight.shape[0]
        return torch.linalg.norm(torch.linalg.norm(torch.eye(size).to(weight.device) -
                              torch.mm(weight, torch.t(weight)), 2) +
                           torch.linalg.norm(torch.eye(size).to(weight.device) -
                                      torch.mm(torch.t(weight), weight), 2), 2)
//  [python]
def reg_error(self):
        return torch.max(torch.norm(self.weight, self.p) - self.sigma_max, torch.zeros(1)) + \
               torch.max(self.sigma_min - torch.norm(self.weight, self.p), torch.zeros(1))
//  [python]
def reg_error(self):
        return torch.max(torch.linalg.norm(self.weight, self.p) - self.sigma_max, torch.zeros(1)) + \
               torch.max(self.sigma_min - torch.linalg.norm(self.weight, self.p), torch.zeros(1))
//  [python]
def reps_cosine_sim(sent1_reps: torch.Tensor, sent2_reps: torch.Tensor) -> torch.Tensor:
    """
    calculate representation cosine similarity, note that this is different from torch version(that compute parwisely)
    :param sent1_reps: (N, sent1_len, reps_dim)
    :param sent2_reps: (N, sent2_len, reps_dim)
    :return: (N, sent1_len, sent2_len)
    """
    dot_sim = torch.bmm(sent1_reps, torch.transpose(sent2_reps, -1, -2))  # shape: (batch, seq_len1, seq_len2)
    sent1_reps_norm = torch.norm(sent1_reps, dim=-1, keepdim=True)  # shape: (batch, seq_len1, 1)
    sent2_reps_norm = torch.norm(sent2_reps, dim=-1, keepdim=True)  # shape: (batch, seq_len2, 1)
    norm_product = torch.bmm(sent1_reps_norm,
                             torch.transpose(sent2_reps_norm, -1, -2))  # shape: (batch, seq_len1, seq_len2)
    sim_predicts = dot_sim / norm_product  # shape: (batch, seq_len1, seq_len2)
    return sim_predicts
//  [python]
def reps_cosine_sim(sent1_reps: torch.Tensor, sent2_reps: torch.Tensor) -> torch.Tensor:
    """
    calculate representation cosine similarity, note that this is different from torch version(that compute parwisely)
    :param sent1_reps: (N, sent1_len, reps_dim)
    :param sent2_reps: (N, sent2_len, reps_dim)
    :return: (N, sent1_len, sent2_len)
    """
    dot_sim = torch.bmm(sent1_reps, torch.transpose(sent2_reps, -1, -2))  # shape: (batch, seq_len1, seq_len2)
    sent1_reps_norm = torch.linalg.norm(sent1_reps, dim=-1, keepdim=True)  # shape: (batch, seq_len1, 1)
    sent2_reps_norm = torch.linalg.norm(sent2_reps, dim=-1, keepdim=True)  # shape: (batch, seq_len2, 1)
    norm_product = torch.bmm(sent1_reps_norm,
                             torch.transpose(sent2_reps_norm, -1, -2))  # shape: (batch, seq_len1, seq_len2)
    sim_predicts = dot_sim / norm_product  # shape: (batch, seq_len1, seq_len2)
    return sim_predicts
//  [python]
def gdupdateWexact(K, V, Ktarget1, Vtarget1, W, device='cuda'):
    input_ = K
    output = V
    C = input_.T@input_
    d = []
    lu, piv = lu_factor(C.cpu().numpy())
    for i in range(Ktarget1.size(0)):
        sol = lu_solve((lu, piv), Ktarget1[i].reshape(-1, 1).cpu().numpy())
        d.append(torch.from_numpy(sol).to(K.device))

    d = torch.cat(d, 1).T

    e2 = d@Ktarget1.T
    e1 = (Vtarget1.T - W@Ktarget1.T)
    delta = e1@torch.linalg.inv(e2)

    Wnew = W + delta@d
    lambda_split1 = Vtarget1.size(0)

    input_ = torch.cat([Ktarget1.T, K.T], dim=1)
    output = torch.cat([Vtarget1, V], dim=0)

    loss = torch.norm((Wnew@input_).T - output, 2, dim=1)
    print(loss[:lambda_split1].mean().item(), loss[lambda_split1:].mean().item())

    return Wnew
//  [python]
def gdupdateWexact(K, V, Ktarget1, Vtarget1, W, device='cuda'):
    input_ = K
    output = V
    C = input_.T@input_
    d = []
    lu, piv = lu_factor(C.cpu().numpy())
    for i in range(Ktarget1.size(0)):
        sol = lu_solve((lu, piv), Ktarget1[i].reshape(-1, 1).cpu().numpy())
        d.append(torch.from_numpy(sol).to(K.device))

    d = torch.cat(d, 1).T

    e2 = d@Ktarget1.T
    e1 = (Vtarget1.T - W@Ktarget1.T)
    delta = e1@torch.linalg.inv(e2)

    Wnew = W + delta@d
    lambda_split1 = Vtarget1.size(0)

    input_ = torch.cat([Ktarget1.T, K.T], dim=1)
    output = torch.cat([Vtarget1, V], dim=0)

    loss = torch.linalg.norm((Wnew@input_).T - output, 2, dim=1)
    print(loss[:lambda_split1].mean().item(), loss[lambda_split1:].mean().item())

    return Wnew
//  [python]
def compute_beamforming_vector(target_scm: torch.Tensor, noise_scm: torch.Tensor):
        noise_scm_t = noise_scm.permute(0, 3, 1, 2)
        noise_scm_t = condition_scm(noise_scm_t, 1e-6)
        e_val, e_vec = generalized_eigenvalue_decomposition(
            target_scm.permute(0, 3, 1, 2), noise_scm_t
        )
        bf_vect = e_vec[..., -1]
        # Normalize
        bf_vect /= torch.norm(bf_vect, dim=-1, keepdim=True)
        bf_vect = bf_vect.squeeze(-1).transpose(-1, -2)  # -> bft
        return bf_vect
//  [python]
def compute_beamforming_vector(target_scm: torch.Tensor, noise_scm: torch.Tensor):
        noise_scm_t = noise_scm.permute(0, 3, 1, 2)
        noise_scm_t = condition_scm(noise_scm_t, 1e-6)
        e_val, e_vec = generalized_eigenvalue_decomposition(
            target_scm.permute(0, 3, 1, 2), noise_scm_t
        )
        bf_vect = e_vec[..., -1]
        # Normalize
        bf_vect /= torch.linalg.norm(bf_vect, dim=-1, keepdim=True)
        bf_vect = bf_vect.squeeze(-1).transpose(-1, -2)  # -> bft
        return bf_vect
//  [python]
def step(self, closure=None):
        """Performs a single optimization step.

        Arguments:
            closure (callable, optional): A closure that reevaluates the model
                and returns the loss.
        """
        loss = None
        if closure is not None:
            with torch.enable_grad():
                loss = closure()

        for group in self.param_groups:
            weight_decay = group["weight_decay"]
            momentum = group["momentum"]
            eeta = group["eeta"]
            lr = group["lr"]
            lars = group["lars"]
            eps = group["epsilon"]

            for p in group["params"]:
                if p.grad is None:
                    continue
                decayed_grad = p.grad
                scaled_lr = lr
                if lars:
                    w_norm = torch.norm(p)
                    g_norm = torch.norm(p.grad)
                    trust_ratio = torch.where(
                        w_norm > 0 and g_norm > 0,
                        eeta * w_norm / (g_norm + weight_decay * w_norm + eps),
                        torch.ones_like(w_norm),
                    )
                    trust_ratio.clamp_(0.0, 50)
                    scaled_lr *= trust_ratio.item()
                    if weight_decay != 0:
                        decayed_grad = decayed_grad.add(p, alpha=weight_decay)
                decayed_grad = torch.clamp(decayed_grad, -10.0, 10.0)

                if momentum != 0:
                    param_state = self.state[p]
                    if "momentum_buffer" not in param_state:
                        buf = param_state["momentum_buffer"] = torch.clone(decayed_grad).detach()
                    else:
                        buf = param_state["momentum_buffer"]
                        buf.mul_(momentum).add_(decayed_grad)
                    decayed_grad = buf

                p.add_(decayed_grad, alpha=-scaled_lr)

        return loss
//  [python]
def step(self, closure=None):
        """Performs a single optimization step.

        Arguments:
            closure (callable, optional): A closure that reevaluates the model
                and returns the loss.
        """
        loss = None
        if closure is not None:
            with torch.enable_grad():
                loss = closure()

        for group in self.param_groups:
            weight_decay = group["weight_decay"]
            momentum = group["momentum"]
            eeta = group["eeta"]
            lr = group["lr"]
            lars = group["lars"]
            eps = group["epsilon"]

            for p in group["params"]:
                if p.grad is None:
                    continue
                decayed_grad = p.grad
                scaled_lr = lr
                if lars:
                    w_norm = torch.linalg.norm(p)
                    g_norm = torch.linalg.norm(p.grad)
                    trust_ratio = torch.where(
                        w_norm > 0 and g_norm > 0,
                        eeta * w_norm / (g_norm + weight_decay * w_norm + eps),
                        torch.ones_like(w_norm),
                    )
                    trust_ratio.clamp_(0.0, 50)
                    scaled_lr *= trust_ratio.item()
                    if weight_decay != 0:
                        decayed_grad = decayed_grad.add(p, alpha=weight_decay)
                decayed_grad = torch.clamp(decayed_grad, -10.0, 10.0)

                if momentum != 0:
                    param_state = self.state[p]
                    if "momentum_buffer" not in param_state:
                        buf = param_state["momentum_buffer"] = torch.clone(decayed_grad).detach()
                    else:
                        buf = param_state["momentum_buffer"]
                        buf.mul_(momentum).add_(decayed_grad)
                    decayed_grad = buf

                p.add_(decayed_grad, alpha=-scaled_lr)

        return loss
//  [python]
def photo_level_tracking(
    Tji_init, aff_init, vals_i, Pi, dI_dT, img_j, intrinsics, photo_sigma, term_criteria
):
    Tji = Tji_init.clone()
    aff = aff_init.clone()

    A_norm = 1.0 / torch.as_tensor(
        (img_j.shape[-1], img_j.shape[-2]), device=img_j.device, dtype=img_j.dtype
    )

    iter = 0
    done = False
    mean_sq_err_prev = float("inf")
    while not done:
        Tji, aff, delta, mean_sq_err, grad_norm, p_j, valid_reproj_mask, depth_j = (
            tracking_iter(
                Tji, Pi, intrinsics, img_j, aff, vals_i, dI_dT, photo_sigma, A_norm
            )
        )

        iter += 1
        delta_norm = torch.norm(delta)
        abs_decrease = mean_sq_err_prev - mean_sq_err

        # NOTE: Checking for convergence, not if error goes up, so want absolute value!
        rel_decrease = torch.abs(abs_decrease / mean_sq_err_prev)

        # print("Tracking: ", iter, mean_sq_err.item(), delta_norm.item(), rel_decrease.item(), grad_norm.item())
        if (
            iter >= term_criteria["max_iter"]
            or delta_norm < term_criteria["delta_norm"]
            or rel_decrease < term_criteria["rel_tol"]
            or grad_norm < term_criteria["grad_norm"]
        ):
            done = True
            # print(iter, abs_decrease, delta_norm, rel_decrease)
        mean_sq_err_prev = mean_sq_err

    return Tji, aff
//  [python]
def photo_level_tracking(
    Tji_init, aff_init, vals_i, Pi, dI_dT, img_j, intrinsics, photo_sigma, term_criteria
):
    Tji = Tji_init.clone()
    aff = aff_init.clone()

    A_norm = 1.0 / torch.as_tensor(
        (img_j.shape[-1], img_j.shape[-2]), device=img_j.device, dtype=img_j.dtype
    )

    iter = 0
    done = False
    mean_sq_err_prev = float("inf")
    while not done:
        Tji, aff, delta, mean_sq_err, grad_norm, p_j, valid_reproj_mask, depth_j = (
            tracking_iter(
                Tji, Pi, intrinsics, img_j, aff, vals_i, dI_dT, photo_sigma, A_norm
            )
        )

        iter += 1
        delta_norm = torch.linalg.norm(delta)
        abs_decrease = mean_sq_err_prev - mean_sq_err

        # NOTE: Checking for convergence, not if error goes up, so want absolute value!
        rel_decrease = torch.abs(abs_decrease / mean_sq_err_prev)

        # print("Tracking: ", iter, mean_sq_err.item(), delta_norm.item(), rel_decrease.item(), grad_norm.item())
        if (
            iter >= term_criteria["max_iter"]
            or delta_norm < term_criteria["delta_norm"]
            or rel_decrease < term_criteria["rel_tol"]
            or grad_norm < term_criteria["grad_norm"]
        ):
            done = True
            # print(iter, abs_decrease, delta_norm, rel_decrease)
        mean_sq_err_prev = mean_sq_err

    return Tji, aff
//  [python]
def run_model_test(
    trainer_options,
    model: LightningModule,
    data: LightningDataModule = None,
    version=None,
    with_hpc: bool = True,
    min_acc: float = 0.25,
    min_change_ratio: float = 0.03,
):
    save_dir = trainer_options["default_root_dir"]

    # logger file to get meta
    logger = get_default_logger(save_dir, version=version)
    trainer_options.update(logger=logger)
    trainer = Trainer(**trainer_options)
    with torch.no_grad():
        initial_values = torch.cat([x.view(-1) for x in model.parameters()])
    trainer.fit(model, datamodule=data)
    with torch.no_grad():
        post_train_values = torch.cat([x.view(-1) for x in model.parameters()])

    # Check that the model has changed post-training
    change_ratio = torch.norm(initial_values - post_train_values) / torch.norm(initial_values)
    assert change_ratio >= min_change_ratio, (
        f"The change in the model's parameter norm is {change_ratio:.1f}"
        f" relative to the initial norm, but expected a change by >={min_change_ratio}"
    )

    if trainer.world_size != trainer.num_devices:
        # we're in multinode. unless the filesystem is shared, only the main node will have access to the checkpoint
        # since we cannot know this, the code below needs to be skipped
        return

    # test model loading
    _ = load_model_from_checkpoint(trainer.checkpoint_callback.best_model_path, type(model))

    # test new model accuracy
    test_loaders = model.test_dataloader() if not data else data.test_dataloader()
    if not isinstance(test_loaders, list):
        test_loaders = [test_loaders]

    if not isinstance(model, BoringModel):
        for dataloader in test_loaders:
            run_model_prediction(model, dataloader, min_acc=min_acc)

    if with_hpc:
        # test HPC saving
        # save logger to make sure we get all the metrics
        if logger:
            logger.finalize("finished")
        hpc_save_path = trainer._checkpoint_connector.hpc_save_path(save_dir)
        trainer.save_checkpoint(hpc_save_path)
        # test HPC loading
        checkpoint_path = trainer._checkpoint_connector._CheckpointConnector__get_max_ckpt_path_from_folder(save_dir)
        trainer._checkpoint_connector.restore(checkpoint_path)
//  [python]
def run_model_test(
    trainer_options,
    model: LightningModule,
    data: LightningDataModule = None,
    version=None,
    with_hpc: bool = True,
    min_acc: float = 0.25,
    min_change_ratio: float = 0.03,
):
    save_dir = trainer_options["default_root_dir"]

    # logger file to get meta
    logger = get_default_logger(save_dir, version=version)
    trainer_options.update(logger=logger)
    trainer = Trainer(**trainer_options)
    with torch.no_grad():
        initial_values = torch.cat([x.view(-1) for x in model.parameters()])
    trainer.fit(model, datamodule=data)
    with torch.no_grad():
        post_train_values = torch.cat([x.view(-1) for x in model.parameters()])

    # Check that the model has changed post-training
    change_ratio = torch.linalg.norm(initial_values - post_train_values) / torch.linalg.norm(initial_values)
    assert change_ratio >= min_change_ratio, (
        f"The change in the model's parameter norm is {change_ratio:.1f}"
        f" relative to the initial norm, but expected a change by >={min_change_ratio}"
    )

    if trainer.world_size != trainer.num_devices:
        # we're in multinode. unless the filesystem is shared, only the main node will have access to the checkpoint
        # since we cannot know this, the code below needs to be skipped
        return

    # test model loading
    _ = load_model_from_checkpoint(trainer.checkpoint_callback.best_model_path, type(model))

    # test new model accuracy
    test_loaders = model.test_dataloader() if not data else data.test_dataloader()
    if not isinstance(test_loaders, list):
        test_loaders = [test_loaders]

    if not isinstance(model, BoringModel):
        for dataloader in test_loaders:
            run_model_prediction(model, dataloader, min_acc=min_acc)

    if with_hpc:
        # test HPC saving
        # save logger to make sure we get all the metrics
        if logger:
            logger.finalize("finished")
        hpc_save_path = trainer._checkpoint_connector.hpc_save_path(save_dir)
        trainer.save_checkpoint(hpc_save_path)
        # test HPC loading
        checkpoint_path = trainer._checkpoint_connector._CheckpointConnector__get_max_ckpt_path_from_folder(save_dir)
        trainer._checkpoint_connector.restore(checkpoint_path)
//  [python]
def get_kmeans_mu(self, x, n_centers, init_times=50, min_delta=1e-3):
        """
        Find an initial value for the mean. Requires a threshold min_delta for the k-means algorithm to stop iterating.
        The algorithm is repeated init_times often, after which the best centerpoint is returned.
        args:
            x:            torch.FloatTensor (n, d) or (n, 1, d)
            init_times:   init
            min_delta:    int
        """
        if len(x.size()) == 3:
            x = x.squeeze(1)
        x_min, x_max = x.min(), x.max()
        x = (x - x_min) / (x_max - x_min)

        min_cost = np.inf

        for i in range(init_times):
            tmp_center = x[np.random.choice(np.arange(x.shape[0]), size=n_centers, replace=False), ...]
            l2_dis = torch.norm((x.unsqueeze(1).repeat(1, n_centers, 1) - tmp_center), p=2, dim=2)
            l2_cls = torch.argmin(l2_dis, dim=1)

            cost = 0
            for c in range(n_centers):
                cost += torch.norm(x[l2_cls == c] - tmp_center[c], p=2, dim=1).mean()

            if cost < min_cost:
                min_cost = cost
                center = tmp_center

        delta = np.inf

        while delta > min_delta:
            l2_dis = torch.norm((x.unsqueeze(1).repeat(1, n_centers, 1) - center), p=2, dim=2)
            l2_cls = torch.argmin(l2_dis, dim=1)
            center_old = center.clone()

            for c in range(n_centers):
                center[c] = x[l2_cls == c].mean(dim=0)

            delta = torch.norm((center_old - center), dim=1).max()

        return (center.unsqueeze(0) * (x_max - x_min) + x_min)