jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_franka_kinematics(cfg):
    tensor_args = TensorDeviceType()

    robot_model = CudaRobotModel(cfg)
    q_test = torch.as_tensor(
        [0.0, -1.2, 0.0, -2.0, 0.0, 1.0, 0.0], **(tensor_args.as_torch_dict())
    ).view(1, -1)
    ee_position = torch.as_tensor(
        [6.0860e-02, -4.7547e-12, 7.6373e-01], **(tensor_args.as_torch_dict())
    ).view(1, -1)
    ee_quat = torch.as_tensor(
        [0.0382, 0.9193, 0.3808, 0.0922], **(tensor_args.as_torch_dict())
    ).view(1, -1)
    b_list = [1, 10, 100, 5000][:1]
    for b in b_list:
        state = robot_model.get_state(q_test.repeat(b, 1).clone())
        pos_err = torch.norm(state.ee_position - ee_position)
        q_err = torch.norm(state.ee_quaternion - ee_quat)
        # check if all values are equal to position and quaternion
        assert pos_err < 1e-3
        assert q_err < 1e-1
//  [python]
def test_franka_kinematics(cfg):
    tensor_args = TensorDeviceType()

    robot_model = CudaRobotModel(cfg)
    q_test = torch.as_tensor(
        [0.0, -1.2, 0.0, -2.0, 0.0, 1.0, 0.0], **(tensor_args.as_torch_dict())
    ).view(1, -1)
    ee_position = torch.as_tensor(
        [6.0860e-02, -4.7547e-12, 7.6373e-01], **(tensor_args.as_torch_dict())
    ).view(1, -1)
    ee_quat = torch.as_tensor(
        [0.0382, 0.9193, 0.3808, 0.0922], **(tensor_args.as_torch_dict())
    ).view(1, -1)
    b_list = [1, 10, 100, 5000][:1]
    for b in b_list:
        state = robot_model.get_state(q_test.repeat(b, 1).clone())
        pos_err = torch.linalg.norm(state.ee_position - ee_position)
        q_err = torch.linalg.norm(state.ee_quaternion - ee_quat)
        # check if all values are equal to position and quaternion
        assert pos_err < 1e-3
        assert q_err < 1e-1
//  [python]
def _evaluate_batch(self, batch: et.SolutionBatch):
        x = batch.values
        norms = torch.norm(x, dim=-1)
        batch.set_evals(norms)
//  [python]
def _evaluate_batch(self, batch: et.SolutionBatch):
        x = batch.values
        norms = torch.linalg.norm(x, dim=-1)
        batch.set_evals(norms)
//  [python]
def test_sh():
    from gsplat import _torch_impl
    from gsplat import sh

    num_points = 1
    degree = 4
    gt_colors = torch.ones(num_points, 3, device=device) * 0.5
    viewdirs = torch.randn(num_points, 3, device=device)
    viewdirs /= torch.norm(viewdirs, dim=-1, keepdim=True)
    sh_coeffs = torch.rand(
        num_points, sh.num_sh_bases(degree), 3, device=device, requires_grad=True
    )
    optim = torch.optim.Adam([sh_coeffs], lr=1e-2)

    num_iters = 1000
    for _ in range(num_iters):
        optim.zero_grad()

        # compute PyTorch's color and grad
        check_colors = _torch_impl.compute_sh_color(viewdirs, sh_coeffs)
        check_loss = torch.square(check_colors - gt_colors).mean()
        check_loss.backward()
        check_grad = sh_coeffs.grad.detach()

        optim.zero_grad()

        # compute our colors and grads
        colors = sh.spherical_harmonics(degree, viewdirs, sh_coeffs)
        loss = torch.square(colors - gt_colors).mean()
        loss.backward()
        grad = sh_coeffs.grad.detach()
        optim.step()

        torch.testing.assert_close(check_grad, grad)
        torch.testing.assert_close(check_colors, colors)

    # check final optimized color
    torch.testing.assert_close(check_colors, gt_colors)
//  [python]
def test_sh():
    from gsplat import _torch_impl
    from gsplat import sh

    num_points = 1
    degree = 4
    gt_colors = torch.ones(num_points, 3, device=device) * 0.5
    viewdirs = torch.randn(num_points, 3, device=device)
    viewdirs /= torch.linalg.norm(viewdirs, dim=-1, keepdim=True)
    sh_coeffs = torch.rand(
        num_points, sh.num_sh_bases(degree), 3, device=device, requires_grad=True
    )
    optim = torch.optim.Adam([sh_coeffs], lr=1e-2)

    num_iters = 1000
    for _ in range(num_iters):
        optim.zero_grad()

        # compute PyTorch's color and grad
        check_colors = _torch_impl.compute_sh_color(viewdirs, sh_coeffs)
        check_loss = torch.square(check_colors - gt_colors).mean()
        check_loss.backward()
        check_grad = sh_coeffs.grad.detach()

        optim.zero_grad()

        # compute our colors and grads
        colors = sh.spherical_harmonics(degree, viewdirs, sh_coeffs)
        loss = torch.square(colors - gt_colors).mean()
        loss.backward()
        grad = sh_coeffs.grad.detach()
        optim.step()

        torch.testing.assert_close(check_grad, grad)
        torch.testing.assert_close(check_colors, colors)

    # check final optimized color
    torch.testing.assert_close(check_colors, gt_colors)
//  [python]
def _norm(self, X, G):
        L_inv = torch.inverse(torch.cholesky(X))
        return torch.norm(torch.mm(torch.mm(L_inv, G), L_inv.t()), 'fro')
//  [python]
def _norm(self, X, G):
        L_inv = torch.inverse(torch.cholesky(X))
        return torch.linalg.norm(torch.mm(torch.mm(L_inv, G), L_inv.t()), ord='fro')
