jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def forward(self, data):
        """
        Input:
            - corr_pos:   [bs, num_corr, 6]
            - src_keypts: [bs, num_corr, 3]
            - tgt_keypts: [bs, num_corr, 3]
            - testing:    flag for test phase, if False will not calculate M and post-refinement.
        Output: (dict)
            - final_trans:   [bs, 4, 4], the predicted transformation matrix. 
            - final_labels:  [bs, num_corr], the predicted inlier/outlier label (0,1), for classification loss calculation.
            - M:             [bs, num_corr, num_corr], feature similarity matrix, for SM loss calculation.
            - seed_trans:    [bs, num_seeds, 4, 4],  the predicted transformation matrix associated with each seeding point, deprecated.
            - corr_features: [bs, num_corr, num_channels], the feature for each correspondence, for circle loss calculation, deprecated.
            - confidence:    [bs], confidence of returned results, for safe guard, deprecated.
        """
        corr_pos, src_keypts, tgt_keypts = data['corr_pos'], data['src_keypts'], data['tgt_keypts']
        bs, num_corr = corr_pos.shape[0], corr_pos.shape[1]
        testing = 'testing' in data.keys()
        with torch.no_grad():
            src_dist = torch.linalg.norm(src_keypts[:, :, None, :] - src_keypts[:, None, :, :], dim=-1)
            corr_compatibility = src_dist - torch.linalg.norm(tgt_keypts[:, :, None, :] - tgt_keypts[:, None, :, :], dim=-1)
            corr_compatibility = torch.clamp(1.0 - corr_compatibility ** 2 / self.sigma_spat ** 2, min=0)
        corr_features = self.encoder(corr_pos.permute(0, 2, 1), corr_compatibility).permute(0, 2, 1)
        normed_corr_features = F.normalize(corr_features, p=2, dim=-1)
        if not testing:
            M = torch.matmul(normed_corr_features, normed_corr_features.permute(0, 2, 1))
            M = torch.clamp(1 - (1 - M) / self.sigma ** 2, min=0, max=1)
            M[:, torch.arange(M.shape[1]), torch.arange(M.shape[1])] = 0
        else:
            M = None
        confidence = self.classification(corr_features.permute(0, 2, 1)).squeeze(1)
        if testing:
            seeds = self.pick_seeds(src_dist, confidence, R=self.nms_radius, max_num=int(num_corr * self.ratio))
        else:
            seeds = torch.argsort(confidence, dim=1, descending=True)[:, 0:int(num_corr * self.ratio)]
        seed_trans, seed_fitness, final_trans, final_labels = self.cal_seed_trans(seeds, normed_corr_features, src_keypts, tgt_keypts)
        if testing:
            final_trans = self.post_refinement(final_trans, src_keypts, tgt_keypts)
        if not testing:
            final_labels = confidence
        res = {'final_trans': final_trans, 'final_labels': final_labels, 'M': M}
        return res
//  [python]
def regress_meta(self, replay_loader, step):
        obs, reward = [], []
        batch_size = 0
        while batch_size < self.lstsq_batch_size:
            batch = replay_loader.sample(self.cfg.batch_size)
            batch_obs, _, batch_reward, *_ = utils.to_torch(batch, self.device)
            obs.append(batch_obs)
            reward.append(batch_reward)
            batch_size += batch_obs.size(0)
        obs, reward = torch.cat(obs, 0), torch.cat(reward, 0)

        obs = self.aug_and_encode(obs)
        rep = self.aps(obs)
        task = torch.linalg.lstsq(reward, rep)[0][:rep.size(1), :][0]
        task = task / torch.norm(task)
        task = task.cpu().numpy()
        meta = OrderedDict()
        meta['task'] = task

        # save for evaluation
        self.solved_meta = meta
        return meta
//  [python]
def regress_meta(self, replay_loader, step):
        obs, reward = [], []
        batch_size = 0
        while batch_size < self.lstsq_batch_size:
            batch = replay_loader.sample(self.cfg.batch_size)
            batch_obs, _, batch_reward, *_ = utils.to_torch(batch, self.device)
            obs.append(batch_obs)
            reward.append(batch_reward)
            batch_size += batch_obs.size(0)
        obs, reward = torch.cat(obs, 0), torch.cat(reward, 0)

        obs = self.aug_and_encode(obs)
        rep = self.aps(obs)
        task = torch.linalg.lstsq(reward, rep)[0][:rep.size(1), :][0]
        task = task / torch.linalg.norm(task)
        task = task.cpu().numpy()
        meta = OrderedDict()
        meta['task'] = task

        # save for evaluation
        self.solved_meta = meta
        return meta
//  [python]
def infer_meta_from_obs_and_rewards(self, obs: torch.Tensor, reward: torch.Tensor) -> MetaDict:
        print('max reward: ', reward.max().cpu().item())
        print('99 percentile: ', torch.quantile(reward, 0.99).cpu().item())
        print('median reward: ', reward.median().cpu().item())
        print('min reward: ', reward.min().cpu().item())
        print('mean reward: ', reward.mean().cpu().item())
        print('num reward: ', reward.shape[0])

        rep = self.aps(obs)
        # task = torch.linalg.lstsq(reward, rep)[0][:rep.size(1), :][0]
        task = torch.linalg.lstsq(rep, reward)[0].squeeze()
        task = task / torch.norm(task)
        task = task.cpu().numpy()
        meta = OrderedDict()
        meta['task'] = task

        # self.solved_meta = meta
        return meta
//  [python]
def infer_meta_from_obs_and_rewards(self, obs: torch.Tensor, reward: torch.Tensor) -> MetaDict:
        print('max reward: ', reward.max().cpu().item())
        print('99 percentile: ', torch.quantile(reward, 0.99).cpu().item())
        print('median reward: ', reward.median().cpu().item())
        print('min reward: ', reward.min().cpu().item())
        print('mean reward: ', reward.mean().cpu().item())
        print('num reward: ', reward.shape[0])

        rep = self.aps(obs)
        # task = torch.linalg.lstsq(reward, rep)[0][:rep.size(1), :][0]
        task = torch.linalg.lstsq(rep, reward)[0].squeeze()
        task = task / torch.linalg.norm(task)
        task = task.cpu().numpy()
        meta = OrderedDict()
        meta['task'] = task

        # self.solved_meta = meta
        return meta
//  [python]
def forward(self, x_mag, y_mag):
        """Calculate forward propagation.

        Args:
            x_mag (Tensor): Magnitude spectrogram of predicted signal (B, #frames, #freq_bins).
            y_mag (Tensor): Magnitude spectrogram of groundtruth signal (B, #frames, #freq_bins).

        Returns:
            Tensor: Spectral convergence loss value.

        """
        return torch.norm(y_mag - x_mag, p='fro') / torch.norm(y_mag, p='fro')
