jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def spd_product_manifold_distance_torch(x1, x2, diagonal_distance=False):
    """
    Compute the distance between SPD points considered as a product of a Euclidean and a sphere manifold.

    Parameters
    ----------
    :param x1: set of SPD matrices      (N1 x d x d or b1 x ... x bk x N1 x d x d)
    :param x2: set of SPD matrices      (N2 x d x d or b1 x ... x bk x N1 x d x d)

    Optional parameters
    -------------------
    :param diagonal_distance: Whole distance matrix, or just the diagonal? If True, we must have `x1 == x2`.

    Returns
    -------
    :return: matrix of product-of-manifolds distance between SPD points in x1 and x2  (N1 x N2 or b1 x ... x bk x N1 x N2)
    """
    # If diag, x1 must be equal to x2 and we can return zeros.
    if diagonal_distance is True:
        shape = list(x2.shape)[:-2]
        shape.append(1)
        return torch.zeros(shape, dtype=x1.dtype)

    # Decompose SPD matrices in sphere (from SO(n)) and Euclidean parts using eigenvalue decompositions
    x_euclidean1, x_s1 = spd_to_sphere_and_euclidean_torch(x1)
    x_euclidean2, x_s2 = spd_to_sphere_and_euclidean_torch(x2)

    # Compute sphere distance
    distance_sphere = sphere_distance_torch(x_s1, x_s2)

    # Compute Euclidean distance
    distance_euclidean = torch.linalg.norm(x_euclidean1 - x_euclidean2, dim=-1).unsqueeze(-1)

    # Compute product-of-manifold distance
    distance = torch.sqrt(distance_sphere**2 + distance_euclidean**2)

    return distance.double()
//  [python]
def check_if_mirror_extended(self, lrs):
        """Check whether the input is a mirror-extended sequence.
        If mirror-extended, the i-th (i=0, ..., t-1) frame is equal to the
        (t-1-i)-th frame.
        Args:
            lrs (tensor): Input LR images with shape (n, t, c, h, w)
        """

        self.is_mirror_extended = False
        if lrs.size(1) % 2 == 0:
            lrs_1, lrs_2 = torch.chunk(lrs, 2, dim=1)
            if torch.norm(lrs_1 - lrs_2.flip(1)) == 0:
                self.is_mirror_extended = True
//  [python]
def check_if_mirror_extended(self, lrs):
        """Check whether the input is a mirror-extended sequence.
        If mirror-extended, the i-th (i=0, ..., t-1) frame is equal to the
        (t-1-i)-th frame.
        Args:
            lrs (tensor): Input LR images with shape (n, t, c, h, w)
        """

        self.is_mirror_extended = False
        if lrs.size(1) % 2 == 0:
            lrs_1, lrs_2 = torch.chunk(lrs, 2, dim=1)
            if torch.linalg.norm(lrs_1 - lrs_2.flip(1)) == 0:
                self.is_mirror_extended = True
//  [python]
def extend_lowdim(mod,
                  pose_gt, 
                  num_steps=500,
                  with_tqdm=False):
    """ Directly fitting via frobenius norm. """
    mod.ths = nn.Parameter(mod.ths.new_zeros(len(pose_gt)), requires_grad=True)
    optimizer = Adam(
        params=[mod.ths],
        lr=1e-1)

    pbar = tqdm.tqdm(range(num_steps)) if with_tqdm else range(num_steps)
    for i in pbar:
        optimizer.zero_grad()
        pose_pred = mod.pose()
        rot_diff = pose_gt.rot_mats - pose_pred.rot_mats
        frob_norm = torch.norm(rot_diff, p='fro')
        tot_err = frob_norm
        tot_err.backward()
        optimizer.step()
        if with_tqdm:
            pbar.set_description(str(tot_err))
    del optimizer
    torch.cuda.empty_cache()
    return mod
//  [python]
def extend_lowdim(mod,
                  pose_gt, 
                  num_steps=500,
                  with_tqdm=False):
    """ Directly fitting via frobenius norm. """
    mod.ths = nn.Parameter(mod.ths.new_zeros(len(pose_gt)), requires_grad=True)
    optimizer = Adam(
        params=[mod.ths],
        lr=1e-1)

    pbar = tqdm.tqdm(range(num_steps)) if with_tqdm else range(num_steps)
    for i in pbar:
        optimizer.zero_grad()
        pose_pred = mod.pose()
        rot_diff = pose_gt.rot_mats - pose_pred.rot_mats
        frob_norm = torch.linalg.norm(rot_diff, ord='fro')
        tot_err = frob_norm
        tot_err.backward()
        optimizer.step()
        if with_tqdm:
            pbar.set_description(str(tot_err))
    del optimizer
    torch.cuda.empty_cache()
    return mod
//  [python]
def render(self, mode='human', showerr=False):
    if not showerr:
      print('%%%%%%%%%%%%%%%%%%%%%%')
      print('%f %f'%(self.rho[0],self.rho[1])) 
      for i in range(self.M):
        print('%d %f %f'%(i,self.x0[i],self.x[i]))
      print('%%%%%%%%%%%%%%%%%%%%%%')
      print('%e %e %f'%(self.rho[0],self.rho[1],torch.norm(self.x0-self.x)))
    else:
      print('%e %e %f'%(self.rho[0],self.rho[1],torch.norm(self.x0-self.x)))
//  [python]
def render(self, mode='human', showerr=False):
    if not showerr:
      print('%%%%%%%%%%%%%%%%%%%%%%')
      print('%f %f'%(self.rho[0],self.rho[1])) 
      for i in range(self.M):
        print('%d %f %f'%(i,self.x0[i],self.x[i]))
      print('%%%%%%%%%%%%%%%%%%%%%%')
      print('%e %e %f'%(self.rho[0],self.rho[1],torch.linalg.norm(self.x0-self.x)))
    else:
      print('%e %e %f'%(self.rho[0],self.rho[1],torch.linalg.norm(self.x0-self.x)))
//  [python]
def fit_with_frobenius(version:str,
                       pose_gt, 
                       num_steps=500,
                       with_tqdm=False):
    """ Directly fitting via frobenius norm. """
    if version == 'lowdim':
        mod = LowDimModel(pose_gt.rot_mats, pose_gt.transl, debug=False) # pose_pred = mod.pose()
        optimizer = Adam(
            params=[mod.Q_6d, mod.t2, mod.P_6d, mod.t1, mod.Rw_6d, mod.ths],
            lr=1e-1)
    elif version == 'static':
        mod = FixedModel(pose_gt.rot_mats, pose_gt.transl, debug=False) # pose_pred = mod.pose()
        optimizer = Adam(
            params=[mod.Q_6d, mod.t2],
            lr=1e-1)

    pbar = tqdm.tqdm(range(num_steps)) if with_tqdm else range(num_steps)
    for i in pbar:
        optimizer.zero_grad()
        pose_pred = mod.pose()
        # rot_diff = pose_gt.rot_mats - pose_pred.rot_mats
        pose_diff = pose_gt.pose_mats - pose_pred.pose_mats
        frob_norm = torch.norm(pose_diff, p='fro')
        tot_err = frob_norm
        tot_err.backward()
        optimizer.step()
        if with_tqdm:
            pbar.set_description(str(tot_err))
    del optimizer
    torch.cuda.empty_cache()
    return mod
