jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def solve(self, fn, x0, rtol=1e-6, atol=1e-10, miter=100, ltol=1e-10):
        """
        Solve the nonlinear system
        """
        x = x0
        R, J = fn(x)

        if not self.stored_preconditioner:
            self._update_preconditioner(J)

        nR = torch.norm(R, dim=-1)
        nR0 = nR
        i = 0

        l_miter = J.shape[-1] + 1
        l_check = max(1, (J.shape[-1] + 1) // 10)
        l_refactor = (J.shape[-1] * 2) // 3

        while (i < miter) and torch.any(nR > atol) and torch.any(nR / nR0 > rtol):
            dx, niter = gmres(
                J,
                R,
                x0=x,
                M=self.stored_preconditioner,
                maxiter=l_miter,
                check=l_check,
                return_iters=True,
                tol=ltol,
            )
            x -= dx
            R, J = fn(x)
            if niter > l_refactor:
                self._update_preconditioner(J)
            nR = torch.norm(R, dim=-1)
            i += 1

        if i == miter:
            warnings.warn(
                "Implicit solve did not succeed.  Results may be inaccurate..."
            )

        return x, J
//  [python]
def solve(self, fn, x0, rtol=1e-6, atol=1e-10, miter=100, ltol=1e-10):
        """
        Solve the nonlinear system
        """
        x = x0
        R, J = fn(x)

        if not self.stored_preconditioner:
            self._update_preconditioner(J)

        nR = torch.linalg.norm(R, dim=-1)
        nR0 = nR
        i = 0

        l_miter = J.shape[-1] + 1
        l_check = max(1, (J.shape[-1] + 1) // 10)
        l_refactor = (J.shape[-1] * 2) // 3

        while (i < miter) and torch.any(nR > atol) and torch.any(nR / nR0 > rtol):
            dx, niter = gmres(
                J,
                R,
                x0=x,
                M=self.stored_preconditioner,
                maxiter=l_miter,
                check=l_check,
                return_iters=True,
                tol=ltol,
            )
            x -= dx
            R, J = fn(x)
            if niter > l_refactor:
                self._update_preconditioner(J)
            nR = torch.linalg.norm(R, dim=-1)
            i += 1

        if i == miter:
            warnings.warn(
                "Implicit solve did not succeed.  Results may be inaccurate..."
            )

        return x, J
//  [python]
def step(self):
        for g in self.param_groups:
            for p in g['params']:
                dp = p.grad

                if dp is None:
                    continue

                if g['weight_decay_filter'] is None or not g['weight_decay_filter'](p):
                    dp = dp.add(p, alpha=g['weight_decay'])

                if g['lars_adaptation_filter'] is None or not g['lars_adaptation_filter'](p):
                    param_norm = torch.norm(p)
                    update_norm = torch.norm(dp)
                    one = torch.ones_like(param_norm)
                    q = torch.where(param_norm > 0.,
                                    torch.where(update_norm > 0,
                                                (g['eta'] * param_norm / update_norm), one), one)
                    dp = dp.mul(q)

                param_state = self.state[p]
                if 'mu' not in param_state:
                    param_state['mu'] = torch.zeros_like(p)
                mu = param_state['mu']
                mu.mul_(g['momentum']).add_(dp)

                p.add_(mu, alpha=-g['lr'])
//  [python]
def step(self):
        for g in self.param_groups:
            for p in g['params']:
                dp = p.grad

                if dp is None:
                    continue

                if g['weight_decay_filter'] is None or not g['weight_decay_filter'](p):
                    dp = dp.add(p, alpha=g['weight_decay'])

                if g['lars_adaptation_filter'] is None or not g['lars_adaptation_filter'](p):
                    param_norm = torch.linalg.norm(p)
                    update_norm = torch.linalg.norm(dp)
                    one = torch.ones_like(param_norm)
                    q = torch.where(param_norm > 0.,
                                    torch.where(update_norm > 0,
                                                (g['eta'] * param_norm / update_norm), one), one)
                    dp = dp.mul(q)

                param_state = self.state[p]
                if 'mu' not in param_state:
                    param_state['mu'] = torch.zeros_like(p)
                mu = param_state['mu']
                mu.mul_(g['momentum']).add_(dp)

                p.add_(mu, alpha=-g['lr'])
//  [python]
def newton_raphson(fn, x0, linsolver="lu", rtol=1e-6, atol=1e-10, miter=100):
    """
    Solve a nonlinear system with Newton's method.  Return the
    solution and the last Jacobian

    Args:
      fn (function):        function that returns the residual and Jacobian
      x0 (torch.tensor):    starting point

    Keyword Args:
      linsolver (string):   method to use to solve the linear system, options are
                            "diag" or "lu".  Defaults to "lu".  See
                            :py:func:`pyoptmat.solvers.solve_linear_system`
      rtol (float):         nonlinear relative tolerance
      atol (float):         nonlinear absolute tolerance
      miter (int):          maximum number of nonlinear iterations

    Returns:
      torch.tensor, torch.tensor:   solution to system of equations and
                                    Jacobian evaluated at that point
    """
    # Determine the solver function
    if x0.shape[-1] == 1:
        solver = jacobi_iteration_linear_solve
    elif linsolver == "lu":
        solver = lu_linear_solve
    elif linsolver == "diag":
        solver = jacobi_iteration_linear_solve
    else:
        raise ValueError(f"Unknown linear solver type {linsolver}")

    x = x0
    R, J = fn(x)

    nR = torch.norm(R, dim=-1)
    nR0 = nR
    i = 0

    while (i < miter) and torch.any(nR > atol) and torch.any(nR / nR0 > rtol):
        x -= solver(J, R)
        R, J = fn(x)
        nR = torch.norm(R, dim=-1)
        i += 1

    if i == miter:
        warnings.warn("Implicit solve did not succeed.  Results may be inaccurate...")

    return x, J
//  [python]
def newton_raphson(fn, x0, linsolver="lu", rtol=1e-6, atol=1e-10, miter=100):
    """
    Solve a nonlinear system with Newton's method.  Return the
    solution and the last Jacobian

    Args:
      fn (function):        function that returns the residual and Jacobian
      x0 (torch.tensor):    starting point

    Keyword Args:
      linsolver (string):   method to use to solve the linear system, options are
                            "diag" or "lu".  Defaults to "lu".  See
                            :py:func:`pyoptmat.solvers.solve_linear_system`
      rtol (float):         nonlinear relative tolerance
      atol (float):         nonlinear absolute tolerance
      miter (int):          maximum number of nonlinear iterations

    Returns:
      torch.tensor, torch.tensor:   solution to system of equations and
                                    Jacobian evaluated at that point
    """
    # Determine the solver function
    if x0.shape[-1] == 1:
        solver = jacobi_iteration_linear_solve
    elif linsolver == "lu":
        solver = lu_linear_solve
    elif linsolver == "diag":
        solver = jacobi_iteration_linear_solve
    else:
        raise ValueError(f"Unknown linear solver type {linsolver}")

    x = x0
    R, J = fn(x)

    nR = torch.linalg.norm(R, dim=-1)
    nR0 = nR
    i = 0

    while (i < miter) and torch.any(nR > atol) and torch.any(nR / nR0 > rtol):
        x -= solver(J, R)
        R, J = fn(x)
        nR = torch.linalg.norm(R, dim=-1)
        i += 1

    if i == miter:
        warnings.warn("Implicit solve did not succeed.  Results may be inaccurate...")

    return x, J
//  [python]
def newton_raphson_bt(
    fn, x0, linsolver="lu", rtol=1e-6, atol=1e-10, miter=100, max_bt=5
):
    """
    Solve a nonlinear system with Newton's method.  Return the
    solution and the last Jacobian

    Args:
      fn (function):        function that returns the residual and Jacobian
      x0 (torch.tensor):    starting point

    Keyword Args:
      linsolver (string or function):   method to use to solve the linear system,
                                        strain options are "diag" or "lu".
                                        Defaults to "lu".
      rtol (float):         nonlinear relative tolerance
      atol (float):         nonlinear absolute tolerance
      miter (int):          maximum number of nonlinear iterations
      max_bt (int):         maximum number of backtracking line search iterations

    Returns:
      torch.tensor, torch.tensor:   solution to system of equations and
                                    Jacobian evaluated at that point
    """
    # Determine the solver function
    if x0.shape[-1] == 1:
        solver = jacobi_iteration_linear_solve
    elif linsolver == "lu":
        solver = lu_linear_solve
    elif linsolver == "diag":
        solver = jacobi_iteration_linear_solve
    else:
        raise ValueError(f"Unknown linear solver type {linsolver}")

    x = x0
    R, J = fn(x)

    nR = torch.norm(R, dim=-1)
    nR0 = nR
    i = 0

    alpha = torch.ones_like(x0)
    while (i < miter) and torch.any(nR > atol) and torch.any(nR / nR0 > rtol):
        dx = solver(J, R)
        nR_prev = nR.detach().clone()
        nR = nR_prev
        k = 0
        alpha = torch.ones_like(x0)
        while True:
            xp = x - alpha * dx
            R, J = fn(xp)
            nR = torch.norm(R, dim=-1)

            if torch.all(nR_prev > nR) or k > max_bt:
                break
            alpha[(nR - nR_prev) > atol] /= 2.0
            k += 1

        x = xp
        i += 1

    if i == miter:
        warnings.warn("Implicit solve did not succeed.  Results may be inaccurate...")

    return x, J
//  [python]
def newton_raphson_bt(
    fn, x0, linsolver="lu", rtol=1e-6, atol=1e-10, miter=100, max_bt=5
):
    """
    Solve a nonlinear system with Newton's method.  Return the
    solution and the last Jacobian

    Args:
      fn (function):        function that returns the residual and Jacobian
      x0 (torch.tensor):    starting point

    Keyword Args:
      linsolver (string or function):   method to use to solve the linear system,
                                        strain options are "diag" or "lu".
                                        Defaults to "lu".
      rtol (float):         nonlinear relative tolerance
      atol (float):         nonlinear absolute tolerance
      miter (int):          maximum number of nonlinear iterations
      max_bt (int):         maximum number of backtracking line search iterations

    Returns:
      torch.tensor, torch.tensor:   solution to system of equations and
                                    Jacobian evaluated at that point
    """
    # Determine the solver function
    if x0.shape[-1] == 1:
        solver = jacobi_iteration_linear_solve
    elif linsolver == "lu":
        solver = lu_linear_solve
    elif linsolver == "diag":
        solver = jacobi_iteration_linear_solve
    else:
        raise ValueError(f"Unknown linear solver type {linsolver}")

    x = x0
    R, J = fn(x)

    nR = torch.linalg.norm(R, dim=-1)
    nR0 = nR
    i = 0

    alpha = torch.ones_like(x0)
    while (i < miter) and torch.any(nR > atol) and torch.any(nR / nR0 > rtol):
        dx = solver(J, R)
        nR_prev = nR.detach().clone()
        nR = nR_prev
        k = 0
        alpha = torch.ones_like(x0)
        while True:
            xp = x - alpha * dx
            R, J = fn(xp)
            nR = torch.linalg.norm(R, dim=-1)

            if torch.all(nR_prev > nR) or k > max_bt:
                break
            alpha[(nR - nR_prev) > atol] /= 2.0
            k += 1

        x = xp
        i += 1

    if i == miter:
        warnings.warn("Implicit solve did not succeed.  Results may be inaccurate...")

    return x, J
//  [python]
def magnitude_error(prediction, target, p=2):
    assert prediction.shape[1] > 1
    error = torch.abs(
        torch.norm(prediction, p=p, dim=-1) - torch.norm(target, p=p, dim=-1)
    )
    return {
        "metric": torch.mean(error).item(),
        "total": torch.sum(error).item(),
        "numel": error.numel(),
    }
//  [python]
def magnitude_error(prediction, target, p=2):
    assert prediction.shape[1] > 1
    error = torch.abs(
        torch.linalg.norm(prediction, ord=p, dim=-1) - torch.linalg.norm(target, ord=p, dim=-1)
    )
    return {
        "metric": torch.mean(error).item(),
        "total": torch.sum(error).item(),
        "numel": error.numel(),
    }
//  [python]
def qlog(cls, q, ordering='wxyz'):
        """
        Applies the log map to quaternions.
        """
        n = 0.5*torch.norm(q[:, 1:], p=2, dim=1, keepdim=True)
        n = torch.clamp(n, min=1e-8)
        q = q[:, 1:] * torch.acos(torch.clamp(q[:, :1], min=-1.0, max=1.0))
        r = q / n
        return r
//  [python]
def qlog(cls, q, ordering='wxyz'):
        """
        Applies the log map to quaternions.
        """
        n = 0.5*torch.linalg.norm(q[:, 1:], ord=2, dim=1, keepdim=True)
        n = torch.clamp(n, min=1e-8)
        q = q[:, 1:] * torch.acos(torch.clamp(q[:, :1], min=-1.0, max=1.0))
        r = q / n
        return r
//  [python]
def compute_skate(vert, info):
    skate = []
    for t in range(vert.shape[0] - 1):
        cind = (vert[t, :, 2] <= info["floor_z"]) & (vert[t + 1, :, 2] <=
                                                     info["floor_z"])
        if torch.any(cind):
            offset = vert[t + 1, cind, :2] - vert[t, cind, :2]
            skate_i = torch.norm(offset, dim=1).mean().item() * 1000
        else:
            skate_i = 0.0
        skate.append(skate_i)
    return skate
//  [python]
def compute_skate(vert, info):
    skate = []
    for t in range(vert.shape[0] - 1):
        cind = (vert[t, :, 2] <= info["floor_z"]) & (vert[t + 1, :, 2] <=
                                                     info["floor_z"])
        if torch.any(cind):
            offset = vert[t + 1, cind, :2] - vert[t, cind, :2]
            skate_i = torch.linalg.norm(offset, dim=1).mean().item() * 1000
        else:
            skate_i = 0.0
        skate.append(skate_i)
    return skate
//  [python]
def brute_cdist(x, y, p=2):
    r1 = x.shape[-2]
    r2 = y.shape[-2]
    if r1 == 0 or r2 == 0:
        return torch.empty(r1, r2, device=x.device)
    return torch.norm(x[..., None, :] - y[..., None, :, :], p=p, dim=-1)
//  [python]
def brute_cdist(x, y, p=2):
    r1 = x.shape[-2]
    r2 = y.shape[-2]
    if r1 == 0 or r2 == 0:
        return torch.empty(r1, r2, device=x.device)
    return torch.linalg.norm(x[..., None, :] - y[..., None, :, :], p=p, dim=-1)
//  [python]
def brute_pdist(inp, p=2):
    """Computes the same as torch.pdist using primitives"""
    n = inp.shape[-2]
    k = n * (n - 1) // 2
    if k == 0:
        # torch complains about empty indices
        return torch.empty(inp.shape[:-2] + (0,), dtype=inp.dtype, device=inp.device)
    square = torch.norm(inp[..., None, :] - inp[..., None, :, :], p=p, dim=-1)
    unroll = square.view(square.shape[:-2] + (n * n,))
    inds = torch.ones(k, dtype=torch.int)
    inds[torch.arange(n - 1, 1, -1, dtype=torch.int).cumsum(0)] += torch.arange(
        2, n, dtype=torch.int
    )
    return unroll[..., inds.cumsum(0)]
//  [python]
def brute_pdist(inp, p=2):
    """Computes the same as torch.pdist using primitives"""
    n = inp.shape[-2]
    k = n * (n - 1) // 2
    if k == 0:
        # torch complains about empty indices
        return torch.empty(inp.shape[:-2] + (0,), dtype=inp.dtype, device=inp.device)
    square = torch.linalg.norm(inp[..., None, :] - inp[..., None, :, :], p=p, dim=-1)
    unroll = square.view(square.shape[:-2] + (n * n,))
    inds = torch.ones(k, dtype=torch.int)
    inds[torch.arange(n - 1, 1, -1, dtype=torch.int).cumsum(0)] += torch.arange(
        2, n, dtype=torch.int
    )
    return unroll[..., inds.cumsum(0)]
//  [python]
def forward(ctx, M, v_k, num_iter=19):
        '''
        :param ctx: used to save meterials for backward.
        :param M: n by n matrix.
        :param v_k: initial guess of leading vector.
        :return: v_k1 leading vector.
        '''
        vk_list = []
        vk_list.append(v_k)
        ctx.num_iter = num_iter
        for _ in range(int(ctx.num_iter)):
            v_k = M.mm(v_k)
            v_k /= th.norm(v_k).clamp(min=1.e-5)
            vk_list.append(v_k)

        ctx.save_for_backward(M, *vk_list)
        return v_k
//  [python]
def forward(ctx, M, v_k, num_iter=19):
        '''
        :param ctx: used to save meterials for backward.
        :param M: n by n matrix.
        :param v_k: initial guess of leading vector.
        :return: v_k1 leading vector.
        '''
        vk_list = []
        vk_list.append(v_k)
        ctx.num_iter = num_iter
        for _ in range(int(ctx.num_iter)):
            v_k = M.mm(v_k)
            v_k /= th.linalg.norm(v_k).clamp(min=1.e-5)
            vk_list.append(v_k)

        ctx.save_for_backward(M, *vk_list)
        return v_k
//  [python]
def calc_mid(M, v_k, v_k1, dL_dvk1):
    I = th.eye(M.shape[-1], out=th.empty_like(M))
    mid = (I - v_k1.mm(th.t(v_k1)))/th.norm(M.mm(v_k)).clamp(min=1.e-5)
    mid = mid.mm(dL_dvk1)
    return mid
//  [python]
def calc_mid(M, v_k, v_k1, dL_dvk1):
    I = th.eye(M.shape[-1], out=th.empty_like(M))
    mid = (I - v_k1.mm(th.t(v_k1)))/th.linalg.norm(M.mm(v_k)).clamp(min=1.e-5)
    mid = mid.mm(dL_dvk1)
    return mid
//  [python]
def backward(ctx, grad_output):
        M, v_k = ctx.saved_tensors
        dL_dvk = grad_output
        I = th.eye(M.shape[-1], out=th.empty_like(M))
        numerator = I - v_k.mm(th.t(v_k))
        denominator = th.norm(M.mm(v_k)).clamp(min=1.e-5)
        ak = numerator / denominator
        term1 = ak
        q = M / denominator
        for i in range(1, ctx.num_iter + 1):
            ak = q.mm(ak)
            term1 += ak
        dL_dM = th.mm(term1.mm(dL_dvk), v_k.t())
        return dL_dM, ak
//  [python]
def backward(ctx, grad_output):
        M, v_k = ctx.saved_tensors
        dL_dvk = grad_output
        I = th.eye(M.shape[-1], out=th.empty_like(M))
        numerator = I - v_k.mm(th.t(v_k))
        denominator = th.linalg.norm(M.mm(v_k)).clamp(min=1.e-5)
        ak = numerator / denominator
        term1 = ak
        q = M / denominator
        for i in range(1, ctx.num_iter + 1):
            ak = q.mm(ak)
            term1 += ak
        dL_dM = th.mm(term1.mm(dL_dvk), v_k.t())
        return dL_dM, ak
//  [python]
def get_kmeans_mu(X, K, min_delta=1e-6):
    """
    input:
    x:              torch.Tensor (n, d)
    output:
    center:         torch.Tensor (1, k, d)
    """
    min,max = X.min(), X.max()
    X = (X-min) / (max - min)
    center = X[np.random.randint(0, X.shape[0], size=K), ...]
    delta = np.inf
    while delta > min_delta:
        l2_dis = torch.norm((X.unsqueeze(1).repeat(1,K,1)-center),p=2,dim=2)
        l2_cls = torch.argmin(l2_dis, dim=1)
        center_old = center.clone()
        for c in range(K):
            center[c] = X[l2_cls==c].mean(dim=0)
        delta = torch.norm((center_old-center),dim=1).max()
        print(delta)
    return (center.unsqueeze(0)*(max-min)+min)
//  [python]
def get_kmeans_mu(X, K, min_delta=1e-6):
    """
    input:
    x:              torch.Tensor (n, d)
    output:
    center:         torch.Tensor (1, k, d)
    """
    min,max = X.min(), X.max()
    X = (X-min) / (max - min)
    center = X[np.random.randint(0, X.shape[0], size=K), ...]
    delta = np.inf
    while delta > min_delta:
        l2_dis = torch.linalg.norm((X.unsqueeze(1).repeat(1,K,1)-center),ord=2,dim=2)
        l2_cls = torch.argmin(l2_dis, dim=1)
        center_old = center.clone()
        for c in range(K):
            center[c] = X[l2_cls==c].mean(dim=0)
        delta = torch.linalg.norm((center_old-center),ord=2,dim=1).max()
        print(delta)
    return (center.unsqueeze(0)*(max-min)+min)
//  [python]
def do_pca(self):
        sems = self.engine._objects_dc.clone().squeeze()
        N, C = sems.shape
        torch.manual_seed(0)
        randint = torch.randint(0, N, [200_000])
        sems /= (torch.norm(sems, dim=1, keepdim=True) + 1e-6)
        sem_chosen = sems[randint, :]
        self.proj_mat = self.pca(sem_chosen, n_components=3)
        print("project mat initialized !")
//  [python]
def do_pca(self):
        sems = self.engine._objects_dc.clone().squeeze()
        N, C = sems.shape
        torch.manual_seed(0)
        randint = torch.randint(0, N, [200_000])
        sems /= (torch.linalg.norm(sems, dim=1, keepdim=True) + 1e-6)
        sem_chosen = sems[randint, :]
        self.proj_mat = self.pca(sem_chosen, n_components=3)
        print("project mat initialized !")
//  [python]
def dohf(self, shallow_hiera, deep_hiera):
        """
        from shallow to deep: order, family, genus, class
        shallow_hiera: N, M*C
        deep_hiera: N, M*C

        return
        """
        if shallow_hiera==None:
            return deep_hiera, deep_hiera
        N1, MC1 = shallow_hiera.shape
        M1 = MC1//self.base_channels
        shallow_hiera_mean = shallow_hiera.reshape(N1, M1, self.base_channels) # N,M1*C -> N,M1,C
        shallow_hiera_mean = shallow_hiera_mean.mean(dim=1) # N, C

        N2, MC2 = deep_hiera.shape
        M2 = MC2//self.base_channels
        deep_hiera_dohf = deep_hiera.reshape(N2, M2, self.base_channels)  # N,M2*C -> N,M2,C
        deep_hiera_dohf = deep_hiera_dohf.permute(0, 2, 1).contiguous() # N,M2,C -> N,C,M2

        shallow_hiera_norm = torch.norm(shallow_hiera_mean, p=2, dim=1) # N
        projection = torch.bmm(shallow_hiera_mean.unsqueeze(1), deep_hiera_dohf)  # N, 1, M2
        projection = torch.bmm(shallow_hiera_mean.unsqueeze(2), projection)  # N, C, M2
        projection = projection / (shallow_hiera_norm * shallow_hiera_norm).view(-1, 1, 1)  # N, C, M2

        orthogonal_comp = deep_hiera_dohf - projection
        deep_hiera_dohf = deep_hiera_dohf + self.add_lambda * orthogonal_comp # N, C, M2
        deep_hiera_dohf = deep_hiera_dohf.permute(0, 2, 1).contiguous() # N, C, M2 -> N,M2,C
        deep_hiera_dohf = deep_hiera_dohf.reshape(N2, -1) # N,M2,C -> N, MC2
        # l2 normalization along dimension M2 and C
        deep_hiera_dohf = F.normalize(deep_hiera_dohf, dim=-1)
        return deep_hiera, deep_hiera_dohf
//  [python]
def dohf(self, shallow_hiera, deep_hiera):
        """
        from shallow to deep: order, family, genus, class
        shallow_hiera: N, M*C
        deep_hiera: N, M*C

        return
        """
        if shallow_hiera==None:
            return deep_hiera, deep_hiera
        N1, MC1 = shallow_hiera.shape
        M1 = MC1//self.base_channels
        shallow_hiera_mean = shallow_hiera.reshape(N1, M1, self.base_channels) # N,M1*C -> N,M1,C
        shallow_hiera_mean = shallow_hiera_mean.mean(dim=1) # N, C

        N2, MC2 = deep_hiera.shape
        M2 = MC2//self.base_channels
        deep_hiera_dohf = deep_hiera.reshape(N2, M2, self.base_channels)  # N,M2*C -> N,M2,C
        deep_hiera_dohf = deep_hiera_dohf.permute(0, 2, 1).contiguous() # N,M2,C -> N,C,M2

        shallow_hiera_norm = torch.linalg.norm(shallow_hiera_mean, ord=2, dim=1) # N
        projection = torch.bmm(shallow_hiera_mean.unsqueeze(1), deep_hiera_dohf)  # N, 1, M2
        projection = torch.bmm(shallow_hiera_mean.unsqueeze(2), projection)  # N, C, M2
        projection = projection / (shallow_hiera_norm * shallow_hiera_norm).view(-1, 1, 1)  # N, C, M2

        orthogonal_comp = deep_hiera_dohf - projection
        deep_hiera_dohf = deep_hiera_dohf + self.add_lambda * orthogonal_comp # N, C, M2
        deep_hiera_dohf = deep_hiera_dohf.permute(0, 2, 1).contiguous() # N, C, M2 -> N,M2,C
        deep_hiera_dohf = deep_hiera_dohf.reshape(N2, -1) # N,M2,C -> N, MC2
        # l2 normalization along dimension M2 and C
        deep_hiera_dohf = F.normalize(deep_hiera_dohf, dim=-1)
        return deep_hiera, deep_hiera_dohf
//  [python]
def test_franka_attached_object_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)
    sph_idx = robot_model.kinematics_config.get_sphere_index_from_link_name("attached_object")
    state = robot_model.get_state(q_test.clone())

    attached_spheres = state.link_spheres_tensor[:, sph_idx, :]
    # check if all values are equal to position and quaternion

    assert torch.norm(attached_spheres[:, :, 0] - 0.061) < 0.01
    assert torch.norm(attached_spheres[:, :, 1]) < 0.0001

    # attach an object:
    new_object = torch.zeros((2, 4), **(tensor_args.as_torch_dict()))
    new_object[:, 3] = 0.01
    new_object[:, 1] = 0.1
    robot_model.kinematics_config.update_link_spheres("attached_object", new_object)
    state = robot_model.get_state(q_test.clone())
    attached_spheres = state.link_spheres_tensor[:, sph_idx, :]
    robot_model.kinematics_config.detach_object()
//  [python]
def test_franka_attached_object_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)
    sph_idx = robot_model.kinematics_config.get_sphere_index_from_link_name("attached_object")
    state = robot_model.get_state(q_test.clone())

    attached_spheres = state.link_spheres_tensor[:, sph_idx, :]
    # check if all values are equal to position and quaternion

    assert torch.linalg.norm(attached_spheres[:, :, 0] - 0.061) < 0.01
    assert torch.linalg.norm(attached_spheres[:, :, 1]) < 0.0001

    # attach an object:
    new_object = torch.zeros((2, 4), **(tensor_args.as_torch_dict()))
    new_object[:, 3] = 0.01
    new_object[:, 1] = 0.1
    robot_model.kinematics_config.update_link_spheres("attached_object", new_object)
    state = robot_model.get_state(q_test.clone())
    attached_spheres = state.link_spheres_tensor[:, sph_idx, :]
    robot_model.kinematics_config.detach_object()