//  [python]
def forward(self, x_mag, y_mag):
        """Calculate forward propagation.

        Args:
            x_mag (Tensor): Magnitude spectrogram of predicted signal (B, #frames, #freq_bins).
            y_mag (Tensor): Magnitude spectrogram of groundtruth signal (B, #frames, #freq_bins).

        Returns:
            Tensor: Spectral convergence loss value.

        """
        return torch.linalg.norm(y_mag - x_mag, ord='fro') / torch.linalg.norm(y_mag, ord='fro')
//  [python]
def get_norms(X, edge):
    X_norm = float(torch.norm(X))
    iX_norm = X_norm ** -1
    A_norm = float(torch.norm(A(X, edge))) * iX_norm
    B_norm = float(torch.norm(B(X, edge))) * iX_norm
    return A_norm, B_norm
//  [python]
def get_norms(X, edge):
    X_norm = float(torch.linalg.norm(X))
    iX_norm = X_norm ** -1
    A_norm = float(torch.linalg.norm(A(X, edge))) * iX_norm
    B_norm = float(torch.linalg.norm(B(X, edge))) * iX_norm
    return A_norm, B_norm
//  [python]
def rerr(U, E, edge, A_norm, B_norm):
    # check_shape(X, E)
    R = A(U, edge) - B(U, edge) * E
    rerr = torch.norm(R, 2, (0, )) * (torch.norm(U, 2, (0, )) * (A_norm + E * B_norm)) ** -1
    return rerr
//  [python]
def rerr(U, E, edge, A_norm, B_norm):
    # check_shape(X, E)
    R = A(U, edge) - B(U, edge) * E
    rerr = torch.linalg.norm(R, 2, (0, )) * (torch.linalg.norm(U, 2, (0, )) * (A_norm + E * B_norm)) ** -1
    return rerr
//  [python]
def compute_mpjpe(preds,
                  target,
                  valid_mask=None,
                  pck_joints=None,
                  sample_wise=True):
    """
    Mean per-joint position error (i.e. mean Euclidean distance)
    often referred to as "Protocol #1" in many papers.
    """
    assert preds.shape == target.shape, print(preds.shape,
                                              target.shape)  # BxJx3
    mpjpe = torch.norm(preds - target, p=2, dim=-1)  # BxJ

    if pck_joints is None:
        if sample_wise:
            mpjpe_seq = ((mpjpe * valid_mask.float()).sum(-1) /
                         valid_mask.float().sum(-1)
                         if valid_mask is not None else mpjpe.mean(-1))
        else:
            mpjpe_seq = mpjpe[valid_mask] if valid_mask is not None else mpjpe
        return mpjpe_seq
    else:
        mpjpe_pck_seq = mpjpe[:, pck_joints]
        return mpjpe_pck_seq
//  [python]
def compute_mpjpe(preds,
                  target,
                  valid_mask=None,
                  pck_joints=None,
                  sample_wise=True):
    """
    Mean per-joint position error (i.e. mean Euclidean distance)
    often referred to as "Protocol #1" in many papers.
    """
    assert preds.shape == target.shape, print(preds.shape,
                                              target.shape)  # BxJx3
    mpjpe = torch.linalg.norm(preds - target, ord=2, dim=-1)  # BxJ

    if pck_joints is None:
        if sample_wise:
            mpjpe_seq = ((mpjpe * valid_mask.float()).sum(-1) /
                         valid_mask.float().sum(-1)
                         if valid_mask is not None else mpjpe.mean(-1))
        else:
            mpjpe_seq = mpjpe[valid_mask] if valid_mask is not None else mpjpe
        return mpjpe_seq
    else:
        mpjpe_pck_seq = mpjpe[:, pck_joints]
        return mpjpe_pck_seq
//  [python]
def cal_confidence(self, M, leading_eig, method='eig_value'):
        """
        Calculate the confidence of the spectral matching solution based on spectral analysis.
        Input:
            - M:          [bs, num_corr, num_corr] the compatibility matrix
            - leading_eig [bs, num_corr]           the leading eigenvector of matrix M
        Output:
            - confidence
        """
        if method == 'eig_value':
            # max eigenvalue as the confidence (Rayleigh quotient)
            max_eig_value = (leading_eig[:, None, :] @ M @ leading_eig[:, :, None]) / (
                        leading_eig[:, None, :] @ leading_eig[:, :, None])
            confidence = max_eig_value.squeeze(-1)
            return confidence
        elif method == 'eig_value_ratio':
            # max eigenvalue / second max eigenvalue as the confidence
            max_eig_value = (leading_eig[:, None, :] @ M @ leading_eig[:, :, None]) / (
                        leading_eig[:, None, :] @ leading_eig[:, :, None])
            # compute the second largest eigen-value
            B = M - max_eig_value * leading_eig[:, :, None] @ leading_eig[:, None, :]
            solution = torch.ones_like(B[:, :, 0:1])
            for i in range(self.num_iterations):
                solution = torch.bmm(B, solution)
                solution = solution / (torch.norm(solution, dim=1, keepdim=True) + 1e-6)
            solution = solution.squeeze(-1)
            second_eig = solution
            second_eig_value = (second_eig[:, None, :] @ B @ second_eig[:, :, None]) / (
                        second_eig[:, None, :] @ second_eig[:, :, None])
            confidence = max_eig_value / second_eig_value
            return confidence
        elif method == 'xMx':
            # max xMx as the confidence (x is the binary solution)
            # rank = torch.argsort(leading_eig, dim=1, descending=True)[:, 0:int(M.shape[1]*self.ratio)]
            # binary_sol = torch.zeros_like(leading_eig)
            # binary_sol[0, rank[0]] = 1
            confidence = leading_eig[:, None, :] @ M @ leading_eig[:, :, None]
            confidence = confidence.squeeze(-1) / M.shape[1]
            return confidence