//  [python]
def fit_with_frobenius(version:str,
                       pose_gt, 
                       num_steps=500,
                       with_tqdm=False):
    """ Directly fitting via frobenius norm. """
    if version == 'lowdim':
        mod = LowDimModel(pose_gt.rot_mats, pose_gt.transl, debug=False) # pose_pred = mod.pose()
        optimizer = Adam(
            params=[mod.Q_6d, mod.t2, mod.P_6d, mod.t1, mod.Rw_6d, mod.ths],
            lr=1e-1)
    elif version == 'static':
        mod = FixedModel(pose_gt.rot_mats, pose_gt.transl, debug=False) # pose_pred = mod.pose()
        optimizer = Adam(
            params=[mod.Q_6d, mod.t2],
            lr=1e-1)

    pbar = tqdm.tqdm(range(num_steps)) if with_tqdm else range(num_steps)
    for i in pbar:
        optimizer.zero_grad()
        pose_pred = mod.pose()
        # rot_diff = pose_gt.rot_mats - pose_pred.rot_mats
        pose_diff = pose_gt.pose_mats - pose_pred.pose_mats
        frob_norm = torch.linalg.norm(pose_diff, ord='fro')
        tot_err = frob_norm
        tot_err.backward()
        optimizer.step()
        if with_tqdm:
            pbar.set_description(str(tot_err))
    del optimizer
    torch.cuda.empty_cache()
    return mod
//  [python]
def reset(self):
    #torch.manual_seed(19)
    # design matrix
    self.A=torch.randn(self.N,self.M,dtype=torch.float32,requires_grad=False,device=mydevice)
    self.A /=torch.norm(self.A) # normalize
    #torch.manual_seed(time.time())
    # ground truth parameters
    self.Mo=int(torch.randint(3,self.M,(1,)).data)
    # ground truth parameters
    z0=torch.randn(self.Mo,dtype=torch.float32,requires_grad=False,device=mydevice)
    self.x0=torch.zeros(self.M,dtype=torch.float32,requires_grad=False,device=mydevice)
    # sparse vector
    self.x0[np.random.randint(0,self.M,self.Mo)]=z0
    self.y0=torch.matmul(self.A,self.x0)

    self.hint=None
    # rho <- rho0
    self.rho=LOW*torch.ones(self.K,dtype=torch.float32)
    observation={'A': self.A.view(-1).cpu(),
            'eig': torch.zeros(self.N,dtype=torch.float32)} 
    return observation
//  [python]
def reset(self):
    #torch.manual_seed(19)
    # design matrix
    self.A=torch.randn(self.N,self.M,dtype=torch.float32,requires_grad=False,device=mydevice)
    self.A /=torch.linalg.norm(self.A) # normalize
    #torch.manual_seed(time.time())
    # ground truth parameters
    self.Mo=int(torch.randint(3,self.M,(1,)).data)
    # ground truth parameters
    z0=torch.randn(self.Mo,dtype=torch.float32,requires_grad=False,device=mydevice)
    self.x0=torch.zeros(self.M,dtype=torch.float32,requires_grad=False,device=mydevice)
    # sparse vector
    self.x0[np.random.randint(0,self.M,self.Mo)]=z0
    self.y0=torch.matmul(self.A,self.x0)

    self.hint=None
    # rho <- rho0
    self.rho=LOW*torch.ones(self.K,dtype=torch.float32)
    observation={'A': self.A.view(-1).cpu(),
            'eig': torch.zeros(self.N,dtype=torch.float32)} 
    return observation
//  [python]
def slerp(val, low, high):
    low_norm = low / torch.norm(low, dim=1, keepdim=True)
    high_norm = high / torch.norm(high, dim=1, keepdim=True)
    omega = torch.acos((low_norm * high_norm).sum(1))
    so = torch.sin(omega)
    res = (torch.sin((1.0 - val) * omega) / so).unsqueeze(1) * low + (torch.sin(val * omega) / so).unsqueeze(1) * high
    return res
//  [python]
def slerp(val, low, high):
    low_norm = low / torch.linalg.norm(low, dim=1, keepdim=True)
    high_norm = high / torch.linalg.norm(high, dim=1, keepdim=True)
    omega = torch.acos((low_norm * high_norm).sum(1))
    so = torch.sin(omega)
    res = (torch.sin((1.0 - val) * omega) / so).unsqueeze(1) * low + (torch.sin(val * omega) / so).unsqueeze(1) * high
    return res
//  [python]
def rigid_from_3_points(N, Ca, C, is_na=None, eps=1e-8):
    dims = N.shape[:-1]

    v1 = C-Ca
    v2 = N-Ca
    e1 = v1/(torch.norm(v1, dim=-1, keepdim=True)+eps)
    u2 = v2-(torch.einsum('...li, ...li -> ...l', e1, v2)[...,None]*e1)
    e2 = u2/(torch.norm(u2, dim=-1, keepdim=True)+eps)
    e3 = torch.cross(e1, e2, dim=-1)
    R = torch.cat([e1[...,None], e2[...,None], e3[...,None]], axis=-1) #[B,L,3,3] - rotation matrix
    
    v2 = v2/(torch.norm(v2, dim=-1, keepdim=True)+eps)
    cosref = torch.sum(e1*v2, dim=-1)
    costgt = torch.full(dims, -0.3616, device=N.device)
    if is_na is not None:
        costgt[is_na] = -0.4929

    cos2del = torch.clamp( cosref*costgt + torch.sqrt((1-cosref*cosref)*(1-costgt*costgt)+eps), min=-1.0, max=1.0 )
    cosdel = torch.sqrt(0.5*(1+cos2del)+eps)
    sindel = torch.sign(costgt-cosref) * torch.sqrt(1-0.5*(1+cos2del)+eps)
    Rp = torch.eye(3, device=N.device).repeat(*dims,1,1)
    Rp[...,0,0] = cosdel
    Rp[...,0,1] = -sindel
    Rp[...,1,0] = sindel
    Rp[...,1,1] = cosdel

    R = torch.einsum('...ij,...jk->...ik', R,Rp)

    return R, Ca