//  [python]
def batch_rodrigues(theta):
    #theta N x 3
    batch_size = theta.shape[0]
    l1norm = torch.norm(theta + 1e-8, p = 2, dim = 1)
    angle = torch.unsqueeze(l1norm, -1)
    normalized = torch.div(theta, angle)
    angle = angle * 0.5
    v_cos = torch.cos(angle)
    v_sin = torch.sin(angle)
    quat = torch.cat([v_cos, v_sin * normalized], dim = 1)
    
    return quat2mat(quat)
//  [python]
def batch_rodrigues(theta):
    #theta N x 3
    batch_size = theta.shape[0]
    l1norm = torch.linalg.norm(theta + 1e-8, ord=2, dim = 1)
    angle = torch.unsqueeze(l1norm, -1)
    normalized = torch.div(theta, angle)
    angle = angle * 0.5
    v_cos = torch.cos(angle)
    v_sin = torch.sin(angle)
    quat = torch.cat([v_cos, v_sin * normalized], dim = 1)
    
    return quat2mat(quat)
//  [python]
def get_kmeans_mu(X, n_centers, min_delta=1e-3, init_times=20):
        """
        input:
        x:              torch.Tensor (n, d),(n, 1, d)
        output:
        center:         torch.Tensor (1, k, d)
        """
        X = X.squeeze(1)
        min, max = X.min(), X.max()
        X = (X-min) / (max - 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).sum()
            if cost < min_cost:
                min_cost = cost
                center = tmp_center

        # center = X[np.random.choice(np.arange(X.shape[0]), size=4, replace=False), ...]
        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)*(max-min)+min)
