jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def newton_minimization(grad: math.Tensor, hessian: math.Tensor, max_loss_change=None) -> math.Tensor:
    """
    Computes a Newton update step that always points towards a minimum of the function.

    Args:
        grad: Gradient vector.
        hessian: Hessian matrix.
        max_loss_change: (Optional) If set, adds a limit to the gradient magnitude, approximately restricting by how much the loss can decrease per step.
            This affects directions in which the objective curvature is small but the gradient is large.
            The limit is computed as `max_loss_change / grad` in principal component space.
            If not set, large updates along these dimensions can occur which may result in divergence.

    Returns:
        Newton direction as `Tensor` matching `grad`.
    """
    assert hessian.shape.non_batch.rank == 2, "Not yet implemented"
    batch = hessian.shape.batch
    h_native = math.reshaped_native(hessian, [*batch, *hessian.shape.non_batch])
    # Transform to eigen space
    S, V = torch.symeig(h_native)  # identity = V.v * V_.v
    S = math.wrap(S, batch, math.channel('vector'))  # scales like 1/global_scale**2
    V = math.reshaped_tensor(V, [*batch, grad.shape.non_batch, math.channel('reduce')])
    V_ = math.reshaped_tensor(V, [*batch, math.channel('reduce'), grad.shape.non_batch])
    grad_eig = grad.vector * V.reduce
    newton_eig = - math.divide_no_nan(grad_eig, S) # scales like 1/global_scale
    S_avg = abs(math.prod(S, 'vector')) ** (1/S.vector.size)
    # S_rel = abs(S / S_avg)
    if max_loss_change is not None:
        inv_grad_eig = abs(math.divide_no_nan(max_loss_change, grad_eig))  # limit step by 1/grad when S is small
        # limit_fac = 1 / (1 + S_rel)  # 1 for small relative curvatures
        # limit = limit_fac * inv_grad_eig + (1 - limit_fac) * abs(newton_eig)
        limit = math.where(S > S_avg, abs(newton_eig), inv_grad_eig)
        newton_eig = math.minimum(abs(newton_eig), limit) * math.sign(-grad_eig)
    else:
        newton_eig = math.sign(-grad_eig) * abs(newton_eig)  # sign flip to descending direction
    # Transform back to parameter space
    newton_min = newton_eig.vector * V_.reduce
    return newton_min
//  [python]
def newton_minimization(grad: math.Tensor, hessian: math.Tensor, max_loss_change=None) -> math.Tensor:
    """
    Computes a Newton update step that always points towards a minimum of the function.

    Args:
        grad: Gradient vector.
        hessian: Hessian matrix.
        max_loss_change: (Optional) If set, adds a limit to the gradient magnitude, approximately restricting by how much the loss can decrease per step.
            This affects directions in which the objective curvature is small but the gradient is large.
            The limit is computed as `max_loss_change / grad` in principal component space.
            If not set, large updates along these dimensions can occur which may result in divergence.

    Returns:
        Newton direction as `Tensor` matching `grad`.
    """
    assert hessian.shape.non_batch.rank == 2, "Not yet implemented"
    batch = hessian.shape.batch
    h_native = math.reshaped_native(hessian, [*batch, *hessian.shape.non_batch])
    # Transform to eigen space
    S, V = torch.linalg.eigh(h_native)  # identity = V.v * V_.v
    S = math.wrap(S, batch, math.channel('vector'))  # scales like 1/global_scale**2
    V = math.reshaped_tensor(V, [*batch, grad.shape.non_batch, math.channel('reduce')])
    V_ = math.reshaped_tensor(V, [*batch, math.channel('reduce'), grad.shape.non_batch])
    grad_eig = grad.vector * V.reduce
    newton_eig = - math.divide_no_nan(grad_eig, S) # scales like 1/global_scale
    S_avg = abs(math.prod(S, 'vector')) ** (1/S.vector.size)
    # S_rel = abs(S / S_avg)
    if max_loss_change is not None:
        inv_grad_eig = abs(math.divide_no_nan(max_loss_change, grad_eig))  # limit step by 1/grad when S is small
        # limit_fac = 1 / (1 + S_rel)  # 1 for small relative curvatures
        # limit = limit_fac * inv_grad_eig + (1 - limit_fac) * abs(newton_eig)
        limit = math.where(S > S_avg, abs(newton_eig), inv_grad_eig)
        newton_eig = math.minimum(abs(newton_eig), limit) * math.sign(-grad_eig)
    else:
        newton_eig = math.sign(-grad_eig) * abs(newton_eig)  # sign flip to descending direction
    # Transform back to parameter space
    newton_min = newton_eig.vector * V_.reduce
    return newton_min
//  [python]
def A_vec_to_quat(A_vec):
    '''
    Input:
        - A_vec [B,10]
    Output:
        - quat [B,4]
    '''
    A = convert_Avec_to_A(A_vec)
    _, eigvector = torch.symeig(A, eigenvectors=True)
    if eigvector.dim() < 3:
        eigvector = eigvector.unsqueeze(0)
    return eigvector[:,:,0]
//  [python]
def A_vec_to_quat(A_vec):
    '''
    Input:
        - A_vec [B,10]
    Output:
        - quat [B,4]
    '''
    A = convert_Avec_to_A(A_vec)
    _, eigvector = torch.linalg.eigh(A, "U")
    if eigvector.dim() < 3:
        eigvector = eigvector.unsqueeze(0)
    return eigvector[:,:,0]
//  [python]
def damped_newton_minimization(grad: math.Tensor, hessian: math.Tensor, damping=0.1, max_loss_change=None) -> math.Tensor:
    """
    Computes a Newton update step that always points towards a minimum of the function.

    Args:
        grad: Gradient vector.
        hessian: Hessian matrix.
        max_loss_change: (Optional) If set, adds a limit to the gradient magnitude, approximately restricting by how much the loss can decrease per step.
            This affects directions in which the objective curvature is small but the gradient is large.
            The limit is computed as `max_loss_change / grad` in principal component space.
            If not set, large updates along these dimensions can occur which may result in divergence.

    Returns:
        Newton direction as `Tensor` matching `grad`.
    """
    assert hessian.shape.non_batch.rank == 2, "Not yet implemented"
    batch = hessian.shape.batch
    h_native = math.reshaped_native(hessian, [*batch, *hessian.shape.non_batch])
    grad_native = math.reshaped_native(grad, [*batch, grad.shape.non_batch])
    identity = torch.diag(torch.ones(grad.shape.non_batch.volume) * damping)
    newton_any = - torch.linalg.solve(h_native + identity, grad_native)
    newton_any = math.reshaped_tensor(newton_any, [*grad.shape.batch, grad.shape.non_batch])
    # Transform to eigen space
    S, V = torch.symeig(h_native)  # identity = V.v * V_.v
    S = math.wrap(S, batch, math.channel('basis'))
    V = math.reshaped_tensor(V, [*batch, grad.shape.non_batch, math.channel('reduce')])
    V_ = math.reshaped_tensor(V, [*batch, math.channel('reduce'), grad.shape.non_batch])
    grad_princ = grad.vector * V.reduce  # scales like global_scale
    newton_princ = newton_any.vector * V.reduce  # scales like 1/global_scale
    if max_loss_change is not None:
        inv_grad_princ = abs(math.divide_no_nan(max_loss_change, grad_princ))  # limit step by 1/grad
        newton_princ = math.minimum(abs(newton_princ), inv_grad_princ) * math.sign(-grad_princ)
    else:
        newton_princ *= math.sign(-grad_princ) * math.sign(newton_princ)  # sign flip to descending direction
    # Transform back to parameter space
    newton_min = newton_princ.vector * V_.reduce
    return newton_min