//  [python]
def rigid_from_3_points(N, Ca, C, is_na=None, eps=1e-8):
    dims = N.shape[:-1]

    v1 = C-Ca
    v2 = N-Ca
    e1 = v1/(torch.linalg.norm(v1, dim=-1, keepdim=True)+eps)
    u2 = v2-(torch.einsum('...li, ...li -> ...l', e1, v2)[...,None]*e1)
    e2 = u2/(torch.linalg.norm(u2, dim=-1, keepdim=True)+eps)
    e3 = torch.cross(e1, e2, dim=-1)
    R = torch.cat([e1[...,None], e2[...,None], e3[...,None]], axis=-1) #[B,L,3,3] - rotation matrix
    
    v2 = v2/(torch.linalg.norm(v2, dim=-1, keepdim=True)+eps)
    cosref = torch.sum(e1*v2, dim=-1)
    costgt = torch.full(dims, -0.3616, device=N.device)
    if is_na is not None:
        costgt[is_na] = -0.4929

    cos2del = torch.clamp( cosref*costgt + torch.sqrt((1-cosref*cosref)*(1-costgt*costgt)+eps), min=-1.0, max=1.0 )
    cosdel = torch.sqrt(0.5*(1+cos2del)+eps)
    sindel = torch.sign(costgt-cosref) * torch.sqrt(1-0.5*(1+cos2del)+eps)
    Rp = torch.eye(3, device=N.device).repeat(*dims,1,1)
    Rp[...,0,0] = cosdel
    Rp[...,0,1] = -sindel
    Rp[...,1,0] = sindel
    Rp[...,1,1] = cosdel

    R = torch.einsum('...ij,...jk->...ik', R,Rp)

    return R, Ca
//  [python]
def initsol(self):
    # generate data (by adding noise to noise-free data)
    #torch.manual_seed(time.time())
    # one observation
    n=torch.randn(self.N,dtype=torch.float32,requires_grad=False,device=mydevice)
    self.y=self.y0+self.SNR*torch.norm(self.y0)/torch.norm(n)*n
    # parameters, initialized to zero
    x=torch.zeros(self.M,requires_grad=True,device=mydevice)

    def lossfunction(A,y,x,alpha=self.rho[0],beta=self.rho[1]):
       Ax=torch.matmul(A,x)
       err=y-Ax
       return torch.norm(err,2)**2+alpha*torch.norm(x,2)**2+beta*torch.norm(x,1)

    opt = LBFGSNew([x],history_size=7,max_iter=10,line_search_fn=True,batch_mode=False)

    # find solution x
    for nepoch in  range(0,20):
       def closure():
         if torch.is_grad_enabled():
           opt.zero_grad()
         loss=lossfunction(self.A,self.y,x,self.rho[0],self.rho[1])
         if loss.requires_grad:
           loss.backward()
           #print(loss.data.item())
         return loss

       opt.step(closure)

    self.x=x
//  [python]
def initsol(self):
    # generate data (by adding noise to noise-free data)
    #torch.manual_seed(time.time())
    # one observation
    n=torch.randn(self.N,dtype=torch.float32,requires_grad=False,device=mydevice)
    self.y=self.y0+self.SNR*torch.linalg.norm(self.y0)/torch.linalg.norm(n)*n
    # parameters, initialized to zero
    x=torch.zeros(self.M,requires_grad=True,device=mydevice)

    def lossfunction(A,y,x,alpha=self.rho[0],beta=self.rho[1]):
       Ax=torch.matmul(A,x)
       err=y-Ax
       return torch.linalg.norm(err,2)**2+alpha*torch.linalg.norm(x,2)**2+beta*torch.linalg.norm(x,1)

    opt = LBFGSNew([x],history_size=7,max_iter=10,line_search_fn=True,batch_mode=False)

    # find solution x
    for nepoch in  range(0,20):
       def closure():
         if torch.is_grad_enabled():
           opt.zero_grad()
         loss=lossfunction(self.A,self.y,x,self.rho[0],self.rho[1])
         if loss.requires_grad:
           loss.backward()
           #print(loss.data.item())
         return loss

       opt.step(closure)

    self.x=x
//  [python]
def get_rays(directions: torch.Tensor, c2w: torch.Tensor, near: float, far: float,
             ray_altitude_range: List[float]) -> torch.Tensor:
    # Rotate ray directions from camera coordinate to the world coordinate
    rays_d = directions @ c2w[:, :3].T  # (H, W, 3)
    rays_d = rays_d / torch.norm(rays_d, dim=-1, keepdim=True)

    # The origin of all rays is the camera origin in world coordinate
    rays_o = c2w[:, 3].expand(rays_d.shape)  # (H, W, 3)

    return _get_rays_inner(rays_o, rays_d, near, far, ray_altitude_range)
//  [python]
def get_rays(directions: torch.Tensor, c2w: torch.Tensor, near: float, far: float,
             ray_altitude_range: List[float]) -> torch.Tensor:
    # Rotate ray directions from camera coordinate to the world coordinate
    rays_d = directions @ c2w[:, :3].T  # (H, W, 3)
    rays_d = rays_d / torch.linalg.norm(rays_d, dim=-1, keepdim=True)

    # The origin of all rays is the camera origin in world coordinate
    rays_o = c2w[:, 3].expand(rays_d.shape)  # (H, W, 3)

    return _get_rays_inner(rays_o, rays_d, near, far, ray_altitude_range)