//  [python]
def get_kmeans_mu(X, n_centers, min_delta=1e-3, init_times=20):
        """
        input:
        x:              torch.Tensor (n, d),(n, 1, d)
        output:
        center:         torch.Tensor (1, k, d)
        """
        X = X.squeeze(1)
        min, max = X.min(), X.max()
        X = (X-min) / (max - 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.linalg.norm((X.unsqueeze(1).repeat(1, n_centers, 1) - tmp_center), ord=2, dim=2)
            l2_cls = torch.argmin(l2_dis, dim=1)
            cost = 0
            for c in range(n_centers):
                cost += torch.linalg.norm(X[l2_cls==c] - tmp_center[c],ord=2,dim=1).sum()
            if cost < min_cost:
                min_cost = cost
                center = tmp_center

        # center = X[np.random.choice(np.arange(X.shape[0]), size=4, replace=False), ...]
        delta = np.inf
        while delta > min_delta:
            l2_dis = torch.linalg.norm((X.unsqueeze(1).repeat(1,n_centers,1)-center),ord=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.linalg.norm((center_old-center),ord=2,dim=1).max()

        return (center.unsqueeze(0)*(max-min)+min)
//  [python]
def forward(self, x):
        """
        Given x, returns z and the log-determinant log|df/dx|.

        Returns
        -------
        """
        if self.h in (F.elu, F.leaky_relu):
            u = self.u
        elif self.h == torch.tanh:
            scal = torch.log(1+torch.exp(self.w @ self.u)) - self.w @ self.u - 1
            u = self.u + scal * self.w / torch.norm(self.w) ** 2
        else:
            raise NotImplementedError("Non-linearity is not supported.")
        lin = torch.unsqueeze(x @ self.w, 1) + self.b
        z = x + u * self.h(lin)
        phi = functional_derivatives[self.h](lin) * self.w
        log_det = torch.log(torch.abs(1 + phi @ u) + 1e-4)
        return z, log_det
//  [python]
def forward(self, x):
        """
        Given x, returns z and the log-determinant log|df/dx|.

        Returns
        -------
        """
        if self.h in (F.elu, F.leaky_relu):
            u = self.u
        elif self.h == torch.tanh:
            scal = torch.log(1+torch.exp(self.w @ self.u)) - self.w @ self.u - 1
            u = self.u + scal * self.w / torch.linalg.norm(self.w) ** 2
        else:
            raise NotImplementedError("Non-linearity is not supported.")
        lin = torch.unsqueeze(x @ self.w, 1) + self.b
        z = x + u * self.h(lin)
        phi = functional_derivatives[self.h](lin) * self.w
        log_det = torch.log(torch.abs(1 + phi @ u) + 1e-4)
        return z, log_det
//  [python]
def test_franka_attach_object_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)
    sph_idx = robot_model.kinematics_config.get_sphere_index_from_link_name("attached_object")
    state = robot_model.get_state(q_test.clone())

    attached_spheres = state.link_spheres_tensor[:, sph_idx, :]
    # check if all values are equal to position and quaternion

    assert torch.norm(attached_spheres[:, :, 0] - 0.0829) < 0.1
    assert torch.norm(attached_spheres[:, :, 1]) < 0.0001

    # attach an object:
    new_object = torch.zeros((4, 4), **(tensor_args.as_torch_dict()))
    new_object[:2, 3] = 0.01
    new_object[:2, 1] = 0.1
    # print(attached_spheres[:, sph_idx,:].shape)

    robot_model.kinematics_config.attach_object(sphere_tensor=new_object)
    state = robot_model.get_state(q_test.clone())
    attached_spheres = state.link_spheres_tensor[:, sph_idx, :]
    assert (torch.norm(attached_spheres[:, :2, 0] - 0.13)) < 0.01
    assert torch.norm(attached_spheres[:, 2:, 0] - 0.061) < 0.01

    robot_model.kinematics_config.detach_object()
    state = robot_model.get_state(q_test.clone())
    attached_spheres = state.link_spheres_tensor[:, sph_idx, :]

    assert torch.norm(attached_spheres[:, :, 0] - 0.061) < 0.01
    assert torch.norm(attached_spheres[:, :, 1]) < 0.0001
    robot_model.kinematics_config.detach_object()
//  [python]
def test_franka_attach_object_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)
    sph_idx = robot_model.kinematics_config.get_sphere_index_from_link_name("attached_object")
    state = robot_model.get_state(q_test.clone())

    attached_spheres = state.link_spheres_tensor[:, sph_idx, :]
    # check if all values are equal to position and quaternion

    assert torch.linalg.norm(attached_spheres[:, :, 0] - 0.0829) < 0.1
    assert torch.linalg.norm(attached_spheres[:, :, 1]) < 0.0001

    # attach an object:
    new_object = torch.zeros((4, 4), **(tensor_args.as_torch_dict()))
    new_object[:2, 3] = 0.01
    new_object[:2, 1] = 0.1
    # print(attached_spheres[:, sph_idx,:].shape)

    robot_model.kinematics_config.attach_object(sphere_tensor=new_object)
    state = robot_model.get_state(q_test.clone())
    attached_spheres = state.link_spheres_tensor[:, sph_idx, :]
    assert (torch.linalg.norm(attached_spheres[:, :2, 0] - 0.13)) < 0.01
    assert torch.linalg.norm(attached_spheres[:, 2:, 0] - 0.061) < 0.01

    robot_model.kinematics_config.detach_object()
    state = robot_model.get_state(q_test.clone())
    attached_spheres = state.link_spheres_tensor[:, sph_idx, :]

    assert torch.linalg.norm(attached_spheres[:, :, 0] - 0.061) < 0.01
    assert torch.linalg.norm(attached_spheres[:, :, 1]) < 0.0001
    robot_model.kinematics_config.detach_object()