//  [python]
def damped_newton_minimization(grad: math.Tensor, hessian: math.Tensor, damping=0.1, max_loss_change=None) -> math.Tensor:
    """
    Computes a Newton update step that always points towards a minimum of the function.

    Args:
        grad: Gradient vector.
        hessian: Hessian matrix.
        max_loss_change: (Optional) If set, adds a limit to the gradient magnitude, approximately restricting by how much the loss can decrease per step.
            This affects directions in which the objective curvature is small but the gradient is large.
            The limit is computed as `max_loss_change / grad` in principal component space.
            If not set, large updates along these dimensions can occur which may result in divergence.

    Returns:
        Newton direction as `Tensor` matching `grad`.
    """
    assert hessian.shape.non_batch.rank == 2, "Not yet implemented"
    batch = hessian.shape.batch
    h_native = math.reshaped_native(hessian, [*batch, *hessian.shape.non_batch])
    grad_native = math.reshaped_native(grad, [*batch, grad.shape.non_batch])
    identity = torch.diag(torch.ones(grad.shape.non_batch.volume) * damping)
    newton_any = - torch.linalg.solve(h_native + identity, grad_native)
    newton_any = math.reshaped_tensor(newton_any, [*grad.shape.batch, grad.shape.non_batch])
    # Transform to eigen space
    S, V = torch.linalg.eigh(h_native)  # identity = V.v * V_.v
    S = math.wrap(S, batch, math.channel('basis'))
    V = math.reshaped_tensor(V, [*batch, grad.shape.non_batch, math.channel('reduce')])
    V_ = math.reshaped_tensor(V, [*batch, math.channel('reduce'), grad.shape.non_batch])
    grad_princ = grad.vector * V.reduce  # scales like global_scale
    newton_princ = newton_any.vector * V.reduce  # scales like 1/global_scale
    if max_loss_change is not None:
        inv_grad_princ = abs(math.divide_no_nan(max_loss_change, grad_princ))  # limit step by 1/grad
        newton_princ = math.minimum(abs(newton_princ), inv_grad_princ) * math.sign(-grad_princ)
    else:
        newton_princ *= math.sign(-grad_princ) * math.sign(newton_princ)  # sign flip to descending direction
    # Transform back to parameter space
    newton_min = newton_princ.vector * V_.reduce
    return newton_min
//  [python]
def __init__(self, scales, cond=None, rcond=None, upper=True):
        """Factorized representation of a batch of scale PSD matrices."""

        self._zero = scales.new_zeros(size=(1,))

        # scales: [B x D x D]

        s, u = torch.symeig(scales, eigenvectors=True, upper=True)
        #s, u = torch.linalg.eigh(scales)

        eps = _eigvalsh_to_eps(s, cond, rcond)

        if torch.any(torch.min(s, dim=-1).values < -eps):
            raise ValueError('scale is not positive semidefinite')

        # probably could use searchsorted
        self.s_mask = s > eps.unsqueeze(dim=-1)

        self.u = u
        self.s = torch.where(self.s_mask, s, self._zero)
        self.s_inv = torch.where(self.s_mask, 1 / s, self._zero)
//  [python]
def __init__(self, scales, cond=None, rcond=None, upper=True):
        """Factorized representation of a batch of scale PSD matrices."""

        self._zero = scales.new_zeros(size=(1,))

        # scales: [B x D x D]

        # s, u = torch.symeig(scales, eigenvectors=True, upper=True)
        s, u = torch.linalg.eigh(scales)

        eps = _eigvalsh_to_eps(s, cond, rcond)

        if torch.any(torch.min(s, dim=-1).values < -eps):
            raise ValueError('scale is not positive semidefinite')

        # probably could use searchsorted
        self.s_mask = s > eps.unsqueeze(dim=-1)

        self.u = u
        self.s = torch.where(self.s_mask, s, self._zero)
        self.s_inv = torch.where(self.s_mask, 1 / s, self._zero)
//  [python]
def rot_to_quat(rot: ROTATION_TENSOR_TYPE) -> QUATERNION_TENSOR_TYPE:
    if rot.shape[-2:] != (3, 3):
        raise ValueError("Input rotation is incorrectly shaped")

    rot = [[rot[..., i, j] for j in range(3)] for i in range(3)]
    [[xx, xy, xz], [yx, yy, yz], [zx, zy, zz]] = rot

    k = [
        [
            xx + yy + zz,
            zy - yz,
            xz - zx,
            yx - xy,
        ],
        [
            zy - yz,
            xx - yy - zz,
            xy + yx,
            xz + zx,
        ],
        [
            xz - zx,
            xy + yx,
            yy - xx - zz,
            yz + zy,
        ],
        [
            yx - xy,
            xz + zx,
            yz + zy,
            zz - xx - yy,
        ],
    ]

    k = (1.0 / 3.0) * torch.stack([torch.stack(t, dim=-1) for t in k], dim=-2)

    _, vectors = torch.symeig(k, eigenvectors=True)
    return vectors[..., -1]
//  [python]
def rot_to_quat(rot: ROTATION_TENSOR_TYPE) -> QUATERNION_TENSOR_TYPE:
    if rot.shape[-2:] != (3, 3):
        raise ValueError("Input rotation is incorrectly shaped")

    rot = [[rot[..., i, j] for j in range(3)] for i in range(3)]
    [[xx, xy, xz], [yx, yy, yz], [zx, zy, zz]] = rot

    k = [
        [
            xx + yy + zz,
            zy - yz,
            xz - zx,
            yx - xy,
        ],
        [
            zy - yz,
            xx - yy - zz,
            xy + yx,
            xz + zx,
        ],
        [
            xz - zx,
            xy + yx,
            yy - xx - zz,
            yz + zy,
        ],
        [
            yx - xy,
            xz + zx,
            yz + zy,
            zz - xx - yy,
        ],
    ]

    k = (1.0 / 3.0) * torch.stack([torch.stack(t, dim=-1) for t in k], dim=-2)

    _, vectors = torch.linalg.eigh(k)
    return vectors[..., -1]