//  [python]
def precompute_terms(cls, base_linear_operator, eager_rhs, logdet_terms=True, include_tmats=True):
        """
        Computes the solves, probe vectors, probe_vector norms, probe vector solves, and probe vector
        tridiagonal matrices to construct a CachedCGLinearOperator

        Set logdet_terms to False if you are not going to compute the logdet of the LinearOperator
        """
        with torch.no_grad():
            if logdet_terms:
                # Generate probe vectors
                num_random_probes = settings.num_trace_samples.value()
                probe_vectors = torch.empty(
                    base_linear_operator.matrix_shape[-1],
                    num_random_probes,
                    dtype=base_linear_operator.dtype,
                    device=base_linear_operator.device,
                )
                probe_vectors.bernoulli_().mul_(2).add_(-1)
                probe_vectors = probe_vectors.expand(
                    *base_linear_operator.batch_shape, base_linear_operator.matrix_shape[-1], num_random_probes
                )
                probe_vector_norms = torch.norm(probe_vectors, 2, dim=-2, keepdim=True)
                probe_vectors = probe_vectors.div(probe_vector_norms)

                # Compute solves
                if include_tmats:
                    all_solves, probe_vector_tmats, = base_linear_operator._solve(
                        torch.cat([probe_vectors, eager_rhs], -1),
                        preconditioner=base_linear_operator._preconditioner()[0],
                        num_tridiag=probe_vectors.size(-1),
                    )
                else:
                    all_solves = base_linear_operator._solve(
                        torch.cat([probe_vectors, eager_rhs], -1),
                        preconditioner=base_linear_operator._preconditioner()[0],
                    )
                    probe_vector_tmats = torch.tensor([])
                probe_vector_solves = all_solves[..., : probe_vectors.size(-1)].detach()
                solves = all_solves[..., probe_vectors.size(-1) :]

                return (
                    solves.detach(),
                    probe_vectors.detach(),
                    probe_vector_norms.detach(),
                    probe_vector_solves.detach(),
                    probe_vector_tmats.detach(),
                )

            else:
                # Compute solves
                if settings.fast_computations.log_prob.on():
                    solves = base_linear_operator._solve(
                        eager_rhs, preconditioner=base_linear_operator._preconditioner()[0]
                    )
                else:
                    solves = base_linear_operator.cholesky()._cholesky_solve(eager_rhs)
                dtype = solves.dtype
                device = solves.device
                return (
                    solves.detach(),
                    torch.tensor([], dtype=dtype, device=device),
                    torch.tensor([], dtype=dtype, device=device),
                    torch.tensor([], dtype=dtype, device=device),
                    torch.tensor([], dtype=dtype, device=device),
                )
//  [python]
def precompute_terms(cls, base_linear_operator, eager_rhs, logdet_terms=True, include_tmats=True):
        """
        Computes the solves, probe vectors, probe_vector norms, probe vector solves, and probe vector
        tridiagonal matrices to construct a CachedCGLinearOperator

        Set logdet_terms to False if you are not going to compute the logdet of the LinearOperator
        """
        with torch.no_grad():
            if logdet_terms:
                # Generate probe vectors
                num_random_probes = settings.num_trace_samples.value()
                probe_vectors = torch.empty(
                    base_linear_operator.matrix_shape[-1],
                    num_random_probes,
                    dtype=base_linear_operator.dtype,
                    device=base_linear_operator.device,
                )
                probe_vectors.bernoulli_().mul_(2).add_(-1)
                probe_vectors = probe_vectors.expand(
                    *base_linear_operator.batch_shape, base_linear_operator.matrix_shape[-1], num_random_probes
                )
                probe_vector_norms = torch.linalg.norm(probe_vectors, 2, dim=-2, keepdim=True)
                probe_vectors = probe_vectors.div(probe_vector_norms)

                # Compute solves
                if include_tmats:
                    all_solves, probe_vector_tmats, = base_linear_operator._solve(
                        torch.cat([probe_vectors, eager_rhs], -1),
                        preconditioner=base_linear_operator._preconditioner()[0],
                        num_tridiag=probe_vectors.size(-1),
                    )
                else:
                    all_solves = base_linear_operator._solve(
                        torch.cat([probe_vectors, eager_rhs], -1),
                        preconditioner=base_linear_operator._preconditioner()[0],
                    )
                    probe_vector_tmats = torch.tensor([])
                probe_vector_solves = all_solves[..., : probe_vectors.size(-1)].detach()
                solves = all_solves[..., probe_vectors.size(-1) :]

                return (
                    solves.detach(),
                    probe_vectors.detach(),
                    probe_vector_norms.detach(),
                    probe_vector_solves.detach(),
                    probe_vector_tmats.detach(),
                )

            else:
                # Compute solves
                if settings.fast_computations.log_prob.on():
                    solves = base_linear_operator._solve(
                        eager_rhs, preconditioner=base_linear_operator._preconditioner()[0]
                    )
                else:
                    solves = base_linear_operator.cholesky()._cholesky_solve(eager_rhs)
                dtype = solves.dtype
                device = solves.device
                return (
                    solves.detach(),
                    torch.tensor([], dtype=dtype, device=device),
                    torch.tensor([], dtype=dtype, device=device),
                    torch.tensor([], dtype=dtype, device=device),
                    torch.tensor([], dtype=dtype, device=device),
                )
//  [python]
def __init__(self, M=5, N=15, provide_hint=False):
    super(ENetEnv, self).__init__()
    # N=rows=datapoints, M=columns=parameters
    # Define action and observation space
    # possible actions: vector Kx1, K=2 to match rho
    self.K=2
    self.N=N
    self.M=M
    # actions: 0,1,..,K-1 : action K : scaled/shifted from actor output
    self.action_space = spaces.Box(low=np.zeros((self.K,1))*LOW,high=np.ones((self.K,1))*HIGH,dtype=np.float32)
    # observation (state space): two parts
    # 'A': design matrix -inf,...,inf : normalized by 1/||A||
    # 'eig': 1+eigenvalues (real), Mx1 real vector -inf,...,inf
    self.observation_space = spaces.Dict({
      'A': spaces.Box(low=np.zeros((self.N,self.M))*(-HIGH),high=np.ones((self.N,self.M))*HIGH,dtype=np.float32), 
      'eig': spaces.Box(low=np.ones((self.N,1))*(-HIGH),high=np.ones((self.N,1))*HIGH,dtype=np.float32)
       })

    self.SNR=0.1 # noise ratio
    # generate design matrix (A), noise-free data vector (y0), ground truth parameter vector (x0)
    #torch.manual_seed(19)
    # design matrix
    self.A=torch.randn(N,M,dtype=torch.float32,requires_grad=False,device=mydevice)
    self.A /=torch.norm(self.A) # normalize
    #torch.manual_seed(time.time())
    # how much sparsity (non-zero values)
    self.Mo=int(torch.randint(3,M,(1,)).data)
    # ground truth parameters
    z0=torch.randn(self.Mo,dtype=torch.float32,requires_grad=False,device=mydevice)
    self.x0=torch.zeros(M,dtype=torch.float32,requires_grad=False,device=mydevice)
    # sparse vector
    self.x0[np.random.randint(0,M,self.Mo)]=z0
    self.y0=torch.matmul(self.A,self.x0)
    # regularization factors for L2 and L1
    self.rho=LOW*torch.ones(self.K,dtype=torch.float32)
    # following used for rendering
    self.x=torch.zeros(M,dtype=torch.float32,requires_grad=False,device=mydevice)
    # following for evaluation (keeping noise the same)
    self.y=None
    self.hint=None
    self.provide_hint=provide_hint