//  [python]
def cal_confidence(self, M, leading_eig, method='eig_value'):
        """
        Calculate the confidence of the spectral matching solution based on spectral analysis.
        Input:
            - M:          [bs, num_corr, num_corr] the compatibility matrix
            - leading_eig [bs, num_corr]           the leading eigenvector of matrix M
        Output:
            - confidence
        """
        if method == 'eig_value':
            # max eigenvalue as the confidence (Rayleigh quotient)
            max_eig_value = (leading_eig[:, None, :] @ M @ leading_eig[:, :, None]) / (
                        leading_eig[:, None, :] @ leading_eig[:, :, None])
            confidence = max_eig_value.squeeze(-1)
            return confidence
        elif method == 'eig_value_ratio':
            # max eigenvalue / second max eigenvalue as the confidence
            max_eig_value = (leading_eig[:, None, :] @ M @ leading_eig[:, :, None]) / (
                        leading_eig[:, None, :] @ leading_eig[:, :, None])
            # compute the second largest eigen-value
            B = M - max_eig_value * leading_eig[:, :, None] @ leading_eig[:, None, :]
            solution = torch.ones_like(B[:, :, 0:1])
            for i in range(self.num_iterations):
                solution = torch.bmm(B, solution)
                solution = solution / (torch.linalg.norm(solution, dim=1, keepdim=True) + 1e-6)
            solution = solution.squeeze(-1)
            second_eig = solution
            second_eig_value = (second_eig[:, None, :] @ B @ second_eig[:, :, None]) / (
                        second_eig[:, None, :] @ second_eig[:, :, None])
            confidence = max_eig_value / second_eig_value
            return confidence
        elif method == 'xMx':
            # max xMx as the confidence (x is the binary solution)
            # rank = torch.argsort(leading_eig, dim=1, descending=True)[:, 0:int(M.shape[1]*self.ratio)]
            # binary_sol = torch.zeros_like(leading_eig)
            # binary_sol[0, rank[0]] = 1
            confidence = leading_eig[:, None, :] @ M @ leading_eig[:, :, None]
            confidence = confidence.squeeze(-1) / M.shape[1]
            return confidence
//  [python]
def post_refinement(self, initial_trans, src_keypts, tgt_keypts, it_num, weights=None):
        """
        Perform post refinement using the initial transformation matrix, only adopted during testing.
        Input
            - initial_trans: [bs, 4, 4]
            - src_keypts:    [bs, num_corr, 3]
            - tgt_keypts:    [bs, num_corr, 3]
            - weights:       [bs, num_corr]
        Output:
            - final_trans:   [bs, 4, 4]
        """
        assert initial_trans.shape[0] == 1
        inlier_threshold = 1.2

        # inlier_threshold_list = [self.inlier_threshold] * it_num

        if self.inlier_threshold == 0.10:  # for 3DMatch
            inlier_threshold_list = [0.10] * it_num
        else:  # for KITTI
            inlier_threshold_list = [1.2] * it_num

        previous_inlier_num = 0
        for inlier_threshold in inlier_threshold_list:
            warped_src_keypts = transform(src_keypts, initial_trans)

            L2_dis = torch.norm(warped_src_keypts - tgt_keypts, dim=-1)
            pred_inlier = (L2_dis < inlier_threshold)[0]  # assume bs = 1
            inlier_num = torch.sum(pred_inlier)
            if abs(int(inlier_num - previous_inlier_num)) < 1:
                break
            else:
                previous_inlier_num = inlier_num
            initial_trans = rigid_transform_3d(
                A=src_keypts[:, pred_inlier, :],
                B=tgt_keypts[:, pred_inlier, :],
                ## https://link.springer.com/article/10.1007/s10589-014-9643-2
                # weights=None,
                weights=1 / (1 + (L2_dis / inlier_threshold) ** 2)[:, pred_inlier],
                # weights=((1-L2_dis/inlier_threshold)**2)[:, pred_inlier]
            )
        return initial_trans
//  [python]
def post_refinement(self, initial_trans, src_keypts, tgt_keypts, it_num, weights=None):
        """
        Perform post refinement using the initial transformation matrix, only adopted during testing.
        Input
            - initial_trans: [bs, 4, 4]
            - src_keypts:    [bs, num_corr, 3]
            - tgt_keypts:    [bs, num_corr, 3]
            - weights:       [bs, num_corr]
        Output:
            - final_trans:   [bs, 4, 4]
        """
        assert initial_trans.shape[0] == 1
        inlier_threshold = 1.2

        # inlier_threshold_list = [self.inlier_threshold] * it_num

        if self.inlier_threshold == 0.10:  # for 3DMatch
            inlier_threshold_list = [0.10] * it_num
        else:  # for KITTI
            inlier_threshold_list = [1.2] * it_num

        previous_inlier_num = 0
        for inlier_threshold in inlier_threshold_list:
            warped_src_keypts = transform(src_keypts, initial_trans)

            L2_dis = torch.linalg.norm(warped_src_keypts - tgt_keypts, dim=-1)
            pred_inlier = (L2_dis < inlier_threshold)[0]  # assume bs = 1
            inlier_num = torch.sum(pred_inlier)
            if abs(int(inlier_num - previous_inlier_num)) < 1:
                break
            else:
                previous_inlier_num = inlier_num
            initial_trans = rigid_transform_3d(
                A=src_keypts[:, pred_inlier, :],
                B=tgt_keypts[:, pred_inlier, :],
                ## https://link.springer.com/article/10.1007/s10589-014-9643-2
                # weights=None,
                weights=1 / (1 + (L2_dis / inlier_threshold) ** 2)[:, pred_inlier],
                # weights=((1-L2_dis/inlier_threshold)**2)[:, pred_inlier]
            )
        return initial_trans
//  [python]
def forward(self, x):
        return torch.norm(x, dim=1, p=self.p)
//  [python]
def forward(self, x):
        return torch.linalg.norm(x, dim=1, ord=self.p)
//  [python]
def batch_rodrigues(param):
    #param N x 3
    batch_size = param.shape[0]

    l1norm = torch.norm(param + 1e-8, p = 2, dim = 1)
    angle = torch.unsqueeze(l1norm, -1)
    normalized = torch.div(param, angle)
    angle = angle * 0.5

    v_cos = torch.cos(angle)
    v_sin = torch.sin(angle)

    quat = torch.cat([v_cos, v_sin * normalized], dim = 1)

    return quat2mat(quat)
//  [python]
def batch_rodrigues(param):
    #param N x 3
    batch_size = param.shape[0]

    l1norm = torch.linalg.norm(param + 1e-8, ord = 2, dim = 1)
    angle = torch.unsqueeze(l1norm, -1)
    normalized = torch.div(param, angle)
    angle = angle * 0.5

    v_cos = torch.cos(angle)
    v_sin = torch.sin(angle)

    quat = torch.cat([v_cos, v_sin * normalized], dim = 1)

    return quat2mat(quat)
//  [python]
def abs(self, x, y):
        num_examples = x.size()[0]

        # Assume uniform mesh
        h = 1.0 / (x.size()[1] - 1.0)

        all_norms = (h ** (self.d / self.p)) * torch.norm(
            x.view(num_examples, -1) - y.view(num_examples, -1), self.p, 1
        )

        if self.reduction:
            if self.size_average:
                return torch.mean(all_norms)
            else:
                return torch.sum(all_norms)

        return all_norms