//  [python]
def forward(self, x):
        """
        Given x, returns z and the log-determinant log|df/dx|.
        """
        m, n = x.shape
        r = torch.norm(x - self.x0)
        h = 1 / (torch.exp(self.log_alpha) + r)
        beta = -torch.exp(self.log_alpha) + torch.log(1 + torch.exp(self.beta))
        z = x + beta * h * (x - self.x0)
        log_det = (n - 1) * torch.log(1 + beta * h) + \
                  torch.log(1 + beta * h - \
                            beta * r / (torch.exp(self.log_alpha) + r) ** 2)
        return z, log_det
//  [python]
def forward(self, x):
        """
        Given x, returns z and the log-determinant log|df/dx|.
        """
        m, n = x.shape
        r = torch.linalg.norm(x - self.x0)
        h = 1 / (torch.exp(self.log_alpha) + r)
        beta = -torch.exp(self.log_alpha) + torch.log(1 + torch.exp(self.beta))
        z = x + beta * h * (x - self.x0)
        log_det = (n - 1) * torch.log(1 + beta * h) + \
                  torch.log(1 + beta * h - \
                            beta * r / (torch.exp(self.log_alpha) + r) ** 2)
        return z, log_det
//  [python]
def train_step(self, dataloader) -> dict:
        """One-step training on the input dataloader.

        Parameters
        ----------
        dataloader : DataLoader
            the training dataloader

        Returns
        -------
        dict
            the output logs, including `loss` and `val_accuracy`, etc.
        """
        loss_fn = self.loss
        model = self.model

        self.reset_metrics()
        model.train()

        binary_reg = self.cfg.get("binary_reg", 0.)

        for epoch, batch in enumerate(dataloader):
            self.callbacks.on_train_batch_begin(epoch)
            x, y, out_index = self.unravel_batch(batch)
            x = self.to_device(x)
            y = self.to_device(y)

            if not isinstance(x, tuple):
                x = x,
            out, z = model(*x)
            if out_index is not None:
                out = out[out_index]
            loss = loss_fn(out, y) + torch.norm(z *
                                                (torch.ones_like(z) - z), p=1) * binary_reg
            loss.backward()
            for metric in self.metrics:
                metric.update_state(y.cpu(), out.detach().cpu())
            self.callbacks.on_train_batch_end(epoch)

        metrics = [metric.result() for metric in self.metrics]
        results = [loss.cpu().item()] + metrics

        return dict(zip(self.metrics_names, results))