//  [python]
def __init__(self, M=5, N=15, provide_hint=False):
    super(ENetEnv, self).__init__()
    # N=rows=datapoints, M=columns=parameters
    # Define action and observation space
    # possible actions: vector Kx1, K=2 to match rho
    self.K=2
    self.N=N
    self.M=M
    # actions: 0,1,..,K-1 : action K : scaled/shifted from actor output
    self.action_space = spaces.Box(low=np.zeros((self.K,1))*LOW,high=np.ones((self.K,1))*HIGH,dtype=np.float32)
    # observation (state space): two parts
    # 'A': design matrix -inf,...,inf : normalized by 1/||A||
    # 'eig': 1+eigenvalues (real), Mx1 real vector -inf,...,inf
    self.observation_space = spaces.Dict({
      'A': spaces.Box(low=np.zeros((self.N,self.M))*(-HIGH),high=np.ones((self.N,self.M))*HIGH,dtype=np.float32), 
      'eig': spaces.Box(low=np.ones((self.N,1))*(-HIGH),high=np.ones((self.N,1))*HIGH,dtype=np.float32)
       })

    self.SNR=0.1 # noise ratio
    # generate design matrix (A), noise-free data vector (y0), ground truth parameter vector (x0)
    #torch.manual_seed(19)
    # design matrix
    self.A=torch.randn(N,M,dtype=torch.float32,requires_grad=False,device=mydevice)
    self.A /=torch.linalg.norm(self.A) # normalize
    #torch.manual_seed(time.time())
    # how much sparsity (non-zero values)
    self.Mo=int(torch.randint(3,M,(1,)).data)
    # ground truth parameters
    z0=torch.randn(self.Mo,dtype=torch.float32,requires_grad=False,device=mydevice)
    self.x0=torch.zeros(M,dtype=torch.float32,requires_grad=False,device=mydevice)
    # sparse vector
    self.x0[np.random.randint(0,M,self.Mo)]=z0
    self.y0=torch.matmul(self.A,self.x0)
    # regularization factors for L2 and L1
    self.rho=LOW*torch.ones(self.K,dtype=torch.float32)
    # following used for rendering
    self.x=torch.zeros(M,dtype=torch.float32,requires_grad=False,device=mydevice)
    # following for evaluation (keeping noise the same)
    self.y=None
    self.hint=None
    self.provide_hint=provide_hint
//  [python]
def get_rays_batch(directions: torch.Tensor, c2w: torch.Tensor, near: float, far: float,
                   ray_altitude_range: List[float]) -> torch.Tensor:
    # Rotate ray directions from camera coordinate to the world coordinate
    rays_d = directions @ c2w[:, :, :3].transpose(1, 2)  # (n, H*W, 3)
    rays_d = rays_d / torch.norm(rays_d, dim=-1, keepdim=True)
    # The origin of all rays is the camera origin in world coordinate
    rays_o = c2w[:, :, 3].unsqueeze(1).expand(rays_d.shape)  # (n, H*W, 3)

    return _get_rays_inner(rays_o, rays_d, near, far, ray_altitude_range)
//  [python]
def get_rays_batch(directions: torch.Tensor, c2w: torch.Tensor, near: float, far: float,
                   ray_altitude_range: List[float]) -> torch.Tensor:
    # Rotate ray directions from camera coordinate to the world coordinate
    rays_d = directions @ c2w[:, :, :3].transpose(1, 2)  # (n, H*W, 3)
    rays_d = rays_d / torch.linalg.norm(rays_d, dim=-1, keepdim=True)
    # The origin of all rays is the camera origin in world coordinate
    rays_o = c2w[:, :, 3].unsqueeze(1).expand(rays_d.shape)  # (n, H*W, 3)

    return _get_rays_inner(rays_o, rays_d, near, far, ray_altitude_range)
//  [python]
def dist(self, x: torch.Tensor, y: torch.Tensor, keepdim) -> torch.Tensor:
        """
        Computes the affine invariant Riemannian metric (AIM)
        """
        inv_sqrt_x = functionals.sym_invsqrtm.apply(x)
        return torch.norm(
            functionals.sym_logm.apply(inv_sqrt_x @ y @ inv_sqrt_x),
            dim=[-1, -2],
            keepdim=keepdim,
        )
//  [python]
def dist(self, x: torch.Tensor, y: torch.Tensor, keepdim) -> torch.Tensor:
        """
        Computes the affine invariant Riemannian metric (AIM)
        """
        inv_sqrt_x = functionals.sym_invsqrtm.apply(x)
        return torch.linalg.norm(
            functionals.sym_logm.apply(inv_sqrt_x @ y @ inv_sqrt_x),
            dim=[-1, -2],
            keepdim=keepdim,
        )
//  [python]
def mpjpe(predicted, target):
    assert predicted.shape == target.shape   
    return torch.mean(torch.norm(predicted - target, dim=len(target.shape)-1))
//  [python]
def mpjpe(predicted, target):
    assert predicted.shape == target.shape   
    return torch.mean(torch.linalg.norm(predicted - target, dim=len(target.shape)-1))
//  [python]
def greedy_cost_assignment(pos1: torch.Tensor, pos2: torch.Tensor) -> torch.Tensor:
    # Compute pairwise distances between positions
    dist = torch.norm(pos1[:, None] - pos2, dim=2)

    # Create a tensor to store the assignments
    assignments = torch.zeros(dist.size(0), 3, dtype=torch.int64)

    # Create a mask to keep track of assigned positions
    mask = torch.ones_like(dist, dtype=torch.bool)

    # Counter for keeping track of the number of assignments
    counter = 0

    # While there are still unassigned positions
    while mask.sum() > 0:
        # Find the minimum distance
        min_val, min_idx = dist[mask].min(dim=0)

        # Get the indices of the two dimensions
        idx = int(min_idx.item())
        ret = mask.nonzero()[idx, :]
        i = ret[0]
        j = ret[1]

        # Add the assignment to the tensor
        assignments[counter, 0] = i
        assignments[counter, 1] = j
        assignments[counter, 2] = min_val

        # Increase the counter
        counter += 1

        # Remove the assigned positions from the distance matrix and the mask
        mask[i, :] = 0
        mask[:, j] = 0

    return assignments[:counter]