//  [python]
def _get_dfrft_evecs(
    N: int,
    *,
    approx_order: int = 2,
    device: torch.device = torch.device("cpu"),
) -> torch.Tensor:
    if N < 1 or approx_order < 2:
        raise ValueError("Necessary conditions for integers: N > 1 and approx_order >= 2.")
    S = _create_hamiltonian(N, approx_order=approx_order, device=device)
    P = _create_odd_even_decomp_matrix(N, device=device)

    CS = torch.einsum("ij,jk,lk->il", P, S, P)
    C2 = CS[: N // 2 + 1, : N // 2 + 1]
    S2 = CS[N // 2 + 1 :, N // 2 + 1 :]

    _, VC = torch.symeig(C2, eigenvectors=True)  # ascending order
    _, VS = torch.symeig(S2, eigenvectors=True)  # ascending order

    N0, N1 = ceil(N / 2 - 1), N // 2 + 1
    qvc = torch.cat((VC, torch.zeros((N0, N1), device=device)))
    qvs = torch.cat((torch.zeros((N1, N0), device=device), VS))

    SC2 = torch.matmul(P, qvc).flip(-1)  # descending order
    SS2 = torch.matmul(P, qvs).flip(-1)  # descending order

    if N % 2 == 0:
        evecs = torch.zeros(N, N + 1, device=device)
        SS2_new = torch.hstack(
            (SS2, torch.zeros((SS2.size(0), 1), dtype=SS2.dtype, device=SS2.device))
        )
        evecs[:, : N + 1 : 2] = SC2
        evecs[:, 1:N:2] = SS2_new
        evecs = torch.hstack((evecs[:, : N - 1], evecs[:, -1].unsqueeze(-1)))
    else:
        evecs = torch.zeros(N, N, device=device)
        evecs[:, : N + 1 : 2] = SC2
        evecs[:, 1:N:2] = SS2
    return evecs
//  [python]
def _get_dfrft_evecs(
    N: int,
    *,
    approx_order: int = 2,
    device: torch.device = torch.device("cpu"),
) -> torch.Tensor:
    if N < 1 or approx_order < 2:
        raise ValueError("Necessary conditions for integers: N > 1 and approx_order >= 2.")
    S = _create_hamiltonian(N, approx_order=approx_order, device=device)
    P = _create_odd_even_decomp_matrix(N, device=device)

    CS = torch.einsum("ij,jk,lk->il", P, S, P)
    C2 = CS[: N // 2 + 1, : N // 2 + 1]
    S2 = CS[N // 2 + 1 :, N // 2 + 1 :]

    _, VC = torch.linalg.eigh(C2)  # ascending order
    _, VS = torch.linalg.eigh(S2)  # ascending order

    N0, N1 = ceil(N / 2 - 1), N // 2 + 1
    qvc = torch.cat((VC, torch.zeros((N0, N1), device=device)))
    qvs = torch.cat((torch.zeros((N1, N0), device=device), VS))

    SC2 = torch.matmul(P, qvc).flip(-1)  # descending order
    SS2 = torch.matmul(P, qvs).flip(-1)  # descending order

    if N % 2 == 0:
        evecs = torch.zeros(N, N + 1, device=device)
        SS2_new = torch.hstack(
            (SS2, torch.zeros((SS2.size(0), 1), dtype=SS2.dtype, device=SS2.device))
        )
        evecs[:, : N + 1 : 2] = SC2
        evecs[:, 1:N:2] = SS2_new
        evecs = torch.hstack((evecs[:, : N - 1], evecs[:, -1].unsqueeze(-1)))
    else:
        evecs = torch.zeros(N, N, device=device)
        evecs[:, : N + 1 : 2] = SC2
        evecs[:, 1:N:2] = SS2
    return evecs
//  [python]
def forward(self, mat_feats):

        e_x, v_x = torch.symeig(mat_feats, eigenvectors=True)
        
        e_x = self.apply_nonlin_to_ev(e_x)
        
        out = v_x @ torch.diag_embed(e_x) @ v_x.transpose(-1, -2)
        
        return out
//  [python]
def forward(self, mat_feats):

        e_x, v_x = torch.linalg.eigh(mat_feats, UPLO='U')
        
        e_x = self.apply_nonlin_to_ev(e_x)
        
        out = v_x @ torch.diag_embed(e_x) @ v_x.transpose(-1, -2)
        
        return out
//  [python]
def transfer_chrom(inp, src):
    '''
    inp (3, H, W):              input image
    src (3, H, W):              source image

    out (3, H, W):              output image
    '''
    
    m_i = torch.mean(inp, dim=(1,2))
    m_s = torch.mean(src, dim=(1,2))
    s_i = torch.maximum(torch.cov(inp.view(3, -1)), 7.5*torch.eye(3).to(inp.get_device())) 
    s_s = torch.maximum(torch.cov(src.view(3, -1)), 7.5*torch.eye(3).to(inp.get_device()))

    w_i, v_i = torch.symeig(s_i, eigenvectors=True)
    w_i[w_i<0] = 0
    da = torch.diag(torch.sqrt(w_i + 1e-10))
    c = torch.mm(da, torch.mm(v_i.T, torch.mm(s_s, torch.mm(v_i, da))))
    w_c, v_c = torch.symeig(c, eigenvectors=True)
    w_c[w_c<0] = 0
    dc = torch.diag(torch.sqrt(w_c + 1e-10))
    da_inv = torch.diag(1/(torch.diag(da)))
    T = torch.mm(v_i, torch.mm(da_inv, torch.mm(v_c, torch.mm(dc, torch.mm(v_c.T, torch.mm(da_inv, v_i.T))))))

    out = torch.matmul(T, (inp.view(3, -1) - m_i[:,None])) + m_s[:,None]

    return out.view(inp.size())
//  [python]
def transfer_chrom(inp, src):
    '''
    inp (3, H, W):              input image
    src (3, H, W):              source image

    out (3, H, W):              output image
    '''
    
    m_i = torch.mean(inp, dim=(1,2))
    m_s = torch.mean(src, dim=(1,2))
    s_i = torch.maximum(torch.cov(inp.view(3, -1)), 7.5*torch.eye(3).to(inp.get_device())) 
    s_s = torch.maximum(torch.cov(src.view(3, -1)), 7.5*torch.eye(3).to(inp.get_device()))

    w_i, v_i = torch.linalg.eigh(s_i)
    w_i[w_i<0] = 0
    da = torch.diag(torch.sqrt(w_i + 1e-10))
    c = torch.mm(da, torch.mm(v_i.T, torch.mm(s_s, torch.mm(v_i, da))))
    w_c, v_c = torch.linalg.eigh(c)
    w_c[w_c<0] = 0
    dc = torch.diag(torch.sqrt(w_c + 1e-10))
    da_inv = torch.diag(1/(torch.diag(da)))
    T = torch.mm(v_i, torch.mm(da_inv, torch.mm(v_c, torch.mm(dc, torch.mm(v_c.T, torch.mm(da_inv, v_i.T))))))

    out = torch.matmul(T, (inp.view(3, -1) - m_i[:,None])) + m_s[:,None]

    return out.view(inp.size())
//  [python]
def _symeig(
        self: Float[LinearOperator, "*batch N N"],
        eigenvectors: bool = False,
        return_evals_as_lazy: Optional[bool] = False,
    ) -> Tuple[Float[Tensor, "*batch M"], Optional[Float[LinearOperator, "*batch N M"]]]:
        r"""
        Method that allows implementing special-cased symeig computation. Should not be called directly
        """
        from linear_operator.operators.dense_linear_operator import DenseLinearOperator

        if settings.verbose_linalg.on():
            settings.verbose_linalg.logger.debug(f"Running symeig on a matrix of size {self.shape}.")

        # potentially perform decomposition in double precision for numerical stability
        dtype = self.dtype
        evals, evecs = torch.symeig(self.to_dense().to(dtype=settings._linalg_dtype_symeig.value()))
        # chop any negative eigenvalues.
        # TODO: warn if evals are significantly negative
        evals = evals.clamp_min(0.0).to(dtype=dtype)
        if eigenvectors:
            evecs = DenseLinearOperator(evecs.to(dtype=dtype))
        else:
            evecs = None
        return evals, evecs
//  [python]
def _symeig(
        self: Float[LinearOperator, "*batch N N"],
        eigenvectors: bool = False,
        return_evals_as_lazy: Optional[bool] = False,
    ) -> Tuple[Float[Tensor, "*batch M"], Optional[Float[LinearOperator, "*batch N M"]]]:
        r"""
        Method that allows implementing special-cased symeig computation. Should not be called directly
        """
        from linear_operator.operators.dense_linear_operator import DenseLinearOperator

        if settings.verbose_linalg.on():
            settings.verbose_linalg.logger.debug(f"Running symeig on a matrix of size {self.shape}.")

        # potentially perform decomposition in double precision for numerical stability
        dtype = self.dtype
        evals, evecs = torch.linalg.eigh(self.to_dense().to(dtype=settings._linalg_dtype_symeig.value()))
        # chop any negative eigenvalues.
        # TODO: warn if evals are significantly negative
        evals = evals.clamp_min(0.0).to(dtype=dtype)
        if eigenvectors:
            evecs = DenseLinearOperator(evecs.to(dtype=dtype))
        else:
            evecs = None
        return evals, evecs
//  [python]
def _check_point_on_manifold(self, x: torch.Tensor, *, atol=1e-05, rtol=1e-05) ->Union[Tuple[bool, Optional[str]], bool]:
        ok = torch.allclose(x, x.transpose(-1, -2), atol=atol, rtol=rtol)
        if not ok:
            return False, '`x != x.transpose` with atol={}, rtol={}'.format(atol, rtol)
        e, _ = torch.symeig(x, eigenvectors=True)
        ok = (e > -atol).min()
        if not ok:
            return False, 'eigenvalues of x are not all greater than 0.'
        return True, None
//  [python]
def _check_point_on_manifold(self, x: torch.Tensor, *, atol=1e-05, rtol=1e-05) ->Union[Tuple[bool, Optional[str]], bool]:
        ok = torch.allclose(x, x.transpose(-1, -2), atol=atol, rtol=rtol)
        if not ok:
            return False, '`x != x.transpose` with atol={}, rtol={}'.format(atol, rtol)
        e, _ = torch.linalg.eigh(x, 'U')
        ok = (e > -atol).min()
        if not ok:
            return False, 'eigenvalues of x are not all greater than 0.'
        return True, None
//  [python]
def m_transfer_chrom(inp, src):
    '''
    inp (3, M):              input image
    src (3, M):              source image

    out (3, M):              output image
    '''
    
    m_i = torch.mean(inp, dim=1)
    m_s = torch.mean(src, dim=1)
    s_i = torch.maximum(torch.cov(inp), 7.5*torch.eye(3).to(inp.get_device())) 
    s_s = torch.maximum(torch.cov(src), 7.5*torch.eye(3).to(inp.get_device()))

    w_i, v_i = torch.symeig(s_i, eigenvectors=True)
    w_i[w_i<0] = 0
    da = torch.diag(torch.sqrt(w_i + 1e-10))
    c = torch.mm(da, torch.mm(v_i.T, torch.mm(s_s, torch.mm(v_i, da))))
    w_c, v_c = torch.symeig(c, eigenvectors=True)
    w_c[w_c<0] = 0
    dc = torch.diag(torch.sqrt(w_c + 1e-10))
    da_inv = torch.diag(1/(torch.diag(da)))
    T = torch.mm(v_i, torch.mm(da_inv, torch.mm(v_c, torch.mm(dc, torch.mm(v_c.T, torch.mm(da_inv, v_i.T))))))

    out = torch.matmul(T, (inp - m_i[:,None])) + m_s[:,None]

    return out
//  [python]
def m_transfer_chrom(inp, src):
    '''
    inp (3, M):              input image
    src (3, M):              source image

    out (3, M):              output image
    '''
    
    m_i = torch.mean(inp, dim=1)
    m_s = torch.mean(src, dim=1)
    s_i = torch.maximum(torch.cov(inp), 7.5*torch.eye(3).to(inp.get_device())) 
    s_s = torch.maximum(torch.cov(src), 7.5*torch.eye(3).to(inp.get_device()))

    w_i, v_i = torch.linalg.eigh(s_i)
    w_i[w_i<0] = 0
    da = torch.diag(torch.sqrt(w_i + 1e-10))
    c = torch.mm(da, torch.mm(v_i.T, torch.mm(s_s, torch.mm(v_i, da))))
    w_c, v_c = torch.linalg.eigh(c)
    w_c[w_c<0] = 0
    dc = torch.diag(torch.sqrt(w_c + 1e-10))
    da_inv = torch.diag(1/(torch.diag(da)))
    T = torch.mm(v_i, torch.mm(da_inv, torch.mm(v_c, torch.mm(dc, torch.mm(v_c.T, torch.mm(da_inv, v_i.T))))))

    out = torch.matmul(T, (inp - m_i[:,None])) + m_s[:,None]

    return out
//  [python]
def rot_to_quat(rot: torch.Tensor) -> torch.Tensor:
    if rot.shape[-2:] != (3, 3):
        raise ValueError("Input rotation is incorrectly shaped")

    [[xx, xy, xz], [yx, yy, yz], [zx, zy, zz]] = [
        [rot[..., i, j] for j in range(3)] for i in range(3)
    ]

    k = [
        [
            xx + yy + zz,
            zy - yz,
            xz - zx,
            yx - xy,
        ],
        [
            zy - yz,
            xx - yy - zz,
            xy + yx,
            xz + zx,
        ],
        [
            xz - zx,
            xy + yx,
            yy - xx - zz,
            yz + zy,
        ],
        [
            yx - xy,
            xz + zx,
            yz + zy,
            zz - xx - yy,
        ],
    ]

    _, vectors = torch.symeig(
        (1.0 / 3.0) * torch.stack([torch.stack(t, dim=-1) for t in k], dim=-2)
    )
    return vectors[..., -1]
//  [python]
def rot_to_quat(rot: torch.Tensor) -> torch.Tensor:
    if rot.shape[-2:] != (3, 3):
        raise ValueError("Input rotation is incorrectly shaped")

    [[xx, xy, xz], [yx, yy, yz], [zx, zy, zz]] = [
        [rot[..., i, j] for j in range(3)] for i in range(3)
    ]

    k = [
        [
            xx + yy + zz,
            zy - yz,
            xz - zx,
            yx - xy,
        ],
        [
            zy - yz,
            xx - yy - zz,
            xy + yx,
            xz + zx,
        ],
        [
            xz - zx,
            xy + yx,
            yy - xx - zz,
            yz + zy,
        ],
        [
            yx - xy,
            xz + zx,
            yz + zy,
            zz - xx - yy,
        ],
    ]

    _, vectors = torch.linalg.eigh(
        (1.0 / 3.0) * torch.stack([torch.stack(t, dim=-1) for t in k], dim=-2)
    )
    return vectors[..., -1]
//  [python]
def z_so3_pca(self, z_so3):
        z_W = (z_so3[..., None] @ z_so3.unsqueeze(2)).sum(1)
        success = False
        while not success:
            try:
                var, z_so3_basis = torch.symeig(z_W.double(), eigenvectors=True)  # B,3,basis
                success = True
            except:
                logging.warning("torch.symeig fail, retry")
        z_so3_basis = z_so3_basis.float()
        z_so3_proj = torch.einsum("bci,bij->bcj", z_so3, z_so3_basis)
        return z_so3_proj, z_so3_basis, var.float()
//  [python]
def z_so3_pca(self, z_so3):
        z_W = (z_so3[..., None] @ z_so3.unsqueeze(2)).sum(1)
        success = False
        while not success:
            try:
                var, z_so3_basis = torch.linalg.eigh(z_W.double())  # B,3,basis
                success = True
            except:
                logging.warning("torch.linalg.eigh fail, retry")
        z_so3_basis = z_so3_basis.float()
        z_so3_proj = torch.einsum("bci,bij->bcj", z_so3, z_so3_basis)
        return z_so3_proj, z_so3_basis, var.float()
//  [python]
def eigh_wrapper(V, use_cpu=True):

    if use_cpu:
        dev = V.device

        # diagonal loading factor
        # dload = 1e-5 * pt.eye(V.shape[-1], device=V.device, dtype=V.dtype)
        # V = (V + dload).cpu()
        V = V.cpu()
        e_val, e_vec = t.linalg.eigh(V)

        return e_val.to(dev), e_vec.to(dev)
    else:
        return t.linalg.eigh(V)
//  [python]
def eigh_wrapper(V, use_cpu=True):

    if use_cpu:
        dev = V.device

        # diagonal loading factor
        # dload = 1e-5 * pt.eye(V.shape[-1], device=V.device, dtype=V.dtype)
        # V = (V + dload).cpu()
        V = V.cpu()
        e_val, e_vec = pt.linalg.eigh(V)

        return e_val.to(dev), e_vec.to(dev)
    else:
        return pt.linalg.eigh(V)
//  [python]
def mat2quat(rmat: torch.Tensor):
    """
    Converts given rotation matrix to quaternion.
    Args:
        rmat: 3x3 rotation matrix
    Returns:
        vec4 float quaternion angles
    """
    M = rmat[:3, :3]

    m00 = float(M[0, 0])
    m01 = float(M[0, 1])
    m02 = float(M[0, 2])
    m10 = float(M[1, 0])
    m11 = float(M[1, 1])
    m12 = float(M[1, 2])
    m20 = float(M[2, 0])
    m21 = float(M[2, 1])
    m22 = float(M[2, 2])
    # symmetric matrix K
    K = torch.tensor(
        [
            [m00 - m11 - m22, 0.0, 0.0, 0.0],
            [m01 + m10, m11 - m00 - m22, 0.0, 0.0],
            [m02 + m20, m12 + m21, m22 - m00 - m11, 0.0],
            [m21 - m12, m02 - m20, m10 - m01, m00 + m11 + m22],
        ]
    ).to(rmat.device)
    K /= 3.0
    # quaternion is Eigen vector of K that corresponds to largest eigenvalue
    w, V = torch.symeig(K,eigenvectors=True)
    inds = torch.tensor([3, 0, 1, 2])
    q1 = V[inds, torch.argmax(w)]
    if q1[0] < 0.0:
        q1 = -q1
    inds = torch.tensor([1, 2, 3, 0])
    return q1[inds]
//  [python]
def mat2quat(rmat: torch.Tensor):
    """
    Converts given rotation matrix to quaternion.
    Args:
        rmat: 3x3 rotation matrix
    Returns:
        vec4 float quaternion angles
    """
    M = rmat[:3, :3]

    m00 = float(M[0, 0])
    m01 = float(M[0, 1])
    m02 = float(M[0, 2])
    m10 = float(M[1, 0])
    m11 = float(M[1, 1])
    m12 = float(M[1, 2])
    m20 = float(M[2, 0])
    m21 = float(M[2, 1])
    m22 = float(M[2, 2])
    # symmetric matrix K
    K = torch.tensor(
        [
            [m00 - m11 - m22, 0.0, 0.0, 0.0],
            [m01 + m10, m11 - m00 - m22, 0.0, 0.0],
            [m02 + m20, m12 + m21, m22 - m00 - m11, 0.0],
            [m21 - m12, m02 - m20, m10 - m01, m00 + m11 + m22],
        ]
    ).to(rmat.device)
    K /= 3.0
    # quaternion is Eigen vector of K that corresponds to largest eigenvalue
    w, V = torch.linalg.eigh(K)
    inds = torch.tensor([3, 0, 1, 2])
    q1 = V[inds, torch.argmax(w)]
    if q1[0] < 0.0:
        q1 = -q1
    inds = torch.tensor([1, 2, 3, 0])
    return q1[inds]
//  [python]
def _null_space(m, kernel_dim):
    """Finds the null space (kernel) basis of the matrix
    Args:
        m: the batch of input matrices, B x N x 12
        kernel_dim: number of dimensions to approximate the kernel
    Returns:
        * a batch of null space basis vectors
            of size B x 4 x 3 x kernel_dim
        * a batch of spectral values where near-0s correspond to actual
            kernel vectors, of size B x kernel_dim
    """
    mTm = torch.bmm(m.transpose(1, 2), m)
    s, v = torch.symeig(mTm)
    return v[:, :, :kernel_dim].reshape(-1, 4, 3, kernel_dim), s[:, :kernel_dim]
//  [python]
def _null_space(m, kernel_dim):
    """Finds the null space (kernel) basis of the matrix
    Args:
        m: the batch of input matrices, B x N x 12
        kernel_dim: number of dimensions to approximate the kernel
    Returns:
        * a batch of null space basis vectors
            of size B x 4 x 3 x kernel_dim
        * a batch of spectral values where near-0s correspond to actual
            kernel vectors, of size B x kernel_dim
    """
    mTm = torch.bmm(m.transpose(1, 2), m)
    s, v = torch.linalg.eigh(mTm)
    return v[:, :, :kernel_dim].reshape(-1, 4, 3, kernel_dim), s[:, :kernel_dim]
//  [python]
def eigs_quiver_with_conds(ax, A_learned, n, x_lims, y_lims):
    x = np.linspace(x_lims[0], x_lims[1], n)
    y = np.linspace(y_lims[0], y_lims[1], n)
    X, Y = np.meshgrid(x, y)
    U1_learned = np.zeros((n, n))  # x-component of vector field
    V1_learned = np.zeros((n, n))  # y-component of vector field
    log_ratios = np.zeros((n, n))
    for i in range(n):
        for j in range(n):
            pt = torch.tensor([[X[i, j], Y[i, j]]], device=device).double()
            M_learned = A_learned(pt).squeeze().detach()
            sorted_evals_learned, sorted_evecs_learned = torch.symeig(M_learned)
            U1_learned[i, j] = sorted_evecs_learned[0, 0]
            V1_learned[i, j] = sorted_evecs_learned[1, 0]
            log_ratios[i,j] = torch.log10(sorted_evals_learned[1]/sorted_evals_learned[0])
    im = ax.pcolor(X, Y, log_ratios, cmap="magma")
    return im
//  [python]
def eigs_quiver_with_conds(ax, A_learned, n, x_lims, y_lims):
    x = np.linspace(x_lims[0], x_lims[1], n)
    y = np.linspace(y_lims[0], y_lims[1], n)
    X, Y = np.meshgrid(x, y)
    U1_learned = np.zeros((n, n))  # x-component of vector field
    V1_learned = np.zeros((n, n))  # y-component of vector field
    log_ratios = np.zeros((n, n))
    for i in range(n):
        for j in range(n):
            pt = torch.tensor([[X[i, j], Y[i, j]]], device=device).double()
            M_learned = A_learned(pt).squeeze().detach()
            sorted_evals_learned, sorted_evecs_learned = torch.linalg.eigh(M_learned)
            U1_learned[i, j] = sorted_evecs_learned[0, 0]
            V1_learned[i, j] = sorted_evecs_learned[1, 0]
            log_ratios[i,j] = torch.log10(sorted_evals_learned[1]/sorted_evals_learned[0])
    im = ax.pcolor(X, Y, log_ratios, cmap="magma")
    return im
//  [python]
def eigs_quiver(R, n, x_lims, y_lims, color="black"):
    x = np.linspace(x_lims[0], x_lims[1], n)
    print(x)
    y = np.linspace(y_lims[0], y_lims[1], n)
    X, Y = np.meshgrid(x, y)
    U1 = np.zeros((n, n))  # x-component of vector field
    V1 = np.zeros((n, n))  # y-component of vector field
    U2 = np.zeros((n, n))
    V2 = np.zeros((n, n))
    log_ratios = np.zeros((n, n))
    ratios = np.zeros((n, n))
    small_eigs = np.zeros((n, n))
    for i in range(n):
        for j in range(n):
            pt = torch.tensor([[X[i, j], Y[i, j]]], device=device).double()
            M = R(pt).squeeze().detach()
            sorted_evals, sorted_evecs = torch.symeig(M)
            small_eigs[i, j] = sorted_evals[0]
            if sorted_evals[0] == 0:
                log_ratios[i,j] = -1
                ratios[i,j] = -1
            else:
                log_ratios[i,j] = torch.log10(sorted_evals[1]/sorted_evals[0])
                ratios[i,j] = sorted_evals[1]/sorted_evals[0]
            U1[i, j] = sorted_evecs[0, 0]
            V1[i, j] = sorted_evecs[1, 0]
            U2[i, j] = sorted_evecs[0, 1]
            V2[i, j] = sorted_evecs[1, 1]
    plt.pcolor(X, Y, log_ratios)
    plt.colorbar()
    plt.quiver(X, Y, U2, V2, color=color) # was U1, V1
    plt.quiver(X, Y, -U2, -V2, color=color)
//  [python]
def eigs_quiver(R, n, x_lims, y_lims, color="black"):
    x = np.linspace(x_lims[0], x_lims[1], n)
    print(x)
    y = np.linspace(y_lims[0], y_lims[1], n)
    X, Y = np.meshgrid(x, y)
    U1 = np.zeros((n, n))  # x-component of vector field
    V1 = np.zeros((n, n))  # y-component of vector field
    U2 = np.zeros((n, n))
    V2 = np.zeros((n, n))
    log_ratios = np.zeros((n, n))
    ratios = np.zeros((n, n))
    small_eigs = np.zeros((n, n))
    for i in range(n):
        for j in range(n):
            pt = torch.tensor([[X[i, j], Y[i, j]]], device=device).double()
            M = R(pt).squeeze().detach()
            sorted_evals, sorted_evecs = torch.linalg.eigh(M)
            small_eigs[i, j] = sorted_evals[0]
            if sorted_evals[0] == 0:
                log_ratios[i,j] = -1
                ratios[i,j] = -1
            else:
                log_ratios[i,j] = torch.log10(sorted_evals[1]/sorted_evals[0])
                ratios[i,j] = sorted_evals[1]/sorted_evals[0]
            U1[i, j] = sorted_evecs[0, 0]
            V1[i, j] = sorted_evecs[1, 0]
            U2[i, j] = sorted_evecs[0, 1]
            V2[i, j] = sorted_evecs[1, 1]
    plt.pcolor(X, Y, log_ratios)
    plt.colorbar()
    plt.quiver(X, Y, U2, V2, color=color) # was U1, V1
    plt.quiver(X, Y, -U2, -V2, color=color)
//  [python]
def forward(self,previousSample,energy_function,numSteps=1,stepSize=0.01,k=5,**kwargs):
        new_draw = previousSample.clone().detach().requires_grad_(True)
        #prev_U,gradient,hess = energy_function(kwargs["decoder"],new_draw)
        energy,hess = energy_function(kwargs["decoder"],new_draw)

        if isinstance(energy,str):
            if energy=="nan":
                return "eig failed"
        #print(check_symmetric(hess))
        #norm_gradient = gradient/torch.norm(gradient)

        #norm_gradient = gradient/torch.norm(gradient)
        for _ in range(1):
            #U,S,V = torch.linalg.svd(hess)
            #Mntm = torch.randn_like(new_draw)
            #w = Mntm - torch.mm(Mntm,torch.unsqueeze(norm_gradient,-1))*norm_gradient

            decomposition = torch.symeig(hess)
            S = decomposition.eigenvalues
            V = decomposition.eigenvectors
            S,indices = torch.sort(S,descending=False)
            V = V[:,indices]
            num_components = 4
            Vk = V[:,:num_components]
            Sk = S[:num_components]
            beta = torch.randn(num_components)
            beta = beta/torch.norm(beta)
            direction = torch.zeros(new_draw.size()).float().cuda()

            for compIndex in range(num_components):
                #if Sk[compIndex]>0:
                direction += beta[compIndex]*Vk[:,compIndex]
                #else:
                #    direction += beta[compIndex]*(-1*Vk[:,compIndex])

            norm_direction = torch.norm(direction)
            direction = direction/norm_direction

            delta = 1e-3
            betasq_lambda = 0.0
            for compIndex in range(num_components):
                betasq_lambda += ((beta[compIndex]/norm_direction)**2 * torch.abs(Sk[compIndex]))
               
            stepSize = torch.sqrt((2.0*delta)/betasq_lambda)
            #stepSize = max(0.3,min(0.5,stepSize))
            #stepSize = max(0.3,min(0.5,stepSize))
            stepSize = min(stepSize,2.0)
            print(stepSize)
            #stepSize = 0.1

            new_draw = (new_draw+stepSize*direction).clone().detach().requires_grad_(True)
            #current_U,gradient,hess = energy_function(kwargs["decoder"],new_draw)
            #print("Energy:",current_U)
            #norm_gradient = gradient/torch.norm(gradient)
        draw_final = new_draw.clone().detach()

        return draw_final
//  [python]
def forward(self,previousSample,energy_function,numSteps=1,stepSize=0.01,k=5,**kwargs):
        new_draw = previousSample.clone().detach().requires_grad_(True)
        #prev_U,gradient,hess = energy_function(kwargs["decoder"],new_draw)
        energy,hess = energy_function(kwargs["decoder"],new_draw)

        if isinstance(energy,str):
            if energy=="nan":
                return "eig failed"
        #print(check_symmetric(hess))
        #norm_gradient = gradient/torch.norm(gradient)

        #norm_gradient = gradient/torch.norm(gradient)
        for _ in range(1):
            #U,S,V = torch.linalg.svd(hess)
            #Mntm = torch.randn_like(new_draw)
            #w = Mntm - torch.mm(Mntm,torch.unsqueeze(norm_gradient,-1))*norm_gradient

            decomposition = torch.linalg.eigh(hess)
            S = decomposition.eigenvalues
            V = decomposition.eigenvectors
            S,indices = torch.sort(S,descending=False)
            V = V[:,indices]
            num_components = 4
            Vk = V[:,:num_components]
            Sk = S[:num_components]
            beta = torch.randn(num_components)
            beta = beta/torch.norm(beta)
            direction = torch.zeros(new_draw.size()).float().cuda()

            for compIndex in range(num_components):
                #if Sk[compIndex]>0:
                direction += beta[compIndex]*Vk[:,compIndex]
                #else:
                #    direction += beta[compIndex]*(-1*Vk[:,compIndex])

            norm_direction = torch.norm(direction)
            direction = direction/norm_direction

            delta = 1e-3
            betasq_lambda = 0.0
            for compIndex in range(num_components):
                betasq_lambda += ((beta[compIndex]/norm_direction)**2 * torch.abs(Sk[compIndex]))
               
            stepSize = torch.sqrt((2.0*delta)/betasq_lambda)
            #stepSize = max(0.3,min(0.5,stepSize))
            #stepSize = max(0.3,min(0.5,stepSize))
            stepSize = min(stepSize,2.0)
            print(stepSize)
            #stepSize = 0.1

            new_draw = (new_draw+stepSize*direction).clone().detach().requires_grad_(True)
            #current_U,gradient,hess = energy_function(kwargs["decoder"],new_draw)
            #print("Energy:",current_U)
            #norm_gradient = gradient/torch.norm(gradient)
        draw_final = new_draw.clone().detach()

        return draw_final
//  [python]
def rotation_10d_to_quaternion(rotation: torch.Tensor) -> torch.Tensor:
    """
    Convert a 10-vector into a symmetric matrix, whose eigenvector corresponding
    to the eigenvalue of minimum modulus is the resulting quaternion.

    Source: https://arxiv.org/abs/2006.01031
    """
    A = _10vec_to_4x4symmetric(rotation)  # A is a symmetric data matrix
    evals, evecs = torch.symeig(A, eigenvectors=True)
    return evecs[..., evals.argmin(dim=-1)]
//  [python]
def rotation_10d_to_quaternion(rotation: torch.Tensor) -> torch.Tensor:
    """
    Convert a 10-vector into a symmetric matrix, whose eigenvector corresponding
    to the eigenvalue of minimum modulus is the resulting quaternion.

    Source: https://arxiv.org/abs/2006.01031
    """
    A = _10vec_to_4x4symmetric(rotation)  # A is a symmetric data matrix
    return torch.linalg.eigh(A).eigenvectors[..., 0]
//  [python]
def __init__(self, mu, Sigma, batch_size, mylen=int(1e9)):
        self.batch_size = batch_size
        self.mu = mu
        self.eigvals, self.eigvects = torch.symeig(Sigma)
        s_half = (self.eigvects @ torch.diag(torch.sqrt(self.eigvals)) @ self.eigvects.t())
        self.s_half = s_half # Sigma^{1/2}
        self.mylen = mylen
//  [python]
def __init__(self, mu, Sigma, batch_size, mylen=int(1e9)):
        self.batch_size = batch_size
        self.mu = mu
        self.eigvals, self.eigvects = torch.linalg.eigh(Sigma)
        s_half = (self.eigvects @ torch.diag(torch.sqrt(self.eigvals)) @ self.eigvects.t())
        self.s_half = s_half # Sigma^{1/2}
        self.mylen = mylen
//  [python]
def eigs_similarity_metric(A_true, A_learned, n, x_lims, y_lims, space_dims): # x_lims, y_lims was box_radius
    #xs = torch.linspace(-box_radius, box_radius, n).to(device)
    #grid_pts = torch.cartesian_prod(*[xs]*space_dims)
    x = torch.linspace(x_lims[0], x_lims[1], n).to(device)
    y = torch.linspace(y_lims[0], y_lims[1], n).to(device)
    grid_pts = torch.cartesian_prod(x, y)
    true_A_at_pts = A_true(grid_pts)
    learned_A_at_pts = A_learned(grid_pts)
    true_A_evals, true_A_evecs = torch.symeig(true_A_at_pts)
    learned_A_evals, learned_A_evecs = torch.symeig(learned_A_at_pts)
    true_A_evecs_T = torch.transpose(true_A_evecs, 1, 2)
    inner_prods = torch.bmm(true_A_evecs_T, torch.flip(learned_A_evecs, dims=(2,))) # reverse order of learned_A_evecs
    diags = torch.diagonal(inner_prods, dim1=1, dim2=2)
    similarity = torch.linalg.vector_norm(diags, ord=1)
    similarity = similarity/(diags.shape[0]*diags.shape[1])
    return similarity
//  [python]
def eigs_similarity_metric(A_true, A_learned, n, x_lims, y_lims, space_dims): # x_lims, y_lims was box_radius
    #xs = torch.linspace(-box_radius, box_radius, n).to(device)
    #grid_pts = torch.cartesian_prod(*[xs]*space_dims)
    x = torch.linspace(x_lims[0], x_lims[1], n).to(device)
    y = torch.linspace(y_lims[0], y_lims[1], n).to(device)
    grid_pts = torch.cartesian_prod(x, y)
    true_A_at_pts = A_true(grid_pts)
    learned_A_at_pts = A_learned(grid_pts)
    true_A_evals, true_A_evecs = torch.linalg.eigh(true_A_at_pts)
    learned_A_evals, learned_A_evecs = torch.linalg.eigh(learned_A_at_pts)
    true_A_evecs_T = torch.transpose(true_A_evecs, 1, 2)
    inner_prods = torch.bmm(true_A_evecs_T, torch.flip(learned_A_evecs, dims=(2,))) # reverse order of learned_A_evecs
    diags = torch.diagonal(inner_prods, dim1=1, dim2=2)
    similarity = torch.linalg.vector_norm(diags, ord=1)
    similarity = similarity/(diags.shape[0]*diags.shape[1])
    return similarity
//  [python]
def eigs_quiver_comparison(ax, A_true, A_learned, n, x_lims, y_lims):
    x = np.linspace(x_lims[0], x_lims[1], n)
    y = np.linspace(y_lims[0], y_lims[1], n)
    X, Y = np.meshgrid(x, y)
    U1_true = np.zeros((n, n))  # x-component of vector field
    V1_true = np.zeros((n, n))  # y-component of vector field
    U1_learned = np.zeros((n, n))  # x-component of vector field
    V1_learned = np.zeros((n, n))  # y-component of vector field
    U2_learned = np.zeros((n, n))  # x-component of vector field
    V2_learned = np.zeros((n, n))  # y-component of vector field
    for i in range(n):
        for j in range(n):
            pt = torch.tensor([[X[i, j], Y[i, j]]], device=device).double()
            M_true = A_true(pt).squeeze().detach()
            sorted_evals_true, sorted_evecs_true = torch.symeig(M_true)
            U1_true[i, j] = sorted_evecs_true[0, 0]
            V1_true[i, j] = sorted_evecs_true[1, 0]
            M_learned = A_learned(pt).squeeze().detach()
            sorted_evals_learned, sorted_evecs_learned = torch.symeig(M_learned)
            U1_learned[i, j] = sorted_evecs_learned[0, 0]
            V1_learned[i, j] = sorted_evecs_learned[1, 0]
            U2_learned[i, j] = sorted_evecs_learned[0, 1]
            V2_learned[i, j] = sorted_evecs_learned[1, 1]
    ax.quiver(X, Y, U1_true, V1_true, color="#5D3A9B")
    ax.quiver(X, Y, -U1_true, -V1_true, color="#5D3A9B")
    ax.quiver(X, Y, U2_learned, V2_learned, color="#E66100") # was U1_learned, V1_learned
    ax.quiver(X, Y, -U2_learned, -V2_learned, color="#E66100")
//  [python]
def eigs_quiver_comparison(ax, A_true, A_learned, n, x_lims, y_lims):
    x = np.linspace(x_lims[0], x_lims[1], n)
    y = np.linspace(y_lims[0], y_lims[1], n)
    X, Y = np.meshgrid(x, y)
    U1_true = np.zeros((n, n))  # x-component of vector field
    V1_true = np.zeros((n, n))  # y-component of vector field
    U1_learned = np.zeros((n, n))  # x-component of vector field
    V1_learned = np.zeros((n, n))  # y-component of vector field
    U2_learned = np.zeros((n, n))  # x-component of vector field
    V2_learned = np.zeros((n, n))  # y-component of vector field
    for i in range(n):
        for j in range(n):
            pt = torch.tensor([[X[i, j], Y[i, j]]], device=device).double()
            M_true = A_true(pt).squeeze().detach()
            sorted_evals_true, sorted_evecs_true = torch.linalg.eigh(M_true)
            U1_true[i, j] = sorted_evecs_true[0, 0]
            V1_true[i, j] = sorted_evecs_true[1, 0]
            M_learned = A_learned(pt).squeeze().detach()
            sorted_evals_learned, sorted_evecs_learned = torch.linalg.eigh(M_learned)
            U1_learned[i, j] = sorted_evecs_learned[0, 0]
            V1_learned[i, j] = sorted_evecs_learned[1, 0]
            U2_learned[i, j] = sorted_evecs_learned[0, 1]
            V2_learned[i, j] = sorted_evecs_learned[1, 1]
    ax.quiver(X, Y, U1_true, V1_true, color="#5D3A9B")
    ax.quiver(X, Y, -U1_true, -V1_true, color="#5D3A9B")
    ax.quiver(X, Y, U2_learned, V2_learned, color="#E66100") # was U1_learned, V1_learned
    ax.quiver(X, Y, -U2_learned, -V2_learned, color="#E66100")
//  [python]
def eigs_quiver_rescaled(R, n, x_lims, y_lims, scale_factor):
    x = np.linspace(x_lims[0], x_lims[1], n)
    y = np.linspace(y_lims[0], y_lims[1], n)
    X, Y = np.meshgrid(x, y)
    U1 = np.zeros((n, n))  # x-component of vector field
    V1 = np.zeros((n, n))  # y-component of vector field
    U2 = np.zeros((n, n))
    V2 = np.zeros((n, n))
    for i in range(n):
        for j in range(n):
            pt = torch.tensor([[X[i, j], Y[i, j]]], device=device).float()
            M = R(pt).squeeze().detach()
            L, Q = torch.symeig(M)
            # Rescale L
            L = torch.tensor([scale_factor, 1]).float().to(device)
            # Reconstruct M with rescaled L
            M = Q @ torch.diag(L) @ Q.T
            sorted_evals, sorted_evecs = torch.symeig(M)
            U1[i, j] = (1 / sorted_evals[0]) * sorted_evecs[0, 0]
            V1[i, j] = (1 / sorted_evals[0]) * sorted_evecs[1, 0]
            U2[i, j] = (1 / sorted_evals[1]) * sorted_evecs[0, 1]
            V2[i, j] = (1 / sorted_evals[1]) * sorted_evecs[1, 1]
    plt.quiver(X, Y, U1, V1, scale=1e6)
    plt.quiver(X, Y, -U1, -V1, scale=1e6)
    plt.quiver(X, Y, U2, V2, scale=1e6)
    plt.quiver(X, Y, -U2, -V2, scale=1e6)
//  [python]
def eigs_quiver_rescaled(R, n, x_lims, y_lims, scale_factor):
    x = np.linspace(x_lims[0], x_lims[1], n)
    y = np.linspace(y_lims[0], y_lims[1], n)
    X, Y = np.meshgrid(x, y)
    U1 = np.zeros((n, n))  # x-component of vector field
    V1 = np.zeros((n, n))  # y-component of vector field
    U2 = np.zeros((n, n))
    V2 = np.zeros((n, n))
    for i in range(n):
        for j in range(n):
            pt = torch.tensor([[X[i, j], Y[i, j]]], device=device).float()
            M = R(pt).squeeze().detach()
            L, Q = torch.linalg.eigh(M)
            # Rescale L
            L = torch.tensor([scale_factor, 1]).float().to(device)
            # Reconstruct M with rescaled L
            M = Q @ torch.diag(L) @ Q.T
            sorted_evals, sorted_evecs = torch.linalg.eigh(M)
            U1[i, j] = (1 / sorted_evals[0]) * sorted_evecs[0, 0]
            V1[i, j] = (1 / sorted_evals[0]) * sorted_evecs[1, 0]
            U2[i, j] = (1 / sorted_evals[1]) * sorted_evecs[0, 1]
            V2[i, j] = (1 / sorted_evals[1]) * sorted_evecs[1, 1]
    plt.quiver(X, Y, U1, V1, scale=1e6)
    plt.quiver(X, Y, -U1, -V1, scale=1e6)
    plt.quiver(X, Y, U2, V2, scale=1e6)
    plt.quiver(X, Y, -U2, -V2, scale=1e6)