//  [python]
def _p2srt(
        self, param: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Applies matrix norm to the parameter vector.

        Args:
            param (torch.Tensor): Parameter vector.

        Returns:
            Tuple[torch.Tensor, torch.Tensor, torch.Tensor]
        """
        t3d = param[:, 3]
        r1 = param[0:1, :3]
        r2 = param[1:2, :3]
        se = (t.norm(r1) + t.norm(r2)) / 2.0
        r1 = r1 / t.norm(r1)
        r2 = r2 / t.norm(r2)
        r3 = torch.cross(r1, r2)
        re = torch.cat((r1, r2, r3), 0)
        return se, re, t3d
//  [python]
def _p2srt(
        self, param: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Applies matrix norm to the parameter vector.

        Args:
            param (torch.Tensor): Parameter vector.

        Returns:
            Tuple[torch.Tensor, torch.Tensor, torch.Tensor]
        """
        t3d = param[:, 3]
        r1 = param[0:1, :3]
        r2 = param[1:2, :3]
        se = (torch.linalg.norm(r1) + torch.linalg.norm(r2)) / 2.0
        r1 = r1 / torch.linalg.norm(r1)
        r2 = r2 / torch.linalg.norm(r2)
        r3 = torch.cross(r1, r2)
        re = torch.cat((r1, r2, r3), 0)
        return se, re, t3d
//  [python]
def test_locked_joints_kinematics():
    tensor_args = TensorDeviceType()

    config_file = load_yaml(join_path(get_robot_configs_path(), "franka.yml"))

    cfg = CudaRobotModelConfig.from_config(
        CudaRobotGeneratorConfig(**config_file["robot_cfg"]["kinematics"], tensor_args=tensor_args)
    )
    robot_model = CudaRobotModel(cfg)

    cspace = CSpaceConfig(**config_file["robot_cfg"]["kinematics"]["cspace"])
    cspace.inplace_reindex(robot_model.joint_names)
    q = cspace.retract_config.unsqueeze(0)

    out = robot_model.get_state(q)
    j_idx = 2
    lock_joint_name = cspace.joint_names[j_idx]
    # lock a joint:
    cfg.lock_joints = {lock_joint_name: cspace.retract_config[j_idx].item()}
    cfg = CudaRobotModelConfig.from_config(
        CudaRobotGeneratorConfig(**config_file["robot_cfg"]["kinematics"], tensor_args=tensor_args)
    )

    robot_model = CudaRobotModel(cfg)
    cspace = CSpaceConfig(**config_file["robot_cfg"]["kinematics"]["cspace"])

    cspace.inplace_reindex(robot_model.joint_names)
    q = cspace.retract_config.unsqueeze(0).clone()
    out_locked = robot_model.get_state(q)
    assert torch.norm(out.ee_position - out_locked.ee_position) < 1e-5
    assert torch.norm(out.ee_quaternion - out_locked.ee_quaternion) < 1e-5
    assert torch.norm(out.link_spheres_tensor - out_locked.link_spheres_tensor) < 1e-5
//  [python]
def test_locked_joints_kinematics():
    tensor_args = TensorDeviceType()

    config_file = load_yaml(join_path(get_robot_configs_path(), "franka.yml"))

    cfg = CudaRobotModelConfig.from_config(
        CudaRobotGeneratorConfig(**config_file["robot_cfg"]["kinematics"], tensor_args=tensor_args)
    )
    robot_model = CudaRobotModel(cfg)

    cspace = CSpaceConfig(**config_file["robot_cfg"]["kinematics"]["cspace"])
    cspace.inplace_reindex(robot_model.joint_names)
    q = cspace.retract_config.unsqueeze(0)

    out = robot_model.get_state(q)
    j_idx = 2
    lock_joint_name = cspace.joint_names[j_idx]
    # lock a joint:
    cfg.lock_joints = {lock_joint_name: cspace.retract_config[j_idx].item()}
    cfg = CudaRobotModelConfig.from_config(
        CudaRobotGeneratorConfig(**config_file["robot_cfg"]["kinematics"], tensor_args=tensor_args)
    )

    robot_model = CudaRobotModel(cfg)
    cspace = CSpaceConfig(**config_file["robot_cfg"]["kinematics"]["cspace"])

    cspace.inplace_reindex(robot_model.joint_names)
    q = cspace.retract_config.unsqueeze(0).clone()
    out_locked = robot_model.get_state(q)
    assert torch.linalg.norm(out.ee_position - out_locked.ee_position) < 1e-5
    assert torch.linalg.norm(out.ee_quaternion - out_locked.ee_quaternion) < 1e-5
    assert torch.linalg.norm(out.link_spheres_tensor - out_locked.link_spheres_tensor) < 1e-5
//  [python]
def wsd2(mu_x,cov_x,mu_y,cov_y, epsilon=torch.tensor(1e-5), device = 'cpu'):
    
    '''
    d(P_X,P_Y) = min_{X,Y marginals}E[|X-Y|^2], that for normally distributed variables becomes
    d = |mu_X-mu_Y|^2 + trace(cov_X+cov_Y -2(cov_X*cov_Y)^(1/2))
    
    Inspired by: https://gist.github.com/Flunzmas/6e359b118b0730ab403753dcc2a447df
    and the paper https://arxiv.org/pdf/2009.14075.pdf
    
    '''
    epsilon=epsilon.to(device)
    
    mu_part = torch.norm(mu_x-mu_y)
    mm = torch.matmul(cov_x, cov_y)
    eye = torch.eye(mm.shape[0]).to(device)
    eps = eye*epsilon
    eps = eps.to(device)
    mm = mm+eps
    s = torch.linalg.eigvals(mm).real
    sq_cov_prod_trace = torch.sqrt(torch.clamp(s,min=1e-5,max = 1e+5)).sum()
    
    trace_part = torch.trace(cov_x+cov_y) - 2*sq_cov_prod_trace
    
    return mu_part + trace_part
//  [python]
def wsd2(mu_x,cov_x,mu_y,cov_y, epsilon=torch.tensor(1e-5), device = 'cpu'):
    
    '''
    d(P_X,P_Y) = min_{X,Y marginals}E[|X-Y|^2], that for normally distributed variables becomes
    d = |mu_X-mu_Y|^2 + trace(cov_X+cov_Y -2(cov_X*cov_Y)^(1/2))
    
    Inspired by: https://gist.github.com/Flunzmas/6e359b118b0730ab403753dcc2a447df
    and the paper https://arxiv.org/pdf/2009.14075.pdf
    
    '''
    epsilon=epsilon.to(device)
    
    mu_part = torch.linalg.norm(mu_x-mu_y)
    mm = torch.matmul(cov_x, cov_y)
    eye = torch.eye(mm.shape[0]).to(device)
    eps = eye*epsilon
    eps = eps.to(device)
    mm = mm+eps
    s = torch.linalg.eigvals(mm).real
    sq_cov_prod_trace = torch.sqrt(torch.clamp(s,min=1e-5,max = 1e+5)).sum()
    
    trace_part = torch.trace(cov_x+cov_y) - 2*sq_cov_prod_trace
    
    return mu_part + trace_part
//  [python]
def SE3_logmap(T, eps=1e-6):
  w = SO3_logmap(T[:,:3,:3])
  theta = torch.norm(w,dim=(1))
  theta = torch.clamp(theta, min=eps)
  w_norm = w/theta
  tan = torch.tan(0.5*theta)

  t = T[:,:3,3]
  wnorm_x_t = torch.cross(w_norm,t)
  V_inv_t = t - (0.5 * t) * wnorm_x_t + (1.0 - theta/(2.0*tan)) * torch.cross(w_norm, wnorm_x_t)
  xi = torch.cat((w, V_inv_t),dim=-1)

  return xi
//  [python]
def SE3_logmap(T, eps=1e-6):
  w = SO3_logmap(T[:,:3,:3])
  theta = torch.linalg.norm(w,dim=(1))
  theta = torch.clamp(theta, min=eps)
  w_norm = w/theta
  tan = torch.tan(0.5*theta)

  t = T[:,:3,3]
  wnorm_x_t = torch.cross(w_norm,t)
  V_inv_t = t - (0.5 * t) * wnorm_x_t + (1.0 - theta/(2.0*tan)) * torch.cross(w_norm, wnorm_x_t)
  xi = torch.cat((w, V_inv_t),dim=-1)

  return xi
//  [python]
def parse_ddim_diffusion(x, seq, model, b, cls_fn=None, rho_scale=1.0, stop=100, ref_path=None):
    parser = FaceParseTool(ref_path=ref_path).cuda()

    # setup iteration variables
    n = x.size(0)
    seq_next = [-1] + list(seq[:-1])
    x0_preds = []
    xs = [x]

    # iterate over the timesteps
    for i, j in tqdm(zip(reversed(seq), reversed(seq_next))):
        t = (torch.ones(n) * i).to(x.device)
        next_t = (torch.ones(n) * j).to(x.device)
        at = compute_alpha(b, t.long())
        at_next = compute_alpha(b, next_t.long())
        xt = xs[-1].to('cuda')
        
        xt.requires_grad = True
        
        if cls_fn == None:
            et = model(xt, t)
        else:
            print("use class_num")
            class_num = 281
            classes = torch.ones(xt.size(0), dtype=torch.long, device=torch.device("cuda"))*class_num
            et = model(xt, t, classes)
            et = et[:, :3]
            et = et - (1 - at).sqrt()[0, 0, 0, 0] * cls_fn(x, t, classes)

        if et.size(1) == 6:
            et = et[:, :3]

        x0_t = (xt - et * (1 - at).sqrt()) / at.sqrt()
        
        residual = parser.get_residual(x0_t)
        norm = torch.norm(residual)
        norm_grad = torch.autograd.grad(outputs=norm, inputs=xt)[0]

        
        eta = 0.5
        c1 = (1 - at_next).sqrt() * eta
        c2 = (1 - at_next).sqrt() * ((1 - eta ** 2) ** 0.5)
        xt_next = at_next.sqrt() * x0_t + c1 * torch.randn_like(x0_t) + c2 * et

        # use guided gradient
        rho = at.sqrt() * rho_scale
        if not i <= stop:
            xt_next -= rho * norm_grad
        
        x0_t = x0_t.detach()
        xt_next = xt_next.detach()
        
        x0_preds.append(x0_t.to('cpu'))
        xs.append(xt_next.to('cpu'))

    # return x0_preds, xs
    return [xs[-1]], [x0_preds[-1]]
//  [python]
def parse_ddim_diffusion(x, seq, model, b, cls_fn=None, rho_scale=1.0, stop=100, ref_path=None):
    parser = FaceParseTool(ref_path=ref_path).cuda()

    # setup iteration variables
    n = x.size(0)
    seq_next = [-1] + list(seq[:-1])
    x0_preds = []
    xs = [x]

    # iterate over the timesteps
    for i, j in tqdm(zip(reversed(seq), reversed(seq_next))):
        t = (torch.ones(n) * i).to(x.device)
        next_t = (torch.ones(n) * j).to(x.device)
        at = compute_alpha(b, t.long())
        at_next = compute_alpha(b, next_t.long())
        xt = xs[-1].to('cuda')
        
        xt.requires_grad = True
        
        if cls_fn == None:
            et = model(xt, t)
        else:
            print("use class_num")
            class_num = 281
            classes = torch.ones(xt.size(0), dtype=torch.long, device=torch.device("cuda"))*class_num
            et = model(xt, t, classes)
            et = et[:, :3]
            et = et - (1 - at).sqrt()[0, 0, 0, 0] * cls_fn(x, t, classes)

        if et.size(1) == 6:
            et = et[:, :3]

        x0_t = (xt - et * (1 - at).sqrt()) / at.sqrt()
        
        residual = parser.get_residual(x0_t)
        norm = torch.linalg.norm(residual)
        norm_grad = torch.autograd.grad(outputs=norm, inputs=xt)[0]

        
        eta = 0.5
        c1 = (1 - at_next).sqrt() * eta
        c2 = (1 - at_next).sqrt() * ((1 - eta ** 2) ** 0.5)
        xt_next = at_next.sqrt() * x0_t + c1 * torch.randn_like(x0_t) + c2 * et

        # use guided gradient
        rho = at.sqrt() * rho_scale
        if not i <= stop:
            xt_next -= rho * norm_grad
        
        x0_t = x0_t.detach()
        xt_next = xt_next.detach()
        
        x0_preds.append(x0_t.to('cpu'))
        xs.append(xt_next.to('cpu'))

    # return x0_preds, xs
    return [xs[-1]], [x0_preds[-1]]
//  [python]
def sketch_ddim_diffusion(x, seq, model, b, cls_fn=None, rho_scale=1.0, stop=100, ref_path=None):
    img2sketch = FaceSketchTool(ref_path=ref_path).cuda()

    # setup iteration variables
    n = x.size(0)
    seq_next = [-1] + list(seq[:-1])
    x0_preds = []
    xs = [x]

    # iterate over the timesteps
    for i, j in tqdm(zip(reversed(seq), reversed(seq_next))):
        t = (torch.ones(n) * i).to(x.device)
        next_t = (torch.ones(n) * j).to(x.device)
        at = compute_alpha(b, t.long())
        at_next = compute_alpha(b, next_t.long())
        xt = xs[-1].to('cuda')
        
        xt.requires_grad = True
        
        if cls_fn == None:
            et = model(xt, t)
        else:
            # print("use class_num")
            class_num = 7
            classes = torch.ones(xt.size(0), dtype=torch.long, device=torch.device("cuda"))*class_num
            et = model(xt, t, classes)
            et = et[:, :3]
            et = et - (1 - at).sqrt()[0, 0, 0, 0] * cls_fn(x, t, classes)

        if et.size(1) == 6:
            et = et[:, :3]

        x0_t = (xt - et * (1 - at).sqrt()) / at.sqrt()
        
        residual = img2sketch.get_residual(x0_t)
        norm = torch.norm(residual)
        norm_grad = torch.autograd.grad(outputs=norm, inputs=xt)[0]
        
        eta = 0.5
        c1 = (1 - at_next).sqrt() * eta
        c2 = (1 - at_next).sqrt() * ((1 - eta ** 2) ** 0.5)
        xt_next = at_next.sqrt() * x0_t + c1 * torch.randn_like(x0_t) + c2 * et
        
        # use guided gradient
        rho = at.sqrt() * rho_scale
        if not i <= stop:
            xt_next -= rho * norm_grad
        
        x0_t = x0_t.detach()
        xt_next = xt_next.detach()
        
        x0_preds.append(x0_t.to('cpu'))
        xs.append(xt_next.to('cpu'))

    return [xs[-1]], [x0_preds[-1]]
//  [python]
def sketch_ddim_diffusion(x, seq, model, b, cls_fn=None, rho_scale=1.0, stop=100, ref_path=None):
    img2sketch = FaceSketchTool(ref_path=ref_path).cuda()

    # setup iteration variables
    n = x.size(0)
    seq_next = [-1] + list(seq[:-1])
    x0_preds = []
    xs = [x]

    # iterate over the timesteps
    for i, j in tqdm(zip(reversed(seq), reversed(seq_next))):
        t = (torch.ones(n) * i).to(x.device)
        next_t = (torch.ones(n) * j).to(x.device)
        at = compute_alpha(b, t.long())
        at_next = compute_alpha(b, next_t.long())
        xt = xs[-1].to('cuda')
        
        xt.requires_grad = True
        
        if cls_fn == None:
            et = model(xt, t)
        else:
            # print("use class_num")
            class_num = 7
            classes = torch.ones(xt.size(0), dtype=torch.long, device=torch.device("cuda"))*class_num
            et = model(xt, t, classes)
            et = et[:, :3]
            et = et - (1 - at).sqrt()[0, 0, 0, 0] * cls_fn(x, t, classes)

        if et.size(1) == 6:
            et = et[:, :3]

        x0_t = (xt - et * (1 - at).sqrt()) / at.sqrt()
        
        residual = img2sketch.get_residual(x0_t)
        norm = torch.linalg.norm(residual)
        norm_grad = torch.autograd.grad(outputs=norm, inputs=xt)[0]
        
        eta = 0.5
        c1 = (1 - at_next).sqrt() * eta
        c2 = (1 - at_next).sqrt() * ((1 - eta ** 2) ** 0.5)
        xt_next = at_next.sqrt() * x0_t + c1 * torch.randn_like(x0_t) + c2 * et
        
        # use guided gradient
        rho = at.sqrt() * rho_scale
        if not i <= stop:
            xt_next -= rho * norm_grad
        
        x0_t = x0_t.detach()
        xt_next = xt_next.detach()
        
        x0_preds.append(x0_t.to('cpu'))
        xs.append(xt_next.to('cpu'))

    return [xs[-1]], [x0_preds[-1]]
//  [python]
def plan_path_mpnet(model, enc, start, goal, device, collisionFn, worldBounds):
    '''
    Plan a path using MPNet.
    :param model: The NN model used for planning.
    :param enc: The latent space representation of the map.
    :param start: A numpy array of size 2 with the normalized start state.
    :param goal: A numpy array of size 2 with the normalized goal state.
    :param device: The device used for torch computations.
    :param collisionFn: A function to check if a state is in collision
    :param worldBounds: A numpy array of worldBounds
    :returns np.array: A set of scaled trajectories
    '''
    goalS = torch.tensor(normalize_state(goal, worldBounds)[None,:], device=device, dtype=torch.float)
    startS = torch.tensor(normalize_state(start, worldBounds)[None,:], device=device, dtype=torch.float)
    reachedGoal = False
    normPredTraj = [startS.cpu().numpy().squeeze()]
    vertex = 0
    for _ in range(10):
        # Check if we can connect start and goal state
        if not check_edge_collision(scale_state(startS.cpu().numpy().squeeze()), goal, collisionFn):
            reachedGoal = True
            break
        inputs = torch.cat([enc, startS, goalS], dim=1)
        for _ in range(10):
            with torch.no_grad():
                tempS = model(inputs)
                temp = tempS.cpu().numpy().squeeze()
            if collisionFn(scale_state(temp)):
                vertex +=1                
                startS = tempS
                normPredTraj.append(temp)
                break
        # Check if we can connect to the goal or reached near it
        if torch.norm(startS-goalS)*24<0.1:
            reachedGoal = True
            break
    if reachedGoal:
        normPredTraj.append(normalize_state(goal, worldBounds))
    return scale_state(np.array(normPredTraj)), reachedGoal, vertex
//  [python]
def plan_path_mpnet(model, enc, start, goal, device, collisionFn, worldBounds):
    '''
    Plan a path using MPNet.
    :param model: The NN model used for planning.
    :param enc: The latent space representation of the map.
    :param start: A numpy array of size 2 with the normalized start state.
    :param goal: A numpy array of size 2 with the normalized goal state.
    :param device: The device used for torch computations.
    :param collisionFn: A function to check if a state is in collision
    :param worldBounds: A numpy array of worldBounds
    :returns np.array: A set of scaled trajectories
    '''
    goalS = torch.tensor(normalize_state(goal, worldBounds)[None,:], device=device, dtype=torch.float)
    startS = torch.tensor(normalize_state(start, worldBounds)[None,:], device=device, dtype=torch.float)
    reachedGoal = False
    normPredTraj = [startS.cpu().numpy().squeeze()]
    vertex = 0
    for _ in range(10):
        # Check if we can connect start and goal state
        if not check_edge_collision(scale_state(startS.cpu().numpy().squeeze()), goal, collisionFn):
            reachedGoal = True
            break
        inputs = torch.cat([enc, startS, goalS], dim=1)
        for _ in range(10):
            with torch.no_grad():
                tempS = model(inputs)
                temp = tempS.cpu().numpy().squeeze()
            if collisionFn(scale_state(temp)):
                vertex +=1                
                startS = tempS
                normPredTraj.append(temp)
                break
        # Check if we can connect to the goal or reached near it
        if torch.linalg.norm(startS-goalS)*24<0.1:
            reachedGoal = True
            break
    if reachedGoal:
        normPredTraj.append(normalize_state(goal, worldBounds))
    return scale_state(np.array(normPredTraj)), reachedGoal, vertex
//  [python]
def clip_ddim_diffusion(x, seq, model, b, cls_fn=None, rho_scale=1.0, prompt=None, stop=100, domain="face"):
    clip_encoder = CLIPEncoder().cuda()

    # setup iteration variables
    n = x.size(0)
    seq_next = [-1] + list(seq[:-1])
    x0_preds = []
    xs = [x]

    # iterate over the timesteps
    for i, j in tqdm(zip(reversed(seq), reversed(seq_next))):
        t = (torch.ones(n) * i).to(x.device)
        next_t = (torch.ones(n) * j).to(x.device)
        at = compute_alpha(b, t.long())
        at_next = compute_alpha(b, next_t.long())
        xt = xs[-1].to('cuda')

        if domain == "face":
            repeat = 1
        elif domain == "imagenet":
            if 800 >= i >= 500:
                repeat = 10
            else:
                repeat = 1
        
        for idx in range(repeat):
        
            xt.requires_grad = True
            
            et = model(xt, t)

            if et.size(1) == 6:
                et = et[:, :3]

            x0_t = (xt - et * (1 - at).sqrt()) / at.sqrt()
            
            # get guided gradient
            residual = clip_encoder.get_residual(x0_t, prompt)
            norm = torch.norm(residual)
            norm_grad = torch.autograd.grad(outputs=norm, inputs=xt)[0]

            c1 = at_next.sqrt() * (1 - at / at_next) / (1 - at)
            c2 = (at / at_next).sqrt() * (1 - at_next) / (1 - at)
            c3 = (1 - at_next) * (1 - at / at_next) / (1 - at)
            c3 = (c3.log() * 0.5).exp()
            xt_next = c1 * x0_t + c2 * xt + c3 * torch.randn_like(x0_t)
            
            l1 = ((et * et).mean().sqrt() * (1 - at).sqrt() / at.sqrt() * c1).item()
            l2 = l1 * 0.02
            rho = l2 / (norm_grad * norm_grad).mean().sqrt().item()
            
            xt_next -= rho * norm_grad
            
            x0_t = x0_t.detach()
            xt_next = xt_next.detach()
            
            x0_preds.append(x0_t.to('cpu'))
            xs.append(xt_next.to('cpu'))

            if idx + 1 < repeat:
                bt = at / at_next
                xt = bt.sqrt() * xt_next + (1 - bt).sqrt() * torch.randn_like(xt_next)

    # return x0_preds, xs
    return [xs[-1]], [x0_preds[-1]]
//  [python]
def clip_ddim_diffusion(x, seq, model, b, cls_fn=None, rho_scale=1.0, prompt=None, stop=100, domain="face"):
    clip_encoder = CLIPEncoder().cuda()

    # setup iteration variables
    n = x.size(0)
    seq_next = [-1] + list(seq[:-1])
    x0_preds = []
    xs = [x]

    # iterate over the timesteps
    for i, j in tqdm(zip(reversed(seq), reversed(seq_next))):
        t = (torch.ones(n) * i).to(x.device)
        next_t = (torch.ones(n) * j).to(x.device)
        at = compute_alpha(b, t.long())
        at_next = compute_alpha(b, next_t.long())
        xt = xs[-1].to('cuda')

        if domain == "face":
            repeat = 1
        elif domain == "imagenet":
            if 800 >= i >= 500:
                repeat = 10
            else:
                repeat = 1
        
        for idx in range(repeat):
        
            xt.requires_grad = True
            
            et = model(xt, t)

            if et.size(1) == 6:
                et = et[:, :3]

            x0_t = (xt - et * (1 - at).sqrt()) / at.sqrt()
            
            # get guided gradient
            residual = clip_encoder.get_residual(x0_t, prompt)
            norm = torch.linalg.norm(residual)
            norm_grad = torch.autograd.grad(outputs=norm, inputs=xt)[0]

            c1 = at_next.sqrt() * (1 - at / at_next) / (1 - at)
            c2 = (at / at_next).sqrt() * (1 - at_next) / (1 - at)
            c3 = (1 - at_next) * (1 - at / at_next) / (1 - at)
            c3 = (c3.log() * 0.5).exp()
            xt_next = c1 * x0_t + c2 * xt + c3 * torch.randn_like(x0_t)
            
            l1 = ((et * et).mean().sqrt() * (1 - at).sqrt() / at.sqrt() * c1).item()
            l2 = l1 * 0.02
            rho = l2 / (norm_grad * norm_grad).mean().sqrt().item()
            
            xt_next -= rho * norm_grad
            
            x0_t = x0_t.detach()
            xt_next = xt_next.detach()
            
            x0_preds.append(x0_t.to('cpu'))
            xs.append(xt_next.to('cpu'))

            if idx + 1 < repeat:
                bt = at / at_next
                xt = bt.sqrt() * xt_next + (1 - bt).sqrt() * torch.randn_like(xt_next)

    # return x0_preds, xs
    return [xs[-1]], [x0_preds[-1]]
//  [python]
def landmark_ddim_diffusion(x, seq, model, b, cls_fn=None, rho_scale=1.0, stop=100, ref_path=None):
    img2landmark = FaceLandMarkTool(ref_path=ref_path).cuda()

    # setup iteration variables
    n = x.size(0)
    seq_next = [-1] + list(seq[:-1])
    x0_preds = []
    xs = [x]

    # iterate over the timesteps
    for i, j in tqdm(zip(reversed(seq), reversed(seq_next))):
        t = (torch.ones(n) * i).to(x.device)
        next_t = (torch.ones(n) * j).to(x.device)
        at = compute_alpha(b, t.long())
        at_next = compute_alpha(b, next_t.long())
        xt = xs[-1].to('cuda')
        
        xt.requires_grad = True
        
        if cls_fn == None:
            et = model(xt, t)
        else:
            print("use class_num")
            class_num = 281
            classes = torch.ones(xt.size(0), dtype=torch.long, device=torch.device("cuda"))*class_num
            et = model(xt, t, classes)
            et = et[:, :3]
            et = et - (1 - at).sqrt()[0, 0, 0, 0] * cls_fn(x, t, classes)

        if et.size(1) == 6:
            et = et[:, :3]
        
        x0_t = (xt - et * (1 - at).sqrt()) / at.sqrt()
        
        residual = img2landmark.get_residual(x0_t)
        norm = torch.norm(residual)
        norm_grad = torch.autograd.grad(outputs=norm, inputs=xt)[0]

        
        eta = 0.5
        c1 = (1 - at_next).sqrt() * eta
        c2 = (1 - at_next).sqrt() * ((1 - eta ** 2) ** 0.5)
        xt_next = at_next.sqrt() * x0_t + c1 * torch.randn_like(x0_t) + c2 * et
        
        # use guided gradient
        rho = at.sqrt() * rho_scale
        if not i <= stop:
            xt_next -= rho * norm_grad
        
        x0_t = x0_t.detach()
        xt_next = xt_next.detach()
        
        x0_preds.append(x0_t.to('cpu'))
        xs.append(xt_next.to('cpu'))

    return [xs[-1]], [x0_preds[-1]]
//  [python]
def landmark_ddim_diffusion(x, seq, model, b, cls_fn=None, rho_scale=1.0, stop=100, ref_path=None):
    img2landmark = FaceLandMarkTool(ref_path=ref_path).cuda()

    # setup iteration variables
    n = x.size(0)
    seq_next = [-1] + list(seq[:-1])
    x0_preds = []
    xs = [x]

    # iterate over the timesteps
    for i, j in tqdm(zip(reversed(seq), reversed(seq_next))):
        t = (torch.ones(n) * i).to(x.device)
        next_t = (torch.ones(n) * j).to(x.device)
        at = compute_alpha(b, t.long())
        at_next = compute_alpha(b, next_t.long())
        xt = xs[-1].to('cuda')
        
        xt.requires_grad = True
        
        if cls_fn == None:
            et = model(xt, t)
        else:
            print("use class_num")
            class_num = 281
            classes = torch.ones(xt.size(0), dtype=torch.long, device=torch.device("cuda"))*class_num
            et = model(xt, t, classes)
            et = et[:, :3]
            et = et - (1 - at).sqrt()[0, 0, 0, 0] * cls_fn(x, t, classes)

        if et.size(1) == 6:
            et = et[:, :3]
        
        x0_t = (xt - et * (1 - at).sqrt()) / at.sqrt()
        
        residual = img2landmark.get_residual(x0_t)
        norm = torch.linalg.norm(residual)
        norm_grad = torch.autograd.grad(outputs=norm, inputs=xt)[0]

        
        eta = 0.5
        c1 = (1 - at_next).sqrt() * eta
        c2 = (1 - at_next).sqrt() * ((1 - eta ** 2) ** 0.5)
        xt_next = at_next.sqrt() * x0_t + c1 * torch.randn_like(x0_t) + c2 * et
        
        # use guided gradient
        rho = at.sqrt() * rho_scale
        if not i <= stop:
            xt_next -= rho * norm_grad
        
        x0_t = x0_t.detach()
        xt_next = xt_next.detach()
        
        x0_preds.append(x0_t.to('cpu'))
        xs.append(xt_next.to('cpu'))

    return [xs[-1]], [x0_preds[-1]]
//  [python]
def arcface_ddim_diffusion(x, seq, model, b, cls_fn=None, rho_scale=1.0, stop=100, ref_path=None):
    idloss = IDLoss(ref_path=ref_path).cuda()

    # setup iteration variables
    n = x.size(0)
    seq_next = [-1] + list(seq[:-1])
    x0_preds = []
    xs = [x]

    # iterate over the timesteps
    for i, j in tqdm(zip(reversed(seq), reversed(seq_next))):
        t = (torch.ones(n) * i).to(x.device)
        next_t = (torch.ones(n) * j).to(x.device)
        at = compute_alpha(b, t.long())
        at_next = compute_alpha(b, next_t.long())
        xt = xs[-1].to('cuda')
        
        xt.requires_grad = True
        
        if cls_fn == None:
            et = model(xt, t)
        else:
            print("use class_num")
            class_num = 281
            classes = torch.ones(xt.size(0), dtype=torch.long, device=torch.device("cuda"))*class_num
            et = model(xt, t, classes)
            et = et[:, :3]
            et = et - (1 - at).sqrt()[0, 0, 0, 0] * cls_fn(x, t, classes)

        if et.size(1) == 6:
            et = et[:, :3]
        
        x0_t = (xt - et * (1 - at).sqrt()) / at.sqrt()
        
        residual = idloss.get_residual(x0_t)
        norm = torch.norm(residual)
        norm_grad = torch.autograd.grad(outputs=norm, inputs=xt)[0]

        
        eta = 0.5
        c1 = (1 - at_next).sqrt() * eta
        c2 = (1 - at_next).sqrt() * ((1 - eta ** 2) ** 0.5)
        xt_next = at_next.sqrt() * x0_t + c1 * torch.randn_like(x0_t) + c2 * et
        
        # use guided gradient
        rho = at.sqrt() * rho_scale
        if not i <= stop:
            xt_next -= rho * norm_grad
        
        x0_t = x0_t.detach()
        xt_next = xt_next.detach()
        
        x0_preds.append(x0_t.to('cpu'))
        xs.append(xt_next.to('cpu'))

    return [xs[-1]], [x0_preds[-1]]
//  [python]
def arcface_ddim_diffusion(x, seq, model, b, cls_fn=None, rho_scale=1.0, stop=100, ref_path=None):
    idloss = IDLoss(ref_path=ref_path).cuda()

    # setup iteration variables
    n = x.size(0)
    seq_next = [-1] + list(seq[:-1])
    x0_preds = []
    xs = [x]

    # iterate over the timesteps
    for i, j in tqdm(zip(reversed(seq), reversed(seq_next))):
        t = (torch.ones(n) * i).to(x.device)
        next_t = (torch.ones(n) * j).to(x.device)
        at = compute_alpha(b, t.long())
        at_next = compute_alpha(b, next_t.long())
        xt = xs[-1].to('cuda')
        
        xt.requires_grad = True
        
        if cls_fn == None:
            et = model(xt, t)
        else:
            print("use class_num")
            class_num = 281
            classes = torch.ones(xt.size(0), dtype=torch.long, device=torch.device("cuda"))*class_num
            et = model(xt, t, classes)
            et = et[:, :3]
            et = et - (1 - at).sqrt()[0, 0, 0, 0] * cls_fn(x, t, classes)

        if et.size(1) == 6:
            et = et[:, :3]
        
        x0_t = (xt - et * (1 - at).sqrt()) / at.sqrt()
        
        residual = idloss.get_residual(x0_t)
        norm = torch.linalg.norm(residual)
        norm_grad = torch.autograd.grad(outputs=norm, inputs=xt)[0]

        
        eta = 0.5
        c1 = (1 - at_next).sqrt() * eta
        c2 = (1 - at_next).sqrt() * ((1 - eta ** 2) ** 0.5)
        xt_next = at_next.sqrt() * x0_t + c1 * torch.randn_like(x0_t) + c2 * et
        
        # use guided gradient
        rho = at.sqrt() * rho_scale
        if not i <= stop:
            xt_next -= rho * norm_grad
        
        x0_t = x0_t.detach()
        xt_next = xt_next.detach()
        
        x0_preds.append(x0_t.to('cpu'))
        xs.append(xt_next.to('cpu'))

    return [xs[-1]], [x0_preds[-1]]
//  [python]
def l2_error(input_, target):
    return torch.norm(input_-target, p=2)/torch.norm(input_, p=2)
//  [python]
def l2_error(input_, target):
    return torch.linalg.norm(input_-target, ord=2)/torch.linalg.norm(input_, ord=2)
//  [python]
def compute_weight_diff(delta_weight):
        weight_diff = torch.norm(delta_weight, 2)
        return weight_diff
//  [python]
def compute_weight_diff(delta_weight):
        weight_diff = torch.linalg.norm(delta_weight, 2)
        return weight_diff
//  [python]
def normalize(x: torch.Tensor) -> Float[Tensor, "*batch"]:
    """Returns a normalized vector."""
    return x / torch.norm(x)
//  [python]
def normalize(x: torch.Tensor) -> Float[Tensor, "*batch"]:
    """Returns a normalized vector."""
    return x / torch.linalg.norm(x)
//  [python]
def is_converge(self, delta_weight, weight=None):
        weight_diff = self.compute_weight_diff(delta_weight)
        if weight is None:
            # avoid tensor[bool]
            if weight_diff < self.eps:
                return True
            return False
        if self.pre_weight is None:
            self.pre_weight = weight
            return False
        if weight_diff < self.eps * max([torch.norm(weight, 2), 1]):
            return True
        return False
//  [python]
def is_converge(self, delta_weight, weight=None):
        weight_diff = self.compute_weight_diff(delta_weight)
        if weight is None:
            # avoid tensor[bool]
            if weight_diff < self.eps:
                return True
            return False
        if self.pre_weight is None:
            self.pre_weight = weight
            return False
        if weight_diff < self.eps * max([torch.linalg.norm(weight, 2), 1]):
            return True
        return False
//  [python]
def contract_to_unisphere(
    x: torch.Tensor,
    aabb: torch.Tensor,
    ord: Union[str, int] = 2,
    #  ord: Union[float, int] = float("inf"),
    eps: float = 1e-6,
    derivative: bool = False,
):
    aabb_min, aabb_max = torch.split(aabb, 3, dim=-1)
    x = (x - aabb_min) / (aabb_max - aabb_min)
    x = x * 2 - 1  # aabb is at [-1, 1]
    mag = torch.norm(x, ord=ord, dim=-1, keepdim=True)
    mask = mag.squeeze(-1) > 1

    if derivative:
        dev = (2 * mag - 1) / mag**2 + 2 * x**2 * (
            1 / mag**3 - (2 * mag - 1) / mag**4
        )
        dev[~mask] = 1.0
        dev = torch.clamp(dev, min=eps)
        return dev
    else:
        x[mask] = (2 - 1 / mag[mask]) * (x[mask] / mag[mask])
        x = x / 4 + 0.5  # [-inf, inf] is at [0, 1]
        return x
//  [python]
def contract_to_unisphere(
    x: torch.Tensor,
    aabb: torch.Tensor,
    ord: Union[str, int] = 2,
    #  ord: Union[float, int] = float("inf"),
    eps: float = 1e-6,
    derivative: bool = False,
):
    aabb_min, aabb_max = torch.split(aabb, 3, dim=-1)
    x = (x - aabb_min) / (aabb_max - aabb_min)
    x = x * 2 - 1  # aabb is at [-1, 1]
    mag = torch.linalg.norm(x, ord=ord, dim=-1, keepdim=True)
    mask = mag.squeeze(-1) > 1

    if derivative:
        dev = (2 * mag - 1) / mag**2 + 2 * x**2 * (
            1 / mag**3 - (2 * mag - 1) / mag**4
        )
        dev[~mask] = 1.0
        dev = torch.clamp(dev, min=eps)
        return dev
    else:
        x[mask] = (2 - 1 / mag[mask]) * (x[mask] / mag[mask])
        x = x / 4 + 0.5  # [-inf, inf] is at [0, 1]
        return x
//  [python]
def _compute_side_length(self, dir_1, dir_2, dir_3):
        side_length_1 = t.norm(dir_1, dim=1).view(-1, 1, 1)
        side_length_2 = t.norm(dir_2, dim=1).view(-1, 1, 1)
        side_length_3 = t.norm(dir_3, dim=1).view(-1, 1, 1)
        return side_length_1, side_length_2, side_length_3
//  [python]
def _compute_side_length(self, dir_1, dir_2, dir_3):
        side_length_1 = torch.linalg.norm(dir_1, dim=1).view(-1, 1, 1)
        side_length_2 = torch.linalg.norm(dir_2, dim=1).view(-1, 1, 1)
        side_length_3 = torch.linalg.norm(dir_3, dim=1).view(-1, 1, 1)
        return side_length_1, side_length_2, side_length_3
//  [python]
def _compute_barycentric_grid(self, n, dir_1, dir_2, device):
        # since we have a triangle, we rmove later the points with bary.-
        # coordinates bigger one. Therefore double the number of points:
        scaled_n = 2 * n
        side_length_1 = torch.norm(dir_1, dim=1)
        side_length_2 = torch.norm(dir_2, dim=1)
        # scale the number of point w.r.t. the 'form' of the parallelogram
        n_1 = int(torch.sqrt(scaled_n*side_length_1/side_length_2))
        n_2 = int(torch.sqrt(scaled_n*side_length_2/side_length_1))
        x = torch.linspace(0, 1, n_1+2, device=device)[1:-1] # create inner grid, so remove
        y = torch.linspace(0, 1, n_2+2, device=device)[1:-1] # first and last value
        bary_coords = torch.stack(torch.meshgrid((x, y))).T.reshape(-1, 2)
        index = torch.where(bary_coords.sum(axis=1) <= 1)
        return bary_coords[index]
//  [python]
def _compute_barycentric_grid(self, n, dir_1, dir_2, device):
        # since we have a triangle, we rmove later the points with bary.-
        # coordinates bigger one. Therefore double the number of points:
        scaled_n = 2 * n
        side_length_1 = torch.linalg.norm(dir_1, dim=1)
        side_length_2 = torch.linalg.norm(dir_2, dim=1)
        # scale the number of point w.r.t. the 'form' of the parallelogram
        n_1 = int(torch.sqrt(scaled_n*side_length_1/side_length_2))
        n_2 = int(torch.sqrt(scaled_n*side_length_2/side_length_1))
        x = torch.linspace(0, 1, n_1+2, device=device)[1:-1] # create inner grid, so remove
        y = torch.linspace(0, 1, n_2+2, device=device)[1:-1] # first and last value
        bary_coords = torch.stack(torch.meshgrid((x, y))).T.reshape(-1, 2)
        index = torch.where(bary_coords.sum(axis=1) <= 1)
        return bary_coords[index]
//  [python]
def _get_normal_direction(self, direction, device):
        # to get normal vector in 2d switch x and y coordinate and multiply 
        # one coordinate with -1
        normal = torch.index_select(direction, 1,
                                    torch.tensor([1, 0], device=device))
        normal[:, 1:] *= -1
        return torch.divide(normal, torch.norm(normal, dim=1).reshape(-1, 1))
//  [python]
def _get_normal_direction(self, direction, device):
        # to get normal vector in 2d switch x and y coordinate and multiply 
        # one coordinate with -1
        normal = torch.index_select(direction, 1,
                                    torch.tensor([1, 0], device=device))
        normal[:, 1:] *= -1
        return torch.divide(normal, torch.linalg.norm(normal, dim=1).reshape(-1, 1))
//  [python]
def eikonal_loss(sdf_gradients):
    gradient_error = (torch.norm(sdf_gradients.reshape(-1, 3), p=2, dim=-1) - 1.0) ** 2
    return gradient_error.mean()
//  [python]
def eikonal_loss(sdf_gradients):
    gradient_error = (torch.linalg.norm(sdf_gradients.reshape(-1, 3), ord=2, dim=-1) - 1.0) ** 2
    return gradient_error.mean()
//  [python]
def rotation_matrix_between(a: Float[Tensor, "3"], b: Float[Tensor, "3"]) -> Float[Tensor, "3 3"]:
    """Compute the rotation matrix that rotates vector a to vector b.

    Args:
        a: The vector to rotate.
        b: The vector to rotate to.
    Returns:
        The rotation matrix.
    """
    a = a / t.norm(a)
    b = b / t.norm(b)
    v = t.cross(a, b)  # Axis of rotation.

    # Handle cases where `a` and `b` are parallel.
    eps = 1e-6
    if t.sum(t.abs(v)) < eps:
        x = t.tensor([1.0, 0, 0]) if abs(a[0]) < eps else t.tensor([0, 1.0, 0])
        v = t.cross(a, x)

    v = v / t.norm(v)
    skew_sym_mat = t.Tensor(
        [
            [0, -v[2], v[1]],
            [v[2], 0, -v[0]],
            [-v[1], v[0], 0],
        ]
    )
    theta = t.acos(t.clip(t.dot(a, b), -1, 1))

    # Rodrigues rotation formula. https://en.wikipedia.org/wiki/Rodrigues%27_rotation_formula
    return t.eye(3) + t.sin(theta) * skew_sym_mat + (1 - t.cos(theta)) * (skew_sym_mat @ skew_sym_mat)
//  [python]
def rotation_matrix_between(a: Float[Tensor, "3"], b: Float[Tensor, "3"]) -> Float[Tensor, "3 3"]:
    """Compute the rotation matrix that rotates vector a to vector b.

    Args:
        a: The vector to rotate.
        b: The vector to rotate to.
    Returns:
        The rotation matrix.
    """
    a = a / torch.linalg.norm(a)
    b = b / torch.linalg.norm(b)
    v = torch.linalg.cross(a, b)  # Axis of rotation.

    # Handle cases where `a` and `b` are parallel.
    eps = 1e-6
    if torch.sum(torch.abs(v)) < eps:
        x = torch.tensor([1.0, 0, 0]) if abs(a[0]) < eps else torch.tensor([0, 1.0, 0])
        v = torch.linalg.cross(a, x)

    v = v / torch.linalg.norm(v)
    skew_sym_mat = torch.Tensor(
        [
            [0, -v[2], v[1]],
            [v[2], 0, -v[0]],
            [-v[1], v[0], 0],
        ]
    )
    theta = torch.acos(torch.clip(torch.dot(a, b), -1, 1))

    # Rodrigues rotation formula. https://en.wikipedia.org/wiki/Rodrigues%27_rotation_formula
    return torch.eye(3) + torch.sin(theta) * skew_sym_mat + (1 - torch.cos(theta)) * (skew_sym_mat @ skew_sym_mat)
//  [python]
def round_tt(tt_cores,R,eps,Rmax,is_ttm=False):
    """
    Rounds a TT-tensor (tt_cores have to be orthogonal)

    Parameters
    ----------
    tt_cores : list of torch tensors.
        Orthogonal TT cores.
    R : list of integers of length d+1.
        ranks of the TT-decomposition.
    eps : double.
        desired rounding accuracy.
    Rmax : list of integers
        the maximum rank that is allowed.

    Returns
    -------
    tt_cores : list of torch tensors.
        The TT-cores of the rounded tensor.
    R : list of inteders of length d+1.
        rounded ranks.

    """
    d = len(tt_cores)
    if d == 1:
        tt_cores = [tt_cores[0].clone()]
        return tt_cores, R
    tt_cores, R = lr_orthogonal(tt_cores, R, is_ttm)
    core_now = tt_cores[-1]
    eps = eps / np.sqrt(d-1) 

    
    for i in range(d-1,0,-1):
        core_next = tt_cores[i-1]
        
        core_now = tn.reshape(core_now,[R[i],-1])
        core_next = tn.reshape(core_next,[-1,R[i]])
        
        
        U, S, V = SVD(core_now)
        if S.is_cuda:
            r_now = min([Rmax[i],rank_chop(S.cpu().numpy(),tn.linalg.norm(S).cpu().numpy()*eps)])
        else:
            r_now = min([Rmax[i],rank_chop(S.numpy(),tn.linalg.norm(S).numpy()*eps)])
    
        U = U[:,:r_now]
        S = S[:r_now]
        V = V[:r_now,:]
        
        U = U @ tn.diag(S)
        R[i] = r_now
        core_next = core_next @ U
        core_now = V
        
        tt_cores[i] = tn.reshape(core_now,[R[i]]+list(tt_cores[i].shape[1:-1])+[R[i+1]])
        tt_cores[i-1] = tn.reshape(core_next,[R[i-1]]+list(tt_cores[i-1].shape[1:-1])+[R[i]])
        
        core_now = core_next
    
    return tt_cores, R
//  [python]
def normal(self, points, params=Points.empty(), device='cpu'):
        points, params, device = \
            self._transform_input_for_normals(points, params, device)
        origin, _, _, dir_1, dir_2, dir_3 = \
            self.domain._construct_triangle(points.join(params), device)
        points = points[:, list(self.space.keys())].as_tensor
        normals = torch.zeros_like(points, device=device)
        bary_x, bary_y = self.domain._solve_lgs(points - origin, dir_1, -dir_3)
        normal_dir_1 = self._get_normal_direction(dir_1, device)
        normal_dir_2 = self._get_normal_direction(dir_2, device)
        normal_dir_3 = self._get_normal_direction(dir_3, device)
        # compute for each point what the normal vector should be, by checking the
        # value of the local barycentric coordinate = 0 or sum = 1
        self._add_local_normal_vector(normals, bary_x, normal_dir_3, 0.0)
        self._add_local_normal_vector(normals, (bary_x + bary_y), normal_dir_2, 1.0)
        self._add_local_normal_vector(normals, bary_y, normal_dir_1, 0.0)
        # scale normal vectors if there where in a corner:
        return torch.div(normals, torch.norm(normals, dim=1).reshape(-1, 1))
//  [python]
def normal(self, points, params=Points.empty(), device='cpu'):
        points, params, device = \
            self._transform_input_for_normals(points, params, device)
        origin, _, _, dir_1, dir_2, dir_3 = \
            self.domain._construct_triangle(points.join(params), device)
        points = points[:, list(self.space.keys())].as_tensor
        normals = torch.zeros_like(points, device=device)
        bary_x, bary_y = self.domain._solve_lgs(points - origin, dir_1, -dir_3)
        normal_dir_1 = self._get_normal_direction(dir_1, device)
        normal_dir_2 = self._get_normal_direction(dir_2, device)
        normal_dir_3 = self._get_normal_direction(dir_3, device)
        # compute for each point what the normal vector should be, by checking the
        # value of the local barycentric coordinate = 0 or sum = 1
        self._add_local_normal_vector(normals, bary_x, normal_dir_3, 0.0)
        self._add_local_normal_vector(normals, (bary_x + bary_y), normal_dir_2, 1.0)
        self._add_local_normal_vector(normals, bary_y, normal_dir_1, 0.0)
        # scale normal vectors if there where in a corner:
        return torch.divide(normals, torch.linalg.norm(normals, dim=1).reshape(-1, 1))
//  [python]
def _CorrCholesky_check(self, value):
    row_norm = torch.norm(value.detach(), dim=-1)
    unit_row_norm = (row_norm - 1.0).abs().le(1e-4).all(dim=-1)
    return torch.distributions.constraints.lower_cholesky.check(value) & unit_row_norm
//  [python]
def _CorrCholesky_check(self, value):
    row_norm = torch.linalg.norm(value.detach(), dim=-1)
    unit_row_norm = (row_norm - 1.0).abs().le(1e-4).all(dim=-1)
    return torch.distributions.constraints.lower_cholesky.check(value) & unit_row_norm
//  [python]
def _contains(self, points, params=Points.empty()):
        center, radius = self._compute_center_and_radius(points.join(params), points.device)
        points = points[:, list(self.space.keys())].as_tensor
        norm = torch.norm(points - center, dim=1).reshape(-1, 1)
        return torch.le(norm[:, None], radius).reshape(-1, 1)
//  [python]
def _contains(self, points, params=Points.empty()):
        center, radius = self._compute_center_and_radius(points.join(params), points.device)
        points = points[:, list(self.space.keys())].as_tensor
        norm = torch.linalg.norm(points - center, dim=1).reshape(-1, 1)
        return torch.le(norm[:, None], radius).reshape(-1, 1)