//  [python]
def greedy_cost_assignment(pos1: torch.Tensor, pos2: torch.Tensor) -> torch.Tensor:
    # Compute pairwise distances between positions
    dist = torch.linalg.norm(pos1[:, None] - pos2, dim=2)

    # Create a tensor to store the assignments
    assignments = torch.zeros(dist.size(0), 3, dtype=torch.int64)

    # Create a mask to keep track of assigned positions
    mask = torch.ones_like(dist, dtype=torch.bool)

    # Counter for keeping track of the number of assignments
    counter = 0

    # While there are still unassigned positions
    while mask.sum() > 0:
        # Find the minimum distance
        min_val, min_idx = dist[mask].min(dim=0)

        # Get the indices of the two dimensions
        idx = int(min_idx.item())
        ret = mask.nonzero()[idx, :]
        i = ret[0]
        j = ret[1]

        # Add the assignment to the tensor
        assignments[counter, 0] = i
        assignments[counter, 1] = j
        assignments[counter, 2] = min_val

        # Increase the counter
        counter += 1

        # Remove the assigned positions from the distance matrix and the mask
        mask[i, :] = 0
        mask[:, j] = 0

    return assignments[:counter]
//  [python]
def update(self, minibatches, unlabeled=None):
        all_x = torch.cat([x for x, y in minibatches])
        all_y = torch.cat([y for x, y in minibatches])
        all_features = self.featurizer(all_x)
        if self._what_is_trainable in ["cla"]:
            all_features = all_features.detach()
        all_features = self.modify_features(all_features)
        objective = F.cross_entropy(self.classifier(all_features), all_y)

        l1_reg = torch.tensor(0., requires_grad=True)

        for name, param in self.classifier.named_parameters():
            if 'weight' in name:
                l1_reg = l1_reg + torch.norm(param, 1)

        loss = objective + self.hparams["l1_reg"] * l1_reg

        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

        return {'loss': loss.item(), "objective": objective.item(), "l1_reg": l1_reg.item()}
//  [python]
def update(self, minibatches, unlabeled=None):
        all_x = torch.cat([x for x, y in minibatches])
        all_y = torch.cat([y for x, y in minibatches])
        all_features = self.featurizer(all_x)
        if self._what_is_trainable in ["cla"]:
            all_features = all_features.detach()
        all_features = self.modify_features(all_features)
        objective = F.cross_entropy(self.classifier(all_features), all_y)

        l1_reg = torch.tensor(0., requires_grad=True)

        for name, param in self.classifier.named_parameters():
            if 'weight' in name:
                l1_reg = l1_reg + torch.linalg.norm(param, 1)

        loss = objective + self.hparams["l1_reg"] * l1_reg

        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

        return {'loss': loss.item(), "objective": objective.item(), "l1_reg": l1_reg.item()}
//  [python]
def conjugate_gradient_block(A, B, x0=None, tol=1e-2, maxit=None, eps=1e-6):
    """
    Solve the linear system A X = B using conjugate gradient, where
     - A is an abstract linear operator implementing an n*n matrix
     - B is a matrix right hand size of size n*s
     - X0 is an initial guess of size n*s

    Essentially, this runs #s classical conjugate gradient algorithms in 'parallel',
    and terminates when the worst of the #s residuals is below tol.
    """
    X = x0
    R = B - A(X)
    P = R
    Rs_old = torch.norm(R, 2., dim=0) ** 2.  
    tol_scale = torch.mean(Rs_old) 
    
    k = 0
    while True:
        k += 1
        AP = A(P)
        alpha = Rs_old / (torch.sum(P * AP, dim=0) + eps)
        X += P * alpha

        if k == maxit:
            break
        
        R -= AP * alpha

        Rs_new = torch.norm(R, 2., dim=0) ** 2.
        res = torch.max(Rs_new)
        if res < (tol ** 2.) * tol_scale:
            break

        P = R + P * (Rs_new / (Rs_old + eps))
        Rs_old = Rs_new

    return X, k
//  [python]
def conjugate_gradient_block(A, B, x0=None, tol=1e-2, maxit=None, eps=1e-6):
    """
    Solve the linear system A X = B using conjugate gradient, where
     - A is an abstract linear operator implementing an n*n matrix
     - B is a matrix right hand size of size n*s
     - X0 is an initial guess of size n*s

    Essentially, this runs #s classical conjugate gradient algorithms in 'parallel',
    and terminates when the worst of the #s residuals is below tol.
    """
    X = x0
    R = B - A(X)
    P = R
    Rs_old = torch.linalg.norm(R, 2., dim=0) ** 2.  
    tol_scale = torch.mean(Rs_old) 
    
    k = 0
    while True:
        k += 1
        AP = A(P)
        alpha = Rs_old / (torch.sum(P * AP, dim=0) + eps)
        X += P * alpha

        if k == maxit:
            break
        
        R -= AP * alpha

        Rs_new = torch.linalg.norm(R, 2., dim=0) ** 2.
        res = torch.max(Rs_new)
        if res < (tol ** 2.) * tol_scale:
            break

        P = R + P * (Rs_new / (Rs_old + eps))
        Rs_old = Rs_new

    return X, k
//  [python]
def __call__(self, X):

        N = X.size(0)
        dist = torch.norm(X.unsqueeze(0) - X.unsqueeze(1), dim=-1, p=2)  # [N, N]
        n_neighbors = min(self.k, N)
        kth_dist = dist.topk(k=n_neighbors, dim=-1, largest=False).values[:, -1]  # compute k^th distance for each point, [N, knn + 1]
        sigma = kth_dist.mean()
        rbf = torch.exp(- dist ** 2 / (2 * sigma ** 2))
        # mask = torch.eye(X.size(0)).to(X.device)
        # rbf = rbf * (1 - mask)
        return rbf
//  [python]
def __call__(self, X):

        N = X.size(0)
        dist = torch.linalg.norm(X.unsqueeze(0) - X.unsqueeze(1), dim=-1, ord=2)  # [N, N]
        n_neighbors = min(self.k, N)
        kth_dist = dist.topk(k=n_neighbors, dim=-1, largest=False).values[:, -1]  # compute k^th distance for each point, [N, knn + 1]
        sigma = kth_dist.mean()
        rbf = torch.exp(- dist ** 2 / (2 * sigma ** 2))
        # mask = torch.eye(X.size(0)).to(X.device)
        # rbf = rbf * (1 - mask)
        return rbf
//  [python]
def check_if_mirror_extended(self, lrs):
        """Check whether the input is a mirror-extended sequence.

        If mirror-extended, the i-th (i=0, ..., t-1) frame is equal to the
        (t-1-i)-th frame.

        Args:
            lrs (tensor): Input LR images with shape (n, t, c, h, w)
        """

        self.is_mirror_extended = False
        if lrs.size(1) % 2 == 0:
            lrs_1, lrs_2 = torch.chunk(lrs, 2, dim=1)
            if torch.norm(lrs_1 - lrs_2.flip(1)) == 0:
                self.is_mirror_extended = True