//  [python]
def train_step(self, dataloader) -> dict:
        """One-step training on the input dataloader.

        Parameters
        ----------
        dataloader : DataLoader
            the training dataloader

        Returns
        -------
        dict
            the output logs, including `loss` and `val_accuracy`, etc.
        """
        loss_fn = self.loss
        model = self.model

        self.reset_metrics()
        model.train()

        binary_reg = self.cfg.get("binary_reg", 0.)

        for epoch, batch in enumerate(dataloader):
            self.callbacks.on_train_batch_begin(epoch)
            x, y, out_index = self.unravel_batch(batch)
            x = self.to_device(x)
            y = self.to_device(y)

            if not isinstance(x, tuple):
                x = x,
            out, z = model(*x)
            if out_index is not None:
                out = out[out_index]
            loss = loss_fn(out, y) + torch.linalg.norm(z *
                                                (torch.ones_like(z) - z), ord=1) * binary_reg
            loss.backward()
            for metric in self.metrics:
                metric.update_state(y.cpu(), out.detach().cpu())
            self.callbacks.on_train_batch_end(epoch)

        metrics = [metric.result() for metric in self.metrics]
        results = [loss.cpu().item()] + metrics

        return dict(zip(self.metrics_names, results))
//  [python]
def avg_quaternions_approx(quats: torch.Tensor, weights=None) -> torch.Tensor:
    """
    Args:
        quats: (N, 4)
    Returns:
        qAvg: (4,)
    """
    if weights is not None and len(quats) != len(weights):
        raise ValueError("Args are of different length")
    if weights is None:
        weights = torch.ones_like(quats[:, 0])
    qAvg = torch.zeros_like(quats[0])
    for i, q in enumerate(quats):
        # Correct for double cover, by ensuring that dot product
        # of quats[i] and quats[0] is positive
        if i > 0 and torch.dot(quats[i], quats[0]) < 0.0:
            weights[i] = -weights[i]
        qAvg += weights[i] * q
    return qAvg / torch.norm(qAvg)