//  [python]
def abs(self, x, y):
        num_examples = x.size()[0]

        # Assume uniform mesh
        h = 1.0 / (x.size()[1] - 1.0)

        all_norms = (h ** (self.d / self.p)) * torch.linalg.norm(
            x.view(num_examples, -1) - y.view(num_examples, -1), ord=self.p, dim=1
        )

        if self.reduction:
            if self.size_average:
                return torch.mean(all_norms)
            else:
                return torch.sum(all_norms)

        return all_norms
//  [python]
def rel(self, x, y):
        num_examples = x.size()[0]

        diff_norms = torch.norm(
            x.reshape(num_examples, -1) - y.reshape(num_examples, -1), self.p, 1
        )
        y_norms = torch.norm(y.reshape(num_examples, -1), self.p, 1)

        if self.reduction:
            if self.size_average:
                return torch.mean(diff_norms / y_norms)
            else:
                return torch.sum(diff_norms / y_norms)

        return diff_norms / y_norms
//  [python]
def rel(self, x, y):
        num_examples = x.size()[0]

        diff_norms = torch.linalg.norm(
            x.reshape(num_examples, -1) - y.reshape(num_examples, -1), ord=self.p, dim=1
        )
        y_norms = torch.linalg.norm(y.reshape(num_examples, -1), ord=self.p, dim=1)

        if self.reduction:
            if self.size_average:
                return torch.mean(diff_norms / y_norms)
            else:
                return torch.sum(diff_norms / y_norms)

        return diff_norms / y_norms
//  [python]
def forward(self, initial_states, controls):
        N, state_dim = initial_states.shape
        assert self.state_dim == state_dim
        
        x = initial_states[:, :3]
        quat = torch.div(initial_states[:, 3:7], torch.norm(initial_states[:, 3:7], dim=1)[:, None]).detach()
        x_dot = initial_states[:, 7:10]
        
        accel = controls[:, :3]
        omega = controls[:, 3:6]
        
        gravity = torch.zeros((N, 3))
        gravity[:, -1] = 9.81
        
        accel_enu = tf.quaternion_apply(quat, accel) - gravity
        accel_enu[:, :] = 0.0
        
#         self.save_data['a_enu'].append(accel_enu[0, :])
        
        quat_dot = compute_quat_dot(quat, omega)
        
#         self.save_data['q_dot'].append(quat_dot[0, :])
        
        predicted_x = x + self.dt*x_dot
        
        _predicted_quat = quat + self.dt*quat_dot
        predicted_quat = torch.div(_predicted_quat, torch.norm(_predicted_quat, dim=1)[:, None])
        
        predicted_x_dot = x_dot + self.dt*accel_enu
        
        predicted_states = torch.cat((predicted_x, predicted_quat, predicted_x_dot), -1)
        
        if self.N_dim > 0:
            N_allsv = initial_states[:, 10:]
            predicted_states = torch.cat((predicted_states, N_allsv), -1)
        
        Q = self.calc_Q()
        
        return predicted_states, Q.expand((N, state_dim, state_dim))
//  [python]
def forward(self, initial_states, controls):
        N, state_dim = initial_states.shape
        assert self.state_dim == state_dim
        
        x = initial_states[:, :3]
        quat = torch.div(initial_states[:, 3:7], torch.linalg.norm(initial_states[:, 3:7], dim=1)[:, None]).detach()
        x_dot = initial_states[:, 7:10]
        
        accel = controls[:, :3]
        omega = controls[:, 3:6]
        
        gravity = torch.zeros((N, 3))
        gravity[:, -1] = 9.81
        
        accel_enu = tf.quaternion_apply(quat, accel) - gravity
        accel_enu[:, :] = 0.0
        
#         self.save_data['a_enu'].append(accel_enu[0, :])
        
        quat_dot = compute_quat_dot(quat, omega)
        
#         self.save_data['q_dot'].append(quat_dot[0, :])
        
        predicted_x = x + self.dt*x_dot
        
        _predicted_quat = quat + self.dt*quat_dot
        predicted_quat = torch.div(_predicted_quat, torch.linalg.norm(_predicted_quat, dim=1)[:, None])
        
        predicted_x_dot = x_dot + self.dt*accel_enu
        
        predicted_states = torch.cat((predicted_x, predicted_quat, predicted_x_dot), -1)
        
        if self.N_dim > 0:
            N_allsv = initial_states[:, 10:]
            predicted_states = torch.cat((predicted_states, N_allsv), -1)
        
        Q = self.calc_Q()
        
        return predicted_states, Q.expand((N, state_dim, state_dim))
//  [python]
def estimate_angles(ratios, outcomes):
    #"ratios" - torch.Tensor, shape=(number of grains, number of neighborhoods), ratios found in each neighborhood
    #"outcomes" - torch.Tensor, shape=(number of grains, number of neighborhoods), which ratio was chosen to flip to
    
    d = ratios.shape[0]
    ij = torch.combinations(torch.arange(d), r=2)
    e = (torch.ones(d)/d)[:,None] #intersecting point when all factors are equal
    e = e/torch.norm(e)
    
    midpoints = []
    angles = []
    for i, j in ij: #for each grain ID pair
    
        #New center
        c = torch.zeros(d)[:,None] 
        c[i] = 0.5
        c[j] = 0.5
        
        #Find max angle from grain IDs i to j
        v = torch.zeros(d)[:,None]; v[i] = 1 #find angle from this vector
        rs = ratios[:,outcomes[i]==1] #for these ratios
        rmi, ami = find_radial_max(rs, c, v) #angle max from i
        # rmi = (v-c)*np.cos(ami/180*np.pi) + (e-c)*np.sin(ami/180*np.pi) + c
        
        #Find max angle from grain j to i
        v = torch.zeros(d)[:,None]; v[j] = 1 #find angle from this vector
        rs = ratios[:,outcomes[j]==1] #for these ratios
        rmj, amj = find_radial_max(rs, c, v) #angle max from j
        # rmj = (v-c)*np.cos(amj/180*np.pi) + (e-c)*np.sin(amj/180*np.pi) + c
        
        #Find a point in the dividing line
        midpoints.append((rmi+rmj)/2)
        
        #Average angles
        angles.append((ami+180-amj)/2)
        
    #Convert to Tensors
    angles = torch.stack(angles)[:,None]
    midpoints = torch.stack(midpoints)
    
    if ratios.shape[0]==3:
        #Replace invalid midpoints
        is_valid = is_valid_3grains(ratios, outcomes)
        midpoints = replace_invalid_3grain(midpoints, is_valid, ij)
    else:
        is_valid = torch.ones(midpoints.shape[0])==1
    
    return angles, midpoints, ij, is_valid