//  [python]
def check_if_mirror_extended(self, lrs):
        """Check whether the input is a mirror-extended sequence.

        If mirror-extended, the i-th (i=0, ..., t-1) frame is equal to the
        (t-1-i)-th frame.

        Args:
            lrs (tensor): Input LR images with shape (n, t, c, h, w)
        """

        self.is_mirror_extended = False
        if lrs.size(1) % 2 == 0:
            lrs_1, lrs_2 = torch.chunk(lrs, 2, dim=1)
            if torch.linalg.norm(lrs_1 - lrs_2.flip(1)) == 0:
                self.is_mirror_extended = True
//  [python]
def compute_matrix_root_inverse_residuals(
    A: Tensor,
    X_hat: Tensor,
    root: int,
    epsilon: float,
    exponent_multiplier: float,
) -> Tuple[Tensor, Tensor]:
    """Compute residual of matrix root inverse for debugging purposes.

        relative error    = ||X - X_hat||_inf / ||X||_inf
        relative residual = ||A X^r - I||_inf

    Args:
        A (Tensor): Matrix of interest.
        X (Tensor): Computed matrix root inverse.
        root (int): Root of interest.
        epsilon (float): Adds epsilon * I to matrix.
        exponent_multiplier (float): Exponent multiplier to be multiplied to the numerator of the inverse root.

    Returns:
        absolute_error (Tensor): absolute error of matrix root inverse
        relative_error (Tensor): relative error of matrix root inverse
        residual (Tensor): residual of matrix root inverse

    """

    # check shape of matrix
    if len(A.shape) != 2:
        raise ValueError("Matrix is not 2-dimensional!")
    elif A.shape[0] != A.shape[1]:
        raise ValueError("Matrix is not square!")
    elif A.shape != X_hat.shape:
        raise ValueError("Matrix shapes do not match!")

    # compute error by comparing against double precision
    X = matrix_inverse_root(
        A.double(),
        root,
        epsilon=epsilon,
        exponent_multiplier=exponent_multiplier,
    )
    relative_error = torch.dist(X, X_hat, p=torch.inf) / torch.norm(
        X, p=torch.inf
    )

    # compute residual
    if exponent_multiplier == 1.0:
        X_invr = torch.linalg.matrix_power(X_hat.double(), n=-root)
    else:
        X_invr, _, _ = _matrix_root_eigen(
            X_hat.double(),
            root=1,
            epsilon=0.0,
            inverse=True,
            make_positive_semidefinite=True,
            exponent_multiplier=root / exponent_multiplier,
        )

    A_reg = A.double() + epsilon * torch.eye(
        A.shape[0], dtype=torch.float64, device=A.device
    )
    relative_residual = torch.dist(X_invr, A_reg, p=torch.inf) / torch.norm(
        A_reg, p=torch.inf
    )

    return relative_error, relative_residual
//  [python]
def compute_matrix_root_inverse_residuals(
    A: Tensor,
    X_hat: Tensor,
    root: int,
    epsilon: float,
    exponent_multiplier: float,
) -> Tuple[Tensor, Tensor]:
    """Compute residual of matrix root inverse for debugging purposes.

        relative error    = ||X - X_hat||_inf / ||X||_inf
        relative residual = ||A X^r - I||_inf

    Args:
        A (Tensor): Matrix of interest.
        X (Tensor): Computed matrix root inverse.
        root (int): Root of interest.
        epsilon (float): Adds epsilon * I to matrix.
        exponent_multiplier (float): Exponent multiplier to be multiplied to the numerator of the inverse root.

    Returns:
        absolute_error (Tensor): absolute error of matrix root inverse
        relative_error (Tensor): relative error of matrix root inverse
        residual (Tensor): residual of matrix root inverse

    """

    # check shape of matrix
    if len(A.shape) != 2:
        raise ValueError("Matrix is not 2-dimensional!")
    elif A.shape[0] != A.shape[1]:
        raise ValueError("Matrix is not square!")
    elif A.shape != X_hat.shape:
        raise ValueError("Matrix shapes do not match!")

    # compute error by comparing against double precision
    X = matrix_inverse_root(
        A.double(),
        root,
        epsilon=epsilon,
        exponent_multiplier=exponent_multiplier,
    )
    relative_error = torch.dist(X, X_hat, p=torch.inf) / torch.linalg.norm(
        X, p=torch.inf
    )

    # compute residual
    if exponent_multiplier == 1.0:
        X_invr = torch.linalg.matrix_power(X_hat.double(), n=-root)
    else:
        X_invr, _, _ = _matrix_root_eigen(
            X_hat.double(),
            root=1,
            epsilon=0.0,
            inverse=True,
            make_positive_semidefinite=True,
            exponent_multiplier=root / exponent_multiplier,
        )

    A_reg = A.double() + epsilon * torch.eye(
        A.shape[0], dtype=torch.float64, device=A.device
    )
    relative_residual = torch.dist(X_invr, A_reg, p=torch.inf) / torch.linalg.norm(
        A_reg, p=torch.inf
    )

    return relative_error, relative_residual
//  [python]
def gcsim(Ua, Sa, Ub, Sb):
    """
    Generalised cosine similarity
    :param Ua: First singular vectors
    :param Sa: First singular values
    :param Ub: Second singular vectors
    :param Sb: Second singular vectors
    :return Generalised cosine similarity
    """
    Sa = torch.diag(Sa)
    Sb = torch.diag(Sb)
    Va = torch.sqrt(Sa).mm(Ua.t())
    Vb = Ub.mm(torch.sqrt(Sb))
    Vab = Va.mm(Vb)
    num = torch.pow(torch.norm(Vab), 2)
    den = torch.norm(torch.diag(Sa), p=2) * torch.norm(torch.diag(Sb), p=2)
    return num / den
//  [python]
def gcsim(Ua, Sa, Ub, Sb):
    """
    Generalised cosine similarity
    :param Ua: First singular vectors
    :param Sa: First singular values
    :param Ub: Second singular vectors
    :param Sb: Second singular vectors
    :return Generalised cosine similarity
    """
    Sa = torch.diag(Sa)
    Sb = torch.diag(Sb)
    Va = torch.sqrt(Sa).mm(Ua.t())
    Vb = Ub.mm(torch.sqrt(Sb))
    Vab = Va.mm(Vb)
    num = torch.pow(torch.linalg.norm(Vab), 2)
    den = torch.linalg.norm(torch.diag(Sa), p=2) * torch.linalg.norm(torch.diag(Sb), p=2)
    return num / den
//  [python]
def grad_norm(self):
        """Norm of weight gradients. Useful for monitoring gradient flow."""
        if self.weight_norm_enabled is False:
            if self.weight.grad is None:
                return 0
            else:
                return torch.norm(
                    self.weight.grad
                ).item() / torch.numel(self.weight.grad)
        else:
            if self.weight_g.grad is None:
                return 0
            else:
                return torch.norm(
                    self.weight_g.grad
                ).item() / torch.numel(self.weight_g.grad)