//  [python]
def avg_quaternions_approx(quats: torch.Tensor, weights=None) -> torch.Tensor:
    """
    Args:
        quats: (N, 4)
    Returns:
        qAvg: (4,)
    """
    if weights is not None and len(quats) != len(weights):
        raise ValueError("Args are of different length")
    if weights is None:
        weights = torch.ones_like(quats[:, 0])
    qAvg = torch.zeros_like(quats[0])
    for i, q in enumerate(quats):
        # Correct for double cover, by ensuring that dot product
        # of quats[i] and quats[0] is positive
        if i > 0 and torch.dot(quats[i], quats[0]) < 0.0:
            weights[i] = -weights[i]
        qAvg += weights[i] * q
    return qAvg / torch.linalg.norm(qAvg)
//  [python]
def measure_avg_cos(W):
    assert(len(W.shape) == 2)
    ## W: embedding matrix, shape=[n_data, dim]
    n_data = W.shape[0]
    W_norm = torch.norm(W, dim=1)
    W /= W_norm.unsqueeze(1)
    WW = torch.matmul(W, W.T) ## WW.shape[n_data, n_data]
    triu_numel = n_data * (n_data-1) / 2
    cos = torch.sum(torch.triu(WW, diagonal=1)) / triu_numel
    return float(cos)