//  [python]
def estimate_angles(ratios, outcomes):
    #"ratios" - torch.Tensor, shape=(number of grains, number of neighborhoods), ratios found in each neighborhood
    #"outcomes" - torch.Tensor, shape=(number of grains, number of neighborhoods), which ratio was chosen to flip to
    
    d = ratios.shape[0]
    ij = torch.combinations(torch.arange(d), r=2)
    e = (torch.ones(d)/d)[:,None] #intersecting point when all factors are equal
    e = e/torch.linalg.norm(e)
    
    midpoints = []
    angles = []
    for i, j in ij: #for each grain ID pair
    
        #New center
        c = torch.zeros(d)[:,None] 
        c[i] = 0.5
        c[j] = 0.5
        
        #Find max angle from grain IDs i to j
        v = torch.zeros(d)[:,None]; v[i] = 1 #find angle from this vector
        rs = ratios[:,outcomes[i]==1] #for these ratios
        rmi, ami = find_radial_max(rs, c, v) #angle max from i
        # rmi = (v-c)*np.cos(ami/180*np.pi) + (e-c)*np.sin(ami/180*np.pi) + c
        
        #Find max angle from grain j to i
        v = torch.zeros(d)[:,None]; v[j] = 1 #find angle from this vector
        rs = ratios[:,outcomes[j]==1] #for these ratios
        rmj, amj = find_radial_max(rs, c, v) #angle max from j
        # rmj = (v-c)*np.cos(amj/180*np.pi) + (e-c)*np.sin(amj/180*np.pi) + c
        
        #Find a point in the dividing line
        midpoints.append((rmi+rmj)/2)
        
        #Average angles
        angles.append((ami+180-amj)/2)
        
    #Convert to Tensors
    angles = torch.stack(angles)[:,None]
    midpoints = torch.stack(midpoints)
    
    if ratios.shape[0]==3:
        #Replace invalid midpoints
        is_valid = is_valid_3grains(ratios, outcomes)
        midpoints = replace_invalid_3grain(midpoints, is_valid, ij)
    else:
        is_valid = torch.ones(midpoints.shape[0])==1
    
    return angles, midpoints, ij, is_valid
//  [python]
def forward(self, x, y):
        """
        computes (1/2) tr |dy/dx|^2
        :param x: input
        :param y: output of the network
        :return:
        """
        B,C = y.shape
        if self.n == -1:
            num_proj = C
        else:
            num_proj = self.n
        J2 = 0
        for ii in range(num_proj):
            if self.n == -1:
                # orthonormal vector, sequentially spanned
                v=torch.zeros(B,C)
                v[:,ii]=1
            else:
                # random properly-normalized vector for each sample
                v = self._random_vector(C=C,B=B)
            if x.is_cuda:
                v = v.cuda()
            Jv = self._jacobian_vector_product(y, x, v, create_graph=True)
            J2 += C*torch.norm(Jv)**2 / (num_proj*B)
        R = (1/2)*J2
        return R
//  [python]
def forward(self, x, y):
        """
        computes (1/2) tr |dy/dx|^2
        :param x: input
        :param y: output of the network
        :return:
        """
        B,C = y.shape
        if self.n == -1:
            num_proj = C
        else:
            num_proj = self.n
        J2 = 0
        for ii in range(num_proj):
            if self.n == -1:
                # orthonormal vector, sequentially spanned
                v=torch.zeros(B,C)
                v[:,ii]=1
            else:
                # random properly-normalized vector for each sample
                v = self._random_vector(C=C,B=B)
            if x.is_cuda:
                v = v.cuda()
            Jv = self._jacobian_vector_product(y, x, v, create_graph=True)
            J2 += C*torch.linalg.norm(Jv)**2 / (num_proj*B)
        R = (1/2)*J2
        return R
//  [python]
def _random_vector(C, B):
        """
        creates a random vector of dimension C with a norm of C^(1/2)
        (as needed for the projection formula to work)

        :param C: int, number of classes
        :param B: int, number of batch elements
        :return:
        """
        if C == 1:
            return torch.ones(B)
        v=torch.randn(B,C)
        arxilirary_zero=torch.zeros(B,C)
        vnorm=torch.norm(v, 2, 1,True)
        v=torch.addcdiv(arxilirary_zero, 1.0, v, vnorm)
        return v
//  [python]
def _random_vector(C, B):
        """
        creates a random vector of dimension C with a norm of C^(1/2)
        (as needed for the projection formula to work)

        :param C: int, number of classes
        :param B: int, number of batch elements
        :return:
        """
        if C == 1:
            return torch.ones(B)
        v=torch.randn(B,C)
        arxilirary_zero=torch.zeros(B,C)
        vnorm=torch.linalg.norm(v, 2, 1,True)
        v=torch.addcdiv(arxilirary_zero, 1.0, v, vnorm)
        return v
//  [python]
def train_lqr_value_approximator(lyapunov_relu, V_lambda, R, x_equilibrium,
                                 x_lo, x_up, num_samples, lqr_S: torch.Tensor,
                                 num_epochs):
    """
    We train both lyapunov_relu and R such that ϕ(x) − ϕ(x*) + λ|R(x−x*)|₁
    approximates the lqr cost-to-go.
    """
    x_samples = utils.uniform_sample_in_box(x_lo, x_up, num_samples)
    V_samples = torch.sum((x_samples.T - x_equilibrium.reshape(
        (6, 1))) * (lqr_S @ (x_samples.T - x_equilibrium.reshape((6, 1)))),
                          dim=0).reshape((-1, 1))
    state_value_dataset = torch.utils.data.TensorDataset(x_samples, V_samples)
    R.requires_grad_(True)

    def compute_v(model, x):
        return model(x) - model(x_equilibrium) + V_lambda * torch.norm(
            R @ (x - x_equilibrium.reshape((1, 6))).T, p=1, dim=0).reshape(
                (-1, 1))

    utils.train_approximator(state_value_dataset,
                             lyapunov_relu,
                             compute_v,
                             batch_size=50,
                             num_epochs=num_epochs,
                             lr=0.001,
                             additional_variable=[R])
    R.requires_grad_(False)
//  [python]
def train_lqr_value_approximator(lyapunov_relu, V_lambda, R, x_equilibrium,
                                 x_lo, x_up, num_samples, lqr_S: torch.Tensor,
                                 num_epochs):
    """
    We train both lyapunov_relu and R such that ϕ(x) − ϕ(x*) + λ|R(x−x*)|₁
    approximates the lqr cost-to-go.
    """
    x_samples = utils.uniform_sample_in_box(x_lo, x_up, num_samples)
    V_samples = torch.sum((x_samples.T - x_equilibrium.reshape(
        (6, 1))) * (lqr_S @ (x_samples.T - x_equilibrium.reshape((6, 1)))),
                          dim=0).reshape((-1, 1))
    state_value_dataset = torch.utils.data.TensorDataset(x_samples, V_samples)
    R.requires_grad_(True)

    def compute_v(model, x):
        return model(x) - model(x_equilibrium) + V_lambda * torch.linalg.norm(
            R @ (x - x_equilibrium.reshape((1, 6))).T, ord=1, dim=0).reshape(
                (-1, 1))

    utils.train_approximator(state_value_dataset,
                             lyapunov_relu,
                             compute_v,
                             batch_size=50,
                             num_epochs=num_epochs,
                             lr=0.001,
                             additional_variable=[R])
    R.requires_grad_(False)