//  [python]
def grad_norm(self):
        """Norm of weight gradients. Useful for monitoring gradient flow."""
        if self.weight_norm_enabled is False:
            if self.weight.grad is None:
                return 0
            else:
                return torch.linalg.norm(
                    self.weight.grad
                ).item() / torch.numel(self.weight.grad)
        else:
            if self.weight_g.grad is None:
                return 0
            else:
                return torch.linalg.norm(
                    self.weight_g.grad
                ).item() / torch.numel(self.weight_g.grad)
//  [python]
def batch_rod2quat(rot_vecs):
    batch_size = rot_vecs.shape[0]

    angle = torch.norm(rot_vecs + 1e-16, dim=1, keepdim=True)
    rot_dir = rot_vecs / angle

    cos = torch.cos(angle / 2)
    sin = torch.sin(angle / 2)

    # Bx1 arrays
    rx, ry, rz = torch.split(rot_dir, 1, dim=1)

    qx = rx * sin
    qy = ry * sin
    qz = rz * sin
    qw = cos - 1.0

    return torch.cat([qx, qy, qz, qw], dim=1)
//  [python]
def batch_rod2quat(rot_vecs):
    batch_size = rot_vecs.shape[0]

    angle = torch.linalg.norm(rot_vecs + 1e-16, dim=1, keepdim=True)
    rot_dir = rot_vecs / angle

    cos = torch.cos(angle / 2)
    sin = torch.sin(angle / 2)

    # Bx1 arrays
    rx, ry, rz = torch.split(rot_dir, 1, dim=1)

    qx = rx * sin
    qy = ry * sin
    qz = rz * sin
    qw = cos - 1.0

    return torch.cat([qx, qy, qz, qw], dim=1)
//  [python]
def score_batch(self, model, data, features_dataset=False):
        data = lib.to_device(data)
        if features_dataset:
            z = data['avg'].float()
        else:
            z = lib.get_features(model, data, layer_names=[self.layer_name])['avg'][-1]

        z = z - self.u
        logits = model.fc.forward(z)
        vlogits = torch.norm(z.matmul(self.R), dim=-1)
        energy = - torch.logsumexp(logits, dim=1)
        return self.alpha * vlogits + energy
//  [python]
def score_batch(self, model, data, features_dataset=False):
        data = lib.to_device(data)
        if features_dataset:
            z = data['avg'].float()
        else:
            z = lib.get_features(model, data, layer_names=[self.layer_name])['avg'][-1]

        z = z - self.u
        logits = model.fc.forward(z)
        vlogits = torch.linalg.norm(z.matmul(self.R), dim=-1)
        energy = - torch.logsumexp(logits, dim=1)
        return self.alpha * vlogits + energy
//  [python]
def forward(self, s_inputs, s_outputs, t_input):

        num_src_domains = len(s_inputs)
        train_losses, discs = [], []
        Phi_t = torch.cat([t_input, torch.ones_like(t_input)], dim=1)
        M_t = torch.mm(Phi_t.transpose(0, 1), Phi_t)

        for i in range(num_src_domains):

            prediction = self.regress_net(s_inputs[i])
            train_loss = torch.mean((prediction - s_outputs[i]) ** 2)
            train_losses.append(train_loss)

            Phi_s = torch.cat([s_inputs[i], torch.ones_like(s_inputs[i])], dim=1)
            M_s = torch.mm(Phi_s.transpose(0, 1), Phi_s)
            disc = torch.max(torch.norm(torch.eig(M_t - M_s)[0], dim=1))
            discs.append(disc)

        train_losses = torch.stack(train_losses)
        discs = torch.stack(discs)

        g = self.gamma * (train_losses + self.mu * discs)
        alpha = self.proj(g)
        loss = torch.dot(g, alpha) + torch.norm(alpha)
        alpha = alpha.detach().cpu().numpy()

        return loss, alpha
//  [python]
def forward(self, s_inputs, s_outputs, t_input):

        num_src_domains = len(s_inputs)
        train_losses, discs = [], []
        Phi_t = torch.cat([t_input, torch.ones_like(t_input)], dim=1)
        M_t = torch.mm(Phi_t.transpose(0, 1), Phi_t)

        for i in range(num_src_domains):

            prediction = self.regress_net(s_inputs[i])
            train_loss = torch.mean((prediction - s_outputs[i]) ** 2)
            train_losses.append(train_loss)

            Phi_s = torch.cat([s_inputs[i], torch.ones_like(s_inputs[i])], dim=1)
            M_s = torch.mm(Phi_s.transpose(0, 1), Phi_s)
            disc = torch.max(torch.linalg.norm(torch.eig(M_t - M_s)[0], dim=1))
            discs.append(disc)

        train_losses = torch.stack(train_losses)
        discs = torch.stack(discs)

        g = self.gamma * (train_losses + self.mu * discs)
        alpha = self.proj(g)
        loss = torch.dot(g, alpha) + torch.linalg.norm(alpha)
        alpha = alpha.detach().cpu().numpy()

        return loss, alpha
//  [python]
def get_compatibility(src_keypts, tgt_keypts, with_dist_corr_compatibility, dist_sigma=None):
    # [n,3], [n,3]
    #################################
    # Step1: corr_compatibily
    #################################
    src_dist = torch.norm((src_keypts.unsqueeze(-2) - src_keypts.unsqueeze(-3)), dim=-1) # [n1, n1]
    tgt_dist = torch.norm((tgt_keypts.unsqueeze(-2) - tgt_keypts.unsqueeze(-3)), dim=-1) # [n1, n1]
    
    corr_compatibility = torch.min(torch.stack((src_dist/(tgt_dist+1e-12),tgt_dist/(src_dist+1e-12)), dim=-1), dim=-1).values + torch.eye(src_dist.shape[-1]).to(src_dist)
    corr_compatibility = torch.clamp((corr_compatibility / 0.9)**2, min=0)
    
    #################################
    # Step2: dist_compatibily
    #################################
    if with_dist_corr_compatibility:
        dist_corr_compatibility = (src_dist + tgt_dist)/2
        dist_corr_compatibility = torch.clamp(1.0 - dist_corr_compatibility ** 2 / dist_sigma ** 2, min=0) # [n1, n1]
        corr_compatibility *= dist_corr_compatibility
        corr_compatibility /= corr_compatibility.max(dim=-1, keepdim=True).values.max(dim=-1, keepdim=True).values
    return corr_compatibility