//  [python]
def measure_avg_cos(W):
    assert(len(W.shape) == 2)
    ## W: embedding matrix, shape=[n_data, dim]
    n_data = W.shape[0]
    W_norm = torch.linalg.norm(W, dim=1)
    W /= W_norm.unsqueeze(1)
    WW = torch.matmul(W, W.T) ## WW.shape[n_data, n_data]
    triu_numel = n_data * (n_data-1) / 2
    cos = torch.sum(torch.triu(WW, diagonal=1)) / triu_numel
    return float(cos)
//  [python]
def make_basis(w: torch.Tensor) -> torch.Tensor:
    """ Make basis from a column vector
    Args:
        w: (3), this will be the last column, i.e. Z axis
    Returns:
        basis: (3, 3) as [u, v, w]
    """
    w = w / torch.norm(w, dim=-1, keepdim=True)
    u = find_perpendicular(w)
    u = u / torch.norm(u, dim=-1, keepdim=True)
    v = torch.cross(w, u)
    basis = torch.zeros((3, 3), device=w.device, dtype=w.dtype)
    basis[:, 0] = u
    basis[:, 1] = v
    basis[:, 2] = w
    return basis
//  [python]
def make_basis(w: torch.Tensor) -> torch.Tensor:
    """ Make basis from a column vector
    Args:
        w: (3), this will be the last column, i.e. Z axis
    Returns:
        basis: (3, 3) as [u, v, w]
    """
    w = w / torch.linalg.norm(w, dim=-1, keepdim=True)
    u = find_perpendicular(w)
    u = u / torch.linalg.norm(u, dim=-1, keepdim=True)
    v = torch.cross(w, u)
    basis = torch.zeros((3, 3), device=w.device, dtype=w.dtype)
    basis[:, 0] = u
    basis[:, 1] = v
    basis[:, 2] = w
    return basis