//  [python]
def add_implant_residual(r,keyp,ikeyp,body_support_0, setpoint = 0.0135, scaling = 1.):
    #  [c_hip,c_ass,c_mid,c_nose,c_tip,c_impl]
    c_impl = body_support_0[5][...,0]
    keyp_implant = (ikeyp == 0)

    n_keyp = sum(keyp_implant)
    if n_keyp > 0:
        # these are n x 3, i.e. n x xyz
        target_keyp = keyp[keyp_implant,:]
        keypoint_distance = torch.norm( c_impl[:,np.newaxis,:] - target_keyp[np.newaxis,:,:] ,dim=2)
        # get the smallest distance

        r_implant = scaling * torch.abs(keypoint_distance - setpoint)
        #         print(r_implant)
        # r = torch.cat([r,r_implant],dim=1)
        r = r+torch.mean(r_implant,dim=1).unsqueeze(1)
    return r / (1.+scaling)
//  [python]
def add_implant_residual(r,keyp,ikeyp,body_support_0, setpoint = 0.0135, scaling = 1.):
    #  [c_hip,c_ass,c_mid,c_nose,c_tip,c_impl]
    c_impl = body_support_0[5][...,0]
    keyp_implant = (ikeyp == 0)

    n_keyp = sum(keyp_implant)
    if n_keyp > 0:
        # these are n x 3, i.e. n x xyz
        target_keyp = keyp[keyp_implant,:]
        keypoint_distance = torch.linalg.norm( c_impl[:,np.newaxis,:] - target_keyp[np.newaxis,:,:] ,dim=2)
        # get the smallest distance

        r_implant = scaling * torch.abs(keypoint_distance - setpoint)
        #         print(r_implant)
        # r = torch.cat([r,r_implant],dim=1)
        r = r+torch.mean(r_implant,dim=1).unsqueeze(1)
    return r / (1.+scaling)
//  [python]
def add_ear_residual(r,keyp,ikeyp,body_support_0,body_support_1,which_keyp = 1, setpoint = 0.015, scaling = 10.):
    #     # stack on first dim?
    #  [c_hip,c_ass,c_mid,c_nose,c_tip,c_impl]
    c_impl = torch.cat(( body_support_0[3][...,0], body_support_1[3][...,0]))

    keyp_implant = (ikeyp == which_keyp)
    n_keyp = sum(keyp_implant)
    if n_keyp > 0:
        # these are n x 3, i.e. n x xyz
        target_keyp = keyp[keyp_implant,:]
        keypoint_distance = torch.norm( c_impl[:,np.newaxis,:] - target_keyp[np.newaxis,:,:] ,dim=2)
        # get the smallest distance
        keypoint_distance = torch.min(keypoint_distance[:r.shape[0],:], keypoint_distance[r.shape[0]:,:])

        r_implant = scaling * torch.abs(keypoint_distance - setpoint)

        r = torch.cat([r,r_implant],dim=1)
    return r
//  [python]
def add_ear_residual(r,keyp,ikeyp,body_support_0,body_support_1,which_keyp = 1, setpoint = 0.015, scaling = 10.):
    #     # stack on first dim?
    #  [c_hip,c_ass,c_mid,c_nose,c_tip,c_impl]
    c_impl = torch.cat(( body_support_0[3][...,0], body_support_1[3][...,0]))

    keyp_implant = (ikeyp == which_keyp)
    n_keyp = sum(keyp_implant)
    if n_keyp > 0:
        # these are n x 3, i.e. n x xyz
        target_keyp = keyp[keyp_implant,:]
        keypoint_distance = torch.linalg.norm( c_impl[:,np.newaxis,:] - target_keyp[np.newaxis,:,:] ,dim=2)
        # get the smallest distance
        keypoint_distance = torch.min(keypoint_distance[:r.shape[0],:], keypoint_distance[r.shape[0]:,:])

        r_implant = scaling * torch.abs(keypoint_distance - setpoint)

        r = torch.cat([r,r_implant],dim=1)
    return r
//  [python]
def add_ass_residual(r,keyp,ikeyp,body_support_0,body_support_1,which_keyp = 3, setpoint = 0.0, scaling = 10.):
    #     # stack on first dim?
    #  [c_hip,c_ass,c_mid,c_nose,c_tip,c_impl]
    c_impl = torch.cat(( body_support_0[1][...,0], body_support_1[1][...,0]))

    keyp_implant = (ikeyp == which_keyp)
    n_keyp = sum(keyp_implant)
    if n_keyp > 0:
    # these are n x 3, i.e. n x xyz
        target_keyp = keyp[keyp_implant,:]
        keypoint_distance = torch.norm( c_impl[:,np.newaxis,:] - target_keyp[np.newaxis,:,:] ,dim=2)
        # get the smallest distance
        keypoint_distance = torch.min(keypoint_distance[:r.shape[0],:], keypoint_distance[r.shape[0]:,:])

        r_implant = scaling * torch.abs(keypoint_distance - setpoint)

        r = torch.cat([r,r_implant],dim=1)
    return r
//  [python]
def add_ass_residual(r,keyp,ikeyp,body_support_0,body_support_1,which_keyp = 3, setpoint = 0.0, scaling = 10.):
    #     # stack on first dim?
    #  [c_hip,c_ass,c_mid,c_nose,c_tip,c_impl]
    c_impl = torch.cat(( body_support_0[1][...,0], body_support_1[1][...,0]))

    keyp_implant = (ikeyp == which_keyp)
    n_keyp = sum(keyp_implant)
    if n_keyp > 0:
    # these are n x 3, i.e. n x xyz
        target_keyp = keyp[keyp_implant,:]
        keypoint_distance = torch.linalg.norm( c_impl[:,np.newaxis,:] - target_keyp[np.newaxis,:,:] ,dim=2)
        # get the smallest distance
        keypoint_distance = torch.min(keypoint_distance[:r.shape[0],:], keypoint_distance[r.shape[0]:,:])

        r_implant = scaling * torch.abs(keypoint_distance - setpoint)

        r = torch.cat([r,r_implant],dim=1)
    return r
//  [python]
def batch_rodrigues(theta):
    """Convert axis-angle representation to rotation matrix.

    Args:
        theta: size = [B, 3]
    Returns:
        Rotation matrix corresponding to the quaternion -- size = [B, 3, 3]
    """
    l1norm = torch.norm(theta + 1e-8, p=2, dim=1)
    angle = torch.unsqueeze(l1norm, -1)
    normalized = torch.div(theta, angle)
    angle = angle * 0.5
    v_cos = torch.cos(angle)
    v_sin = torch.sin(angle)
    quat = torch.cat([v_cos, v_sin * normalized], dim=1)
    return quat_to_rotmat(quat)
//  [python]
def batch_rodrigues(theta):
    """Convert axis-angle representation to rotation matrix.

    Args:
        theta: size = [B, 3]
    Returns:
        Rotation matrix corresponding to the quaternion -- size = [B, 3, 3]
    """
    l1norm = torch.linalg.norm(theta + 1e-8, ord=2, dim=1)
    angle = torch.unsqueeze(l1norm, -1)
    normalized = torch.div(theta, angle)
    angle = angle * 0.5
    v_cos = torch.cos(angle)
    v_sin = torch.sin(angle)
    quat = torch.cat([v_cos, v_sin * normalized], dim=1)
    return quat_to_rotmat(quat)
//  [python]
def add_body_residual(r,keyp,ikeyp,body_support_0,body_support_1,bpart = 'ass', setpoint = 0.0, scaling = 10.):
    #  [c_hip,c_ass,c_mid,c_nose,c_tip,c_impl]
    if bpart == 'ear':
        which_keyp = 1
        which_support = 3
    elif bpart == 'nose':
        which_keyp = 2
        which_support = 4
    elif bpart == 'ass':
        which_keyp = 3
        which_support = 1

    c_impl = torch.cat(( body_support_0[which_support][...,0], body_support_1[which_support][...,0]))

    keyp_implant = (ikeyp == which_keyp)
    n_keyp = sum(keyp_implant)
    if n_keyp > 0:
    # these are n x 3, i.e. n x xyz
        target_keyp = keyp[keyp_implant,:]
        keypoint_distance = torch.norm( c_impl[:,np.newaxis,:] - target_keyp[np.newaxis,:,:] ,dim=2)
        # get the smallest distance
        keypoint_distance = torch.min(keypoint_distance[:r.shape[0],:], keypoint_distance[r.shape[0]:,:])

        r_implant = scaling * torch.abs(keypoint_distance - setpoint)

        # r = torch.cat([r,r_implant],dim=1)
        r = r+torch.mean(r_implant,dim=1).unsqueeze(1)
    return r / (1.+scaling)
//  [python]
def add_body_residual(r,keyp,ikeyp,body_support_0,body_support_1,bpart = 'ass', setpoint = 0.0, scaling = 10.):
    #  [c_hip,c_ass,c_mid,c_nose,c_tip,c_impl]
    if bpart == 'ear':
        which_keyp = 1
        which_support = 3
    elif bpart == 'nose':
        which_keyp = 2
        which_support = 4
    elif bpart == 'ass':
        which_keyp = 3
        which_support = 1

    c_impl = torch.cat(( body_support_0[which_support][...,0], body_support_1[which_support][...,0]))

    keyp_implant = (ikeyp == which_keyp)
    n_keyp = sum(keyp_implant)
    if n_keyp > 0:
    # these are n x 3, i.e. n x xyz
        target_keyp = keyp[keyp_implant,:]
        keypoint_distance = torch.linalg.norm( c_impl[:,np.newaxis,:] - target_keyp[np.newaxis,:,:] ,dim=2)
        # get the smallest distance
        keypoint_distance = torch.min(keypoint_distance[:r.shape[0],:], keypoint_distance[r.shape[0]:,:])

        r_implant = scaling * torch.abs(keypoint_distance - setpoint)

        # r = torch.cat([r,r_implant],dim=1)
        r = r+torch.mean(r_implant,dim=1).unsqueeze(1)
    return r / (1.+scaling)
//  [python]
def rotation_matrix_vec2vec(f,t):
    # from this paper, ffrom math stacj
    # but made batch-able for pytorch
    # https://math.stackexchange.com/questions/180418/calculate-rotation-matrix-to-align-vector-a-to-vector-b-in-3d/476311#476311
    #rotate vector f onto vector t
    # import numpy as np
    # v = np.cross(f, t)
    # u = v/np.linalg.norm(v)
    # c = np.dot(f, t)
    # h = (1 - c)/(1 - c**2)

    # vx, vy, vz = v
    # rot =[[c + h*vx**2, h*vx*vy - vz, h*vx*vz + vy],
    #       [h*vx*vy+vz, c+h*vy**2, h*vy*vz-vx],
    #       [h*vx*vz - vy, h*vy*vz + vx, c+h*vz**2]]

    # rotate f onto t
    # very fast, but slightly numerically unstable, so we add epsilon!
    epsilon = 1e-6
    # f = x_pointer
    # t = nose_pointer
    # cross product
    v = torch.cross(f,t)
    u = v/(torch.norm(v,dim=1).unsqueeze(1) + epsilon)
    # dot product
    c = torch.einsum('ai,ai->a', [f,t])
    # the factor h
    h = (1 - c)/(1 - c**2 + epsilon)

    vx, vy, vz = v[:,0],v[:,1],v[:,2]

    R = torch.stack([torch.stack([c + h*vx**2, h*vx*vy - vz, h*vx*vz + vy], dim=1),
                         torch.stack([h*vx*vy+vz, c+h*vy**2, h*vy*vz-vx],      dim=1),
                         torch.stack([h*vx*vz - vy, h*vy*vz + vx, c+h*vz**2], dim=1)], dim=1)

    return R
//  [python]
def rotation_matrix_vec2vec(f,t):
    # from this paper, ffrom math stacj
    # but made batch-able for pytorch
    # https://math.stackexchange.com/questions/180418/calculate-rotation-matrix-to-align-vector-a-to-vector-b-in-3d/476311#476311
    #rotate vector f onto vector t
    # import numpy as np
    # v = np.cross(f, t)
    # u = v/np.linalg.norm(v)
    # c = np.dot(f, t)
    # h = (1 - c)/(1 - c**2)

    # vx, vy, vz = v
    # rot =[[c + h*vx**2, h*vx*vy - vz, h*vx*vz + vy],
    #       [h*vx*vy+vz, c+h*vy**2, h*vy*vz-vx],
    #       [h*vx*vz - vy, h*vy*vz + vx, c+h*vz**2]]

    # rotate f onto t
    # very fast, but slightly numerically unstable, so we add epsilon!
    epsilon = 1e-6
    # f = x_pointer
    # t = nose_pointer
    # cross product
    v = torch.cross(f,t)
    u = v/(torch.linalg.norm(v,dim=1).unsqueeze(1) + epsilon)
    # dot product
    c = torch.einsum('ai,ai->a', [f,t])
    # the factor h
    h = (1 - c)/(1 - c**2 + epsilon)

    vx, vy, vz = v[:,0],v[:,1],v[:,2]

    R = torch.stack([torch.stack([c + h*vx**2, h*vx*vy - vz, h*vx*vz + vy], dim=1),
                         torch.stack([h*vx*vy+vz, c+h*vy**2, h*vy*vz-vx],      dim=1),
                         torch.stack([h*vx*vz - vy, h*vy*vz + vx, c+h*vz**2], dim=1)], dim=1)

    return R