//  [python]
def get_parameter_image_and_name_to_visualize(self, ISource=None, use_softmax=False, output_preweight=True):
        """
        visualize the regularizer parameters

        :param ISource: not used
        :param use_softmax: true: apply softmax to get pre-weight
        :param output_preweight: true: output the pre-weight of the regularizer , false: output the weight of the regualrizer
        :return:
        """

        # if self.clamp_local_weight:
        #     self.local_weights.data =self.local_weights.clamp(min=-self.local_pre_weight_max,max=self.local_pre_weight_max)
        if use_softmax:
            local_adapt_weights_pre = stable_softmax(self.local_weights, dim=1)  # torch.abs(self.local_weights)
        else:
            local_adapt_weights_pre = torch.abs(self.local_weights)
        if self.smoother.weighting_type == 'w_K_w':
            local_adapt_weights_pre = local_adapt_weights_pre / torch.norm(local_adapt_weights_pre, p=None, dim=1,
                                                                           keepdim=True)
        if output_preweight:
            local_adapt_weights = local_adapt_weights_pre
            name = 'preweight h_0'

        else:
            local_adapt_weights = self.embedded_smoother.smooth(local_adapt_weights_pre)
            name = 'weight w_0'

        dim = len(local_adapt_weights.shape) - 2
        adaptive_smoother_map = local_adapt_weights.detach()
        if self.smoother.weighting_type == 'w_K_w':
            adaptive_smoother_map = adaptive_smoother_map ** 2
        gaussian_stds = self.smoother.multi_gaussian_stds.detach()
        view_sz = [1] + [len(gaussian_stds)] + [1] * dim
        gaussian_stds = gaussian_stds.view(*view_sz)
        smoother_map = adaptive_smoother_map * (gaussian_stds ** 2)
        par_image = torch.sqrt(torch.sum(smoother_map, 1, keepdim=True))
        return par_image[:, 0], name
//  [python]
def get_parameter_image_and_name_to_visualize(self, ISource=None, use_softmax=False, output_preweight=True):
        """
        visualize the regularizer parameters

        :param ISource: not used
        :param use_softmax: true: apply softmax to get pre-weight
        :param output_preweight: true: output the pre-weight of the regularizer , false: output the weight of the regualrizer
        :return:
        """

        # if self.clamp_local_weight:
        #     self.local_weights.data =self.local_weights.clamp(min=-self.local_pre_weight_max,max=self.local_pre_weight_max)
        if use_softmax:
            local_adapt_weights_pre = stable_softmax(self.local_weights, dim=1)  # torch.abs(self.local_weights)
        else:
            local_adapt_weights_pre = torch.abs(self.local_weights)
        if self.smoother.weighting_type == 'w_K_w':
            local_adapt_weights_pre = local_adapt_weights_pre / torch.linalg.norm(local_adapt_weights_pre, ord=None, dim=1,
                                                                           keepdim=True)
        if output_preweight:
            local_adapt_weights = local_adapt_weights_pre
            name = 'preweight h_0'

        else:
            local_adapt_weights = self.embedded_smoother.smooth(local_adapt_weights_pre)
            name = 'weight w_0'

        dim = len(local_adapt_weights.shape) - 2
        adaptive_smoother_map = local_adapt_weights.detach()
        if self.smoother.weighting_type == 'w_K_w':
            adaptive_smoother_map = adaptive_smoother_map ** 2
        gaussian_stds = self.smoother.multi_gaussian_stds.detach()
        view_sz = [1] + [len(gaussian_stds)] + [1] * dim
        gaussian_stds = gaussian_stds.view(*view_sz)
        smoother_map = adaptive_smoother_map * (gaussian_stds ** 2)
        par_image = torch.sqrt(torch.sum(smoother_map, 1, keepdim=True))
        return par_image[:, 0], name
//  [python]
def PowerIter(mat_g, error_tolerance=1e-6, num_iters=100):
  """Power iteration.

  Compute the maximum eigenvalue of mat, for scaling.
  v is a random vector with values in (-1, 1)

  Args:
    mat_g: the symmetric PSD matrix.
    error_tolerance: Iterative exit condition.
    num_iters: Number of iterations.

  Returns:
    eigen vector, eigen value, num_iters
  """
  v = torch.rand(list(mat_g.shape)[0], device=mat_g.get_device()) * 2 - 1
  error = 1
  iters = 0
  singular_val = 0
  while error > error_tolerance and iters < num_iters:
    v = v / torch.norm(v)
    mat_v = torch.mv(mat_g, v)
    s_v = torch.dot(v, mat_v)
    error = torch.abs(s_v - singular_val)
    v = mat_v
    singular_val = s_v
    iters += 1
  return singular_val, v / torch.norm(v), iters
//  [python]
def PowerIter(mat_g, error_tolerance=1e-6, num_iters=100):
  """Power iteration.

  Compute the maximum eigenvalue of mat, for scaling.
  v is a random vector with values in (-1, 1)

  Args:
    mat_g: the symmetric PSD matrix.
    error_tolerance: Iterative exit condition.
    num_iters: Number of iterations.

  Returns:
    eigen vector, eigen value, num_iters
  """
  v = torch.rand(list(mat_g.shape)[0], device=mat_g.get_device()) * 2 - 1
  error = 1
  iters = 0
  singular_val = 0
  while error > error_tolerance and iters < num_iters:
    v = v / torch.linalg.norm(v)
    mat_v = torch.mv(mat_g, v)
    s_v = torch.dot(v, mat_v)
    error = torch.abs(s_v - singular_val)
    v = mat_v
    singular_val = s_v
    iters += 1
  return singular_val, v / torch.linalg.norm(v), iters