//  [python]
def add_nose_residual(r,keyp,ikeyp,body_support_0,body_support_1,which_keyp = 2, setpoint = 0.0, scaling = 10.):
    #     # stack on first dim?
    #  [c_hip,c_ass,c_mid,c_nose,c_tip,c_impl]
    #  [0,    1,    2 ,   3,     4 ,   5     ]
    c_impl = torch.cat(( body_support_0[4][...,0], body_support_1[4][...,0]))

    keyp_implant = (ikeyp == which_keyp)

    n_keyp = sum(keyp_implant)
    if n_keyp > 0:

        # these are n x 3, i.e. n x xyz
        target_keyp = keyp[keyp_implant,:]
        keypoint_distance = torch.norm( c_impl[:,np.newaxis,:] - target_keyp[np.newaxis,:,:] ,dim=2)
        # get the smallest distance
        keypoint_distance = torch.min(keypoint_distance[:r.shape[0],:], keypoint_distance[r.shape[0]:,:])

        r_implant = scaling * torch.abs(keypoint_distance - setpoint)

        r = torch.cat([r,r_implant],dim=1)
    return r
//  [python]
def add_nose_residual(r,keyp,ikeyp,body_support_0,body_support_1,which_keyp = 2, setpoint = 0.0, scaling = 10.):
    #     # stack on first dim?
    #  [c_hip,c_ass,c_mid,c_nose,c_tip,c_impl]
    #  [0,    1,    2 ,   3,     4 ,   5     ]
    c_impl = torch.cat(( body_support_0[4][...,0], body_support_1[4][...,0]))

    keyp_implant = (ikeyp == which_keyp)

    n_keyp = sum(keyp_implant)
    if n_keyp > 0:

        # these are n x 3, i.e. n x xyz
        target_keyp = keyp[keyp_implant,:]
        keypoint_distance = torch.linalg.norm( c_impl[:,np.newaxis,:] - target_keyp[np.newaxis,:,:] ,dim=2)
        # get the smallest distance
        keypoint_distance = torch.min(keypoint_distance[:r.shape[0],:], keypoint_distance[r.shape[0]:,:])

        r_implant = scaling * torch.abs(keypoint_distance - setpoint)

        r = torch.cat([r,r_implant],dim=1)
    return r
//  [python]
c_py = torchsolve(K_xx, b)
sync()
end = time.time()
print("Timing (PyTorch implementation):", round(end - start, 5), "s")
print("Relative error = ", (torch.norm(c - c_py) / torch.norm(c_py)).item())
//  [python]
c_py = torchsolve(K_xx, b)
sync()
end = time.time()
print("Timing (PyTorch implementation):", round(end - start, 5), "s")
print("Relative error = ", (torch.linalg.norm(c - c_py) / torch.linalg.norm(c_py)).item())
//  [python]
def forward(self, mels: torch.FloatTensor):
        """
        Computes the embeddings of a batch of utterance spectrograms.
        :param mels: a batch of mel spectrograms of same duration as a float32 tensor of shape
        (batch_size, n_frames, n_channels)
        :return: the embeddings as a float 32 tensor of shape (batch_size, embedding_size).
        Embeddings are positive and L2-normed, thus they lay in the range [0, 1].
        """
        # Pass the input through the LSTM layers and retrieve the final hidden state of the last
        # layer. Apply a cutoff to 0 for negative values and L2 normalize the embeddings.
        _, (hidden, _) = self.lstm(mels)
        embeds_raw = self.relu(self.linear(hidden[-1]))
        return embeds_raw / torch.norm(embeds_raw, dim=1, keepdim=True)
//  [python]
def forward(self, mels: torch.FloatTensor):
        """
        Computes the embeddings of a batch of utterance spectrograms.
        :param mels: a batch of mel spectrograms of same duration as a float32 tensor of shape
        (batch_size, n_frames, n_channels)
        :return: the embeddings as a float 32 tensor of shape (batch_size, embedding_size).
        Embeddings are positive and L2-normed, thus they lay in the range [0, 1].
        """
        # Pass the input through the LSTM layers and retrieve the final hidden state of the last
        # layer. Apply a cutoff to 0 for negative values and L2 normalize the embeddings.
        _, (hidden, _) = self.lstm(mels)
        embeds_raw = self.relu(self.linear(hidden[-1]))
        return embeds_raw / torch.linalg.norm(embeds_raw, dim=1, keepdim=True)
//  [python]
def compute_sqnr(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    """
    Computes the SQNR between `x` and `y`.

    Args:
        x: Tensor or tuple of tensors
        y: Tensor or tuple of tensors

    Return:
        float or tuple of floats
    """
    Ps = torch.norm(x)
    Pn = torch.norm(x - y)
    return 20 * torch.log10(Ps / Pn)
//  [python]
def compute_sqnr(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    """
    Computes the SQNR between `x` and `y`.

    Args:
        x: Tensor or tuple of tensors
        y: Tensor or tuple of tensors

    Return:
        float or tuple of floats
    """
    Ps = torch.linalg.norm(x)
    Pn = torch.linalg.norm(x - y)
    return 20 * torch.log10(Ps / Pn)
//  [python]
def check_if_mirror_extended(self, lqs):
        """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:
            lqs (tensor): Input low quality (LQ) sequence with
                shape (n, t, c, h, w).
        """

        if lqs.size(1) % 2 == 0:
            lqs_1, lqs_2 = torch.chunk(lqs, 2, dim=1)
            if torch.norm(lqs_1 - lqs_2.flip(1)) == 0:
                self.is_mirror_extended = True
//  [python]
def check_if_mirror_extended(self, lqs):
        """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:
            lqs (tensor): Input low quality (LQ) sequence with
                shape (n, t, c, h, w).
        """

        if lqs.size(1) % 2 == 0:
            lqs_1, lqs_2 = torch.chunk(lqs, 2, dim=1)
            if torch.linalg.norm(lqs_1 - lqs_2.flip(1)) == 0:
                self.is_mirror_extended = True