jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def norm(self, tensor):
        tensor_format = getattr(self, "tensor_format", "pt")
        if tensor_format == "np":
            return np.linalg.norm(tensor)
        elif tensor_format == "pt":
            return torch.norm(tensor.reshape(tensor.shape[0], -1), dim=-1).mean()

        raise ValueError(f"`self.tensor_format`: {self.tensor_format} is not valid.")
//  [python]
def norm(self, tensor):
        tensor_format = getattr(self, "tensor_format", "pt")
        if tensor_format == "np":
            return np.linalg.norm(tensor)
        elif tensor_format == "pt":
            return torch.linalg.norm(tensor.reshape(tensor.shape[0], -1), dim=-1).mean()

        raise ValueError(f"`self.tensor_format`: {self.tensor_format} is not valid.")
//  [python]
def sphere_compliment_case():
    torch.manual_seed(42)
    shape = manifold_shapes[geoopt.manifolds.Sphere]
    complement = torch.rand(shape[-1], 1, dtype=torch.float64)

    Q, _ = geoopt.linalg.batch_linalg.qr(complement, "reduced")
    P = -Q @ Q.transpose(-1, -2)
    P[..., torch.arange(P.shape[-2]), torch.arange(P.shape[-2])] += 1

    ex = torch.randn(*shape, dtype=torch.float64)
    ev = torch.randn(*shape, dtype=torch.float64)
    x = (ex @ P.t()) / torch.norm(ex @ P.t())
    v = (ev - (x @ ev) * x) @ P.t()

    manifold = geoopt.Sphere(complement=complement)
    x = geoopt.ManifoldTensor(x, manifold=manifold)
    case = UnaryCase(shape, x, ex, v, ev, manifold)
    yield case
    manifold = geoopt.SphereExact(complement=complement)
    x = geoopt.ManifoldTensor(x, manifold=manifold)
    case = UnaryCase(shape, x, ex, v, ev, manifold)
    yield case
//  [python]
def sphere_compliment_case():
    torch.manual_seed(42)
    shape = manifold_shapes[geoopt.manifolds.Sphere]
    complement = torch.rand(shape[-1], 1, dtype=torch.float64)

    Q, _ = geoopt.linalg.batch_linalg.qr(complement, "reduced")
    P = -Q @ Q.transpose(-1, -2)
    P[..., torch.arange(P.shape[-2]), torch.arange(P.shape[-2])] += 1

    ex = torch.randn(*shape, dtype=torch.float64)
    ev = torch.randn(*shape, dtype=torch.float64)
    x = (ex @ P.t()) / torch.linalg.norm(ex @ P.t())
    v = (ev - (x @ ev) * x) @ P.t()

    manifold = geoopt.Sphere(complement=complement)
    x = geoopt.ManifoldTensor(x, manifold=manifold)
    case = UnaryCase(shape, x, ex, v, ev, manifold)
    yield case
    manifold = geoopt.SphereExact(complement=complement)
    x = geoopt.ManifoldTensor(x, manifold=manifold)
    case = UnaryCase(shape, x, ex, v, ev, manifold)
    yield case
//  [python]
def weighted_8points(x_in, logits):
    # x_in: batch * 1 * N * 4
    x_shp = x_in.shape
    # Turn into weights for each sample
    weights = torch.sigmoid(logits)
    x_in = x_in.squeeze(1)
    
    # Make input data (num_img_pair x num_corr x 4),  b*4*n, (u,v,u',v')
    xx = torch.reshape(x_in, (x_shp[0], x_shp[2], 4)).permute(0, 2, 1)

    # Create the matrix to be used for the eight-point algorithm, (b,n,9)
    X = torch.stack([
        xx[:, 2] * xx[:, 0], xx[:, 2] * xx[:, 1], xx[:, 2],
        xx[:, 3] * xx[:, 0], xx[:, 3] * xx[:, 1], xx[:, 3],
        xx[:, 0], xx[:, 1], torch.ones_like(xx[:, 0])
    ], dim=1).permute(0, 2, 1)

    wX = torch.reshape(weights, (x_shp[0], x_shp[2], 1)) * X

    XwX = torch.matmul(X.permute(0, 2, 1), wX)    

    # Recover essential matrix from self-adjoing eigen
    v = batch_symeig(XwX)
    e_hat = torch.reshape(v[:, :, 0], (x_shp[0], 9))

    # Make unit norm just in case
    e_hat = e_hat / torch.norm(e_hat, dim=1, keepdim=True)
    return e_hat
//  [python]
def weighted_8points(x_in, logits):
    # x_in: batch * 1 * N * 4
    x_shp = x_in.shape
    # Turn into weights for each sample
    weights = torch.sigmoid(logits)
    x_in = x_in.squeeze(1)
    
    # Make input data (num_img_pair x num_corr x 4),  b*4*n, (u,v,u',v')
    xx = torch.reshape(x_in, (x_shp[0], x_shp[2], 4)).permute(0, 2, 1)

    # Create the matrix to be used for the eight-point algorithm, (b,n,9)
    X = torch.stack([
        xx[:, 2] * xx[:, 0], xx[:, 2] * xx[:, 1], xx[:, 2],
        xx[:, 3] * xx[:, 0], xx[:, 3] * xx[:, 1], xx[:, 3],
        xx[:, 0], xx[:, 1], torch.ones_like(xx[:, 0])
    ], dim=1).permute(0, 2, 1)

    wX = torch.reshape(weights, (x_shp[0], x_shp[2], 1)) * X

    XwX = torch.matmul(X.permute(0, 2, 1), wX)    

    # Recover essential matrix from self-adjoing eigen
    v = batch_symeig(XwX)
    e_hat = torch.reshape(v[:, :, 0], (x_shp[0], 9))

    # Make unit norm just in case
    e_hat = e_hat / torch.linalg.norm(e_hat, dim=1, keepdim=True)
    return e_hat
//  [python]
def qexp_t(q):
  """
  Applies exponential map to log quaternion
  :param q: N x 3
  :return: N x 4
  """
  n = torch.norm(q, p=2, dim=1, keepdim=True)
  n = torch.clamp(n, min=1e-8)
  q = q * torch.sin(n)
  q = q / n
  q = torch.cat((torch.cos(n), q), dim=1)
  return q
//  [python]
def qexp_t(q):
  """
  Applies exponential map to log quaternion
  :param q: N x 3
  :return: N x 4
  """
  n = torch.linalg.norm(q, ord=2, dim=1, keepdim=True)
  n = torch.clamp(n, min=1e-8)
  q = q * torch.sin(n)
  q = q / n
  q = torch.cat((torch.cos(n), q), dim=1)
  return q
//  [python]
def qlog_t(q):
  """
  Applies the log map to a quaternion
  :param q: N x 4
  :return: N x 3
  """
  n = torch.norm(q[:, 1:], p=2, dim=1, keepdim=True)
  n = torch.clamp(n, min=1e-8)
  q = q[:, 1:] * torch.acos(torch.clamp(q[:, :1], min=-1.0, max=1.0))
  q = q / n
  return q
//  [python]
def qlog_t(q):
  """
  Applies the log map to a quaternion
  :param q: N x 4
  :return: N x 3
  """
  n = torch.linalg.norm(q[:, 1:], ord=2, dim=1, keepdim=True)
  n = torch.clamp(n, min=1e-8)
  q = q[:, 1:] * torch.acos(torch.clamp(q[:, :1], min=-1.0, max=1.0))
  q = q / n
  return q
//  [python]
def fwd_solver(f, z_init, threshold, max_iter, mode='abs'):
    z_prev, z = z_init, f(z_init)
    nstep = 0
    while nstep < max_iter:
        z_prev, z = z, f(z)
        abs_diff = torch.norm(z_prev - z).item()
        rel_diff = abs_diff / (torch.norm(z_prev).item() + 1e-9)
        diff_dict = {'abs': abs_diff, 'rel': rel_diff}
        # print(f'abs_diff: {diff_dict["abs"]}, rel_diff: {rel_diff}')
        if diff_dict[mode] < threshold:
            # print(nstep, abs_diff)
            break
        nstep += 1
    if nstep == max_iter:
        print(f'step {nstep}, not converged, abs_diff: {abs_diff}')
    return z, abs_diff
//  [python]
def fwd_solver(f, z_init, threshold, max_iter, mode='abs'):
    z_prev, z = z_init, f(z_init)
    nstep = 0
    while nstep < max_iter:
        z_prev, z = z, f(z)
        abs_diff = torch.linalg.norm(z_prev - z).item()
        rel_diff = abs_diff / (torch.linalg.norm(z_prev).item() + 1e-9)
        diff_dict = {'abs': abs_diff, 'rel': rel_diff}
        # print(f'abs_diff: {diff_dict["abs"]}, rel_diff: {rel_diff}')
        if diff_dict[mode] < threshold:
            # print(nstep, abs_diff)
            break
        nstep += 1
    if nstep == max_iter:
        print(f'step {nstep}, not converged, abs_diff: {abs_diff}')
    return z, abs_diff
//  [python]
def test_solve(self):
        torch.manual_seed(0)
        a = torch.rand(10, 10)
        b = torch.rand(10)
        a /= 2*torch.norm(a) # ensure that x exists
        x = torch.linalg.inv(torch.eye(*a.shape)-a) @ b
        
        with self.subTest(semiring='RealSemiring'):
            self.assertAlmostEqual(RealSemiring().solve(a, b), x)
        with self.subTest(semiring='LogSemiring'):
            self.assertAlmostEqual(LogSemiring().solve(a.log(), b.log()), x.log())
//  [python]
def test_solve(self):
        torch.manual_seed(0)
        a = torch.rand(10, 10)
        b = torch.rand(10)
        a /= 2*torch.linalg.norm(a) # ensure that x exists
        x = torch.linalg.inv(torch.eye(*a.shape)-a) @ b
        
        with self.subTest(semiring='RealSemiring'):
            self.assertAlmostEqual(RealSemiring().solve(a, b), x)
        with self.subTest(semiring='LogSemiring'):
            self.assertAlmostEqual(LogSemiring().solve(a.log(), b.log()), x.log())
//  [python]
def forward(self, x, mask):
        z, z_out = self.encoder(x, mask)

        x_rec = self.decoder(z, z_out)

        mask = (
            mask.repeat_interleave(self.patch_size[0], 1)
            .repeat_interleave(self.patch_size[1], 2)
            .repeat_interleave(self.patch_size[2], 3)
            .unsqueeze(1)
            .contiguous()
        )

        if self.loss == "mask_only":
            loss_recon = F.l1_loss(x, x_rec, reduction="none")
            loss = (loss_recon * mask).sum() / (mask.sum() + 1e-5) / self.in_chans
            # _, loss_var_l2 = self.get_image_prior_losses(x_rec)
            # bs = x.shape[0]
            # img_l2_loss = torch.norm(x_rec.view(bs, -1), dim=1).mean()
            # loss = loss + 0.0001*loss_var_l2+1e-5*img_l2_loss

        elif self.loss == "all_img":
            loss_recon = F.l1_loss(x, x_rec, reduction="mean")
            loss = loss_recon

        elif self.loss == "l2":
            loss = F.mse_loss(x, x_rec, reduction="mean")
            _, loss_var_l2 = self.get_image_prior_losses(x_rec)
            bs = x.shape[0]
            img_l2_loss = torch.norm(x_rec.view(bs, -1), dim=1).mean()
            loss = loss + 0.0001 * loss_var_l2 + 1e-5 * img_l2_loss

        return loss, x_rec, mask
//  [python]
def forward(self, x, mask):
        z, z_out = self.encoder(x, mask)

        x_rec = self.decoder(z, z_out)

        mask = (
            mask.repeat_interleave(self.patch_size[0], 1)
            .repeat_interleave(self.patch_size[1], 2)
            .repeat_interleave(self.patch_size[2], 3)
            .unsqueeze(1)
            .contiguous()
        )

        if self.loss == "mask_only":
            loss_recon = F.l1_loss(x, x_rec, reduction="none")
            loss = (loss_recon * mask).sum() / (mask.sum() + 1e-5) / self.in_chans
            # _, loss_var_l2 = self.get_image_prior_losses(x_rec)
            # bs = x.shape[0]
            # img_l2_loss = torch.linalg.norm(x_rec.view(bs, -1), dim=1).mean()
            # loss = loss + 0.0001*loss_var_l2+1e-5*img_l2_loss

        elif self.loss == "all_img":
            loss_recon = F.l1_loss(x, x_rec, reduction="mean")
            loss = loss_recon

        elif self.loss == "l2":
            loss = F.mse_loss(x, x_rec, reduction="mean")
            _, loss_var_l2 = self.get_image_prior_losses(x_rec)
            bs = x.shape[0]
            img_l2_loss = torch.linalg.norm(x_rec.view(bs, -1), dim=1).mean()
            loss = loss + 0.0001 * loss_var_l2 + 1e-5 * img_l2_loss

        return loss, x_rec, mask
//  [python]
def print_model_info(model, logfile, summaryfile):
    model_info = ["Model info:"]

    for module_name, module in model.named_modules():
        module_info = []

        # Submodule-specific extra info
        if hasattr(module, "model_info"):
            module_info.extend(module.model_info())

        for frame_name, frame in parseval_frames(module):
            if frame.shape[0] > frame.shape[1]:
                gram = frame.conj().t() @ frame
            else:
                gram = frame @ frame.conj().t()
            singular_values = torch.symeig(gram)[0]
            min_sv, max_sv = singular_values[0], singular_values[-1]
            module_info.append(
                f"\n  - Parseval on {frame_name}: singular values {tensor_summary_stats(singular_values)}")

        for frame_name, frame in quadrature_frames(module):
            #norm_ratio = torch.norm(frame.t() @ frame)/(torch.norm(frame)**2)
            norm_ratio = torch.norm(frame @ frame.t()) / (torch.norm(frame)**2)
            norm_frame = torch.norm(frame)
            #module_info.append(f"\n  - Quadrature on {frame_name}: ratio norm W^T W / norm W^2 {norm_ratio:.3f}")
            module_info.append(f"\n  - Quadrature on {frame_name}: ratio norm W W^T / norm W^2 {norm_ratio:.3f}")
            module_info.append(f"\n  - Frame norm of {frame_name}: {norm_frame:.3f}")

        if len(module_info) > 0:
            model_info.extend([f"\n- {module_name} ({module.__class__.__name__}): "] + module_info)

    print_and_write("".join(model_info), logfile, summaryfile)
//  [python]
def print_model_info(model, logfile, summaryfile):
    model_info = ["Model info:"]

    for module_name, module in model.named_modules():
        module_info = []

        # Submodule-specific extra info
        if hasattr(module, "model_info"):
            module_info.extend(module.model_info())

        for frame_name, frame in parseval_frames(module):
            if frame.shape[0] > frame.shape[1]:
                gram = frame.conj().t() @ frame
            else:
                gram = frame @ frame.conj().t()
            singular_values = torch.symeig(gram)[0]
            min_sv, max_sv = singular_values[0], singular_values[-1]
            module_info.append(
                f"\n  - Parseval on {frame_name}: singular values {tensor_summary_stats(singular_values)}")

        for frame_name, frame in quadrature_frames(module):
            #norm_ratio = torch.linalg.norm(frame.t() @ frame)/(torch.linalg.norm(frame)**2)
            norm_ratio = torch.linalg.norm(frame @ frame.t()) / (torch.linalg.norm(frame)**2)
            norm_frame = torch.linalg.norm(frame)
            #module_info.append(f"\n  - Quadrature on {frame_name}: ratio norm W^T W / norm W^2 {norm_ratio:.3f}")
            module_info.append(f"\n  - Quadrature on {frame_name}: ratio norm W W^T / norm W^2 {norm_ratio:.3f}")
            module_info.append(f"\n  - Frame norm of {frame_name}: {norm_frame:.3f}")

        if len(module_info) > 0:
            model_info.extend([f"\n- {module_name} ({module.__class__.__name__}): "] + module_info)

    print_and_write("".join(model_info), logfile, summaryfile)
//  [python]
def _ncc_c_3dim(x, y):
    den = torch.norm(x, p=2, dim=(0, 1)) * torch.norm(y, p=2, dim=(0, 1))

    if den < 1e-9:
        den = torch.tensor(float("inf"), device="cuda", dtype=torch.float32)

    x_len = x.shape[0]
    fft_size = 1 << (2*x_len-1).bit_length()

    cc = torch.fft.ifft(torch.fft.fft(x, fft_size, dim=0) * torch.conj(torch.fft.fft(y, fft_size, dim=0)), dim=0)
    cc = torch.cat((cc[-(x_len-1):], cc[:x_len]), dim=0)

    return torch.div(torch.sum(torch.real(cc), dim=-1), den)
//  [python]
def _ncc_c_3dim(x, y):
    den = torch.linalg.norm(x, ord=2, dim=(0, 1)) * torch.linalg.norm(y, ord=2, dim=(0, 1))

    if den < 1e-9:
        den = torch.tensor(float("inf"), device="cuda", dtype=torch.float32)

    x_len = x.shape[0]
    fft_size = 1 << (2*x_len-1).bit_length()

    cc = torch.fft.ifft(torch.fft.fft(x, fft_size, dim=0) * torch.conj(torch.fft.fft(y, fft_size, dim=0)), dim=0)
    cc = torch.cat((cc[-(x_len-1):], cc[:x_len]), dim=0)

    return torch.div(torch.sum(torch.real(cc), dim=-1), den)
//  [python]
def test_computes_periodic_function(self):
        a = torch.tensor([4, 2, 8], dtype=torch.float).view(3, 1)
        b = torch.tensor([0, 2], dtype=torch.float).view(2, 1)
        lengthscale = 2
        period = 3
        kernel = PeriodicKernel().initialize(lengthscale=lengthscale, period_length=period)
        kernel.eval()

        actual = torch.zeros(3, 2)
        for i in range(3):
            for j in range(2):
                val = 2 * torch.pow(torch.sin(math.pi * (a[i] - b[j]) / 3), 2) / lengthscale
                actual[i, j] = torch.exp(-val).item()

        res = kernel(a, b).to_dense()
        self.assertLess(torch.norm(res - actual), 1e-5)
//  [python]
def test_computes_periodic_function(self):
        a = torch.tensor([4, 2, 8], dtype=torch.float).view(3, 1)
        b = torch.tensor([0, 2], dtype=torch.float).view(2, 1)
        lengthscale = 2
        period = 3
        kernel = PeriodicKernel().initialize(lengthscale=lengthscale, period_length=period)
        kernel.eval()

        actual = torch.zeros(3, 2)
        for i in range(3):
            for j in range(2):
                val = 2 * torch.pow(torch.sin(math.pi * (a[i] - b[j]) / 3), 2) / lengthscale
                actual[i, j] = torch.exp(-val).item()

        res = kernel(a, b).to_dense()
        self.assertLess(torch.linalg.norm(res - actual), 1e-5)
//  [python]
def train_IRLS(X_train, y_train, X_test=None, y_test=None, L2_param=0, max_iter=MAX_ITER):
    """train Logistic Regression via IRLS algorithm
    X: Nxd
    y: Nx1
    ---

    """
    N, d = X_train.shape
    X_train = torch.as_tensor(X_train, device=device)
    X_test = torch.as_tensor(X_test, device=device)
    y_train = torch.as_tensor(y_train, device=device)
    y_test = torch.as_tensor(y_test, device=device)

    w = torch.full((d, 1), 0.01, device=device)

    print("start training...")
    tic = time.time()
    print("Device: {}".format(device))
    print("L2 param(lambda): {}".format(L2_param))
    i = 0
    # iteration
    while i <= max_iter:
        print("iter: {}".format(i))

        neg_L = neg_log_likelihood(w, X_train, y_train, L2_param)
        print("\t neg log likelihood: {}".format(neg_L.sum()))

        train_acc = compute_acc(X_train, y_train, w)
        test_acc = compute_acc(X_test, y_test, w)
        print("\t train acc: {}, test acc: {}".format(train_acc, test_acc))

        L2_norm_w = torch.norm(w)
        print("\t L2 norm of w: {}".format(L2_norm_w.item()))

        if i > 0:
            diff_w = torch.norm(w - w_old_data)
            print("\t diff of w_old and w: {}".format(diff_w.item()))
            if diff_w < 1e-2:
                break

        w_old_data = w.clone()
        w = update_weight(w, X_train, y_train, L2_param)
        i += 1
    print(f"training done, using {time.time() - tic}s.")
//  [python]
def train_IRLS(X_train, y_train, X_test=None, y_test=None, L2_param=0, max_iter=MAX_ITER):
    """train Logistic Regression via IRLS algorithm
    X: Nxd
    y: Nx1
    ---

    """
    N, d = X_train.shape
    X_train = torch.as_tensor(X_train, device=device)
    X_test = torch.as_tensor(X_test, device=device)
    y_train = torch.as_tensor(y_train, device=device)
    y_test = torch.as_tensor(y_test, device=device)

    w = torch.full((d, 1), 0.01, device=device)

    print("start training...")
    tic = time.time()
    print("Device: {}".format(device))
    print("L2 param(lambda): {}".format(L2_param))
    i = 0
    # iteration
    while i <= max_iter:
        print("iter: {}".format(i))

        neg_L = neg_log_likelihood(w, X_train, y_train, L2_param)
        print("\t neg log likelihood: {}".format(neg_L.sum()))

        train_acc = compute_acc(X_train, y_train, w)
        test_acc = compute_acc(X_test, y_test, w)
        print("\t train acc: {}, test acc: {}".format(train_acc, test_acc))

        L2_norm_w = torch.linalg.norm(w)
        print("\t L2 norm of w: {}".format(L2_norm_w.item()))

        if i > 0:
            diff_w = torch.linalg.norm(w - w_old_data)
            print("\t diff of w_old and w: {}".format(diff_w.item()))
            if diff_w < 1e-2:
                break

        w_old_data = w.clone()
        w = update_weight(w, X_train, y_train, L2_param)
        i += 1
    print(f"training done, using {time.time() - tic}s.")
//  [python]
def test_batch(self):
        a = torch.tensor([[4, 2, 8], [1, 2, 3]], dtype=torch.float).view(2, 3, 1)
        b = torch.tensor([[0, 2], [-1, 2]], dtype=torch.float).view(2, 2, 1)
        period = torch.tensor(1, dtype=torch.float).view(1, 1, 1)
        lengthscale = torch.tensor(2, dtype=torch.float).view(1, 1, 1)
        kernel = PeriodicKernel().initialize(lengthscale=lengthscale, period_length=period)
        kernel.eval()

        actual = torch.zeros(2, 3, 2)
        for k in range(2):
            actual[k] = kernel(a[k], b[k]).to_dense()

        res = kernel(a, b).to_dense()
        self.assertLess(torch.norm(res - actual), 1e-5)
//  [python]
def test_batch(self):
        a = torch.tensor([[4, 2, 8], [1, 2, 3]], dtype=torch.float).view(2, 3, 1)
        b = torch.tensor([[0, 2], [-1, 2]], dtype=torch.float).view(2, 2, 1)
        period = torch.tensor(1, dtype=torch.float).view(1, 1, 1)
        lengthscale = torch.tensor(2, dtype=torch.float).view(1, 1, 1)
        kernel = PeriodicKernel().initialize(lengthscale=lengthscale, period_length=period)
        kernel.eval()

        actual = torch.zeros(2, 3, 2)
        for k in range(2):
            actual[k] = kernel(a[k], b[k]).to_dense()

        res = kernel(a, b).to_dense()
        self.assertLess(torch.linalg.norm(res - actual), 1e-5)
//  [python]
def test_batch_separate(self):
        a = torch.tensor([[4, 2, 8], [1, 2, 3]], dtype=torch.float).view(2, 3, 1)
        b = torch.tensor([[0, 2], [-1, 2]], dtype=torch.float).view(2, 2, 1)
        period = torch.tensor([1, 2], dtype=torch.float).view(2, 1, 1)
        lengthscale = torch.tensor([2, 1], dtype=torch.float).view(2, 1, 1)
        kernel = PeriodicKernel(batch_shape=torch.Size([2])).initialize(lengthscale=lengthscale, period_length=period)
        kernel.eval()

        actual = torch.zeros(2, 3, 2)
        for k in range(2):
            diff = a[k].unsqueeze(1) - b[k].unsqueeze(0)
            diff = diff * math.pi / period[k].unsqueeze(-1)
            actual[k] = diff.sin().pow(2).sum(dim=-1).div(lengthscale[k]).mul(-2.0).exp_()

        res = kernel(a, b).to_dense()
        self.assertLess(torch.norm(res - actual), 1e-5)
//  [python]
def test_batch_separate(self):
        a = torch.tensor([[4, 2, 8], [1, 2, 3]], dtype=torch.float).view(2, 3, 1)
        b = torch.tensor([[0, 2], [-1, 2]], dtype=torch.float).view(2, 2, 1)
        period = torch.tensor([1, 2], dtype=torch.float).view(2, 1, 1)
        lengthscale = torch.tensor([2, 1], dtype=torch.float).view(2, 1, 1)
        kernel = PeriodicKernel(batch_shape=torch.Size([2])).initialize(lengthscale=lengthscale, period_length=period)
        kernel.eval()

        actual = torch.zeros(2, 3, 2)
        for k in range(2):
            diff = a[k].unsqueeze(1) - b[k].unsqueeze(0)
            diff = diff * math.pi / period[k].unsqueeze(-1)
            actual[k] = diff.sin().pow(2).sum(dim=-1).div(lengthscale[k]).mul(-2.0).exp_()

        res = kernel(a, b).to_dense()
        self.assertLess(torch.linalg.norm(res - actual), 1e-5)
//  [python]
def forward(self, query, value, key_padding_mask=None, state=None):
        projected_query = self.query_proj(query).unsqueeze(0)
        key = self.value_proj(value)
        if self.normalize:
            normed_v = self.g * self.v / torch.norm(self.v)
            attn_scores = (normed_v * torch.tanh(projected_query + key + self.b)).sum(dim=2)
        else:
            attn_scores = self.v * torch.tanh(projected_query + key).sum(dim=2)
        if key_padding_mask is not None:
            attn_scores = attn_scores.float().masked_fill_(key_padding_mask, float('-inf')).type_as(attn_scores)
        attn_scores = F.softmax(attn_scores, dim=0)
        context = (attn_scores.unsqueeze(2) * value).sum(dim=0)
        next_state = attn_scores
        return context, attn_scores, next_state
//  [python]
def forward(self, query, value, key_padding_mask=None, state=None):
        projected_query = self.query_proj(query).unsqueeze(0)
        key = self.value_proj(value)
        if self.normalize:
            normed_v = self.g * self.v / torch.linalg.norm(self.v)
            attn_scores = (normed_v * torch.tanh(projected_query + key + self.b)).sum(dim=2)
        else:
            attn_scores = self.v * torch.tanh(projected_query + key).sum(dim=2)
        if key_padding_mask is not None:
            attn_scores = attn_scores.float().masked_fill_(key_padding_mask, float('-inf')).type_as(attn_scores)
        attn_scores = F.softmax(attn_scores, dim=0)
        context = (attn_scores.unsqueeze(2) * value).sum(dim=0)
        next_state = attn_scores
        return context, attn_scores, next_state
//  [python]
def fit_dfl(treatment_feature: torch.Tensor,
                covariate_feature: Optional[torch.Tensor],
                outcome_t: torch.Tensor,
                lam: float, add_intercept: bool
                ):

        # stage1
        feature = DFLModel.augment_feature(treatment_feature,
                                           covariate_feature,
                                           add_intercept)

        weight = fit_linear(outcome_t, feature, lam)
        pred = linear_reg_pred(feature, weight)
        loss = torch.norm((outcome_t - pred)) ** 2 + lam * torch.norm(weight) ** 2

        labels = logit(outcome_t)
        logits = logit(pred)
        log_loss = (-(torch.log(logits) * labels+torch.log(1-logits) * (1-labels))).sum() + lam * torch.norm(weight) ** 2

        return dict(weight=weight, loss=loss, log_loss=log_loss)
//  [python]
def fit_dfl(treatment_feature: torch.Tensor,
                covariate_feature: Optional[torch.Tensor],
                outcome_t: torch.Tensor,
                lam: float, add_intercept: bool
                ):

        # stage1
        feature = DFLModel.augment_feature(treatment_feature,
                                           covariate_feature,
                                           add_intercept)

        weight = fit_linear(outcome_t, feature, lam)
        pred = linear_reg_pred(feature, weight)
        loss = torch.linalg.norm((outcome_t - pred)) ** 2 + lam * torch.linalg.norm(weight) ** 2

        labels = logit(outcome_t)
        logits = logit(pred)
        log_loss = (-(torch.log(logits) * labels+torch.log(1-logits) * (1-labels))).sum() + lam * torch.linalg.norm(weight) ** 2

        return dict(weight=weight, loss=loss, log_loss=log_loss)
//  [python]
def score(self,A=None,B=None,score_method=None,group=None):
        """
        Given an optimal C already computed, calculate the metric

        Parameters
        __________
        A : np.array or torch.tensor or None
            first data matrix, if None defaults to the saved matrix in fit
        B : np.array or torch.tensor or None
            second data matrix if None, defaults to the savec matrix in fit
        score_method : None or {'angular','euclidean'}
            overwrites the score method in the object for this application
        Returns
        _______

        score : float
            similarity of the data under the similarity transform w.r.t C
        """
        assert self.C_star is not None
        A = self.A if A is None else A
        B = self.B if B is None else B 
        assert A is not None
        assert B is not None
        assert A.shape == self.C_star.shape
        assert B.shape == self.C_star.shape
        score_method = self.score_method if score_method is None else score_method
        group = self.group if group is None else group
        with torch.no_grad():
            if not isinstance(A,torch.Tensor):
                A = torch.from_numpy(A).float().to(self.device)
            if not isinstance(B,torch.Tensor):
                B = torch.from_numpy(B).float().to(self.device)
            C = self.C_star.to(self.device)

        if group in {"SO(n)", "O(n)"}:
            Cinv = C.T
        elif group in {"GL(n)"}:
            Cinv = torch.linalg.inv(C)
        else:
            raise AssertionError("Need proper group name")
        if score_method == 'angular':    
            num = torch.trace(A.T @ C @ B @ Cinv) 
            den = torch.norm(A,p = 'fro')*torch.norm(B,p = 'fro')
            score = torch.arccos(num/den).cpu().numpy()
            if np.isnan(score): #around -1 and 1, we sometimes get NaNs due to arccos
                if num/den < 0:
                    score = np.pi
                else:
                    score = 0
        else:
            score = torch.norm(A - C @ B @ Cinv,p='fro').cpu().numpy().item() #/ A.numpy().size
    
        return score
//  [python]
def score(self,A=None,B=None,score_method=None,group=None):
        """
        Given an optimal C already computed, calculate the metric

        Parameters
        __________
        A : np.array or torch.tensor or None
            first data matrix, if None defaults to the saved matrix in fit
        B : np.array or torch.tensor or None
            second data matrix if None, defaults to the savec matrix in fit
        score_method : None or {'angular','euclidean'}
            overwrites the score method in the object for this application
        Returns
        _______

        score : float
            similarity of the data under the similarity transform w.r.t C
        """
        assert self.C_star is not None
        A = self.A if A is None else A
        B = self.B if B is None else B 
        assert A is not None
        assert B is not None
        assert A.shape == self.C_star.shape
        assert B.shape == self.C_star.shape
        score_method = self.score_method if score_method is None else score_method
        group = self.group if group is None else group
        with torch.no_grad():
            if not isinstance(A,torch.Tensor):
                A = torch.from_numpy(A).float().to(self.device)
            if not isinstance(B,torch.Tensor):
                B = torch.from_numpy(B).float().to(self.device)
            C = self.C_star.to(self.device)

        if group in {"SO(n)", "O(n)"}:
            Cinv = C.T
        elif group in {"GL(n)"}:
            Cinv = torch.linalg.inv(C)
        else:
            raise AssertionError("Need proper group name")
        if score_method == 'angular':    
            num = torch.trace(A.T @ C @ B @ Cinv) 
            den = torch.linalg.norm(A,ord='fro')*torch.linalg.norm(B,ord='fro')
            score = torch.arccos(num/den).cpu().numpy()
            if np.isnan(score): #around -1 and 1, we sometimes get NaNs due to arccos
                if num/den < 0:
                    score = np.pi
                else:
                    score = 0
        else:
            score = torch.linalg.norm(A - C @ B @ Cinv,ord='fro').cpu().numpy().item() #/ A.numpy().size
    
        return score
//  [python]
def Degree_matrix(x, y):
    """
    Compute the degree matrix starting from corresponding patches of data organized
    in batches. It first computes the affinity matrices of the two batches and then
    it computes the norm of the difference between the rows of Ax and the rows of Ay.
    Then it is normalized.
        Input:
            x - float, array of [batch_size, patch_size, patch_size, num_channels_x],
                Batch of patches from data domain x.
            y - float, array of [batch_size, patch_size, patch_size, num_channels_y],
                Batch of patches from data domain y.
        Output:
            D - float, array of [batch_size, patch_size^2, patch_size^2], Degree matrix
    """
    ax = affinity(x)
    ay = affinity(y)
    D = torch.norm(torch.unsqueeze(ax, 1) - torch.unsqueeze(ay, 2), 2, -1)
    D = (D - torch.min(D)) / (torch.max(D) - torch.min(D))
    return D
//  [python]
def Degree_matrix(x, y):
    """
    Compute the degree matrix starting from corresponding patches of data organized
    in batches. It first computes the affinity matrices of the two batches and then
    it computes the norm of the difference between the rows of Ax and the rows of Ay.
    Then it is normalized.
        Input:
            x - float, array of [batch_size, patch_size, patch_size, num_channels_x],
                Batch of patches from data domain x.
            y - float, array of [batch_size, patch_size, patch_size, num_channels_y],
                Batch of patches from data domain y.
        Output:
            D - float, array of [batch_size, patch_size^2, patch_size^2], Degree matrix
    """
    ax = affinity(x)
    ay = affinity(y)
    D = torch.linalg.norm(torch.unsqueeze(ax, 1) - torch.unsqueeze(ay, 2), 2, -1)
    D = (D - torch.min(D)) / (torch.max(D) - torch.min(D))
    return D
//  [python]
def affinity(x):
    """
    Compute the affinity matrices of the patches of contained in a batch.
    It first computes the distances between the datapoints within a patch.
    Then it finds the suitable kernel width for each patch.
    Finally, applies the RBF.
        Input:
            x - float, array of [batch_size, patch_size, patch_size, num_channels],
                Batch of patches from data domain x.
        Output:
            A - float, array of [batch_size, patch_size^2, patch_size^2], Affinity matrix
    """
    _, h, w, c = x.shape
    x_1 = torch.unsqueeze(torch.reshape(x, [-1, h * w, c]), 2)
    x_2 = torch.unsqueeze(torch.reshape(x, [-1, h * w, c]), 1)
    A = torch.norm(x_1 - x_2, dim=-1)
    krnl_width, _ = torch.topk(A, k=A.shape[-1])
    krnl_width = torch.mean(krnl_width[:, :, (h * w) // 4], 1)
    krnl_width = torch.reshape(krnl_width, (-1, 1, 1))
    krnl_width = torch.where(torch.eq(krnl_width, torch.zeros_like(krnl_width)), torch.ones_like(krnl_width), krnl_width)
    A = torch.exp(-(torch.div(A, krnl_width) ** 2))
    return A
//  [python]
def affinity(x):
    """
    Compute the affinity matrices of the patches of contained in a batch.
    It first computes the distances between the datapoints within a patch.
    Then it finds the suitable kernel width for each patch.
    Finally, applies the RBF.
        Input:
            x - float, array of [batch_size, patch_size, patch_size, num_channels],
                Batch of patches from data domain x.
        Output:
            A - float, array of [batch_size, patch_size^2, patch_size^2], Affinity matrix
    """
    _, h, w, c = x.shape
    x_1 = torch.unsqueeze(torch.reshape(x, [-1, h * w, c]), 2)
    x_2 = torch.unsqueeze(torch.reshape(x, [-1, h * w, c]), 1)
    A = torch.linalg.norm(x_1 - x_2, dim=-1)
    krnl_width, _ = torch.topk(A, k=A.shape[-1])
    krnl_width = torch.mean(krnl_width[:, :, (h * w) // 4], 1)
    krnl_width = torch.reshape(krnl_width, (-1, 1, 1))
    krnl_width = torch.where(torch.eq(krnl_width, torch.zeros_like(krnl_width)), torch.ones_like(krnl_width), krnl_width)
    A = torch.exp(-(torch.div(A, krnl_width) ** 2))
    return A
//  [python]
def norm(x, ord="fro", axis=None, keepdims=False):
    return torch.norm(x, p=ord, dim=axis, keepdim=keepdims)
//  [python]
def norm(x, ord="fro", axis=None, keepdims=False):
    return torch.linalg.norm(x, ord=ord, dim=axis, keepdim=keepdims)
//  [python]
def test_keep_batch_order(self):
        embeddings = self.create_embedding_with_args(self.default_args)
        embedding_names = embeddings.get_names()

        sentences_1 = [Sentence("First sentence"), Sentence("This is second sentence")]
        sentences_2 = [Sentence("This is second sentence"), Sentence("First sentence")]

        embeddings.embed(sentences_1)
        embeddings.embed(sentences_2)

        assert sentences_1[0].to_original_text() == "First sentence"
        assert sentences_1[1].to_original_text() == "This is second sentence"

        if self.is_document_embedding:
            assert (
                torch.norm(
                    sentences_1[0].get_embedding(embedding_names) - sentences_2[1].get_embedding(embedding_names)
                )
                == 0.0
            )
            assert (
                torch.norm(
                    sentences_1[1].get_embedding(embedding_names) - sentences_2[0].get_embedding(embedding_names)
                )
                == 0.0
            )
        if self.is_token_embedding:
            for i in range(len(sentences_1[0])):
                assert (
                    torch.norm(
                        sentences_1[0][i].get_embedding(embedding_names)
                        - sentences_2[1][i].get_embedding(embedding_names)
                    )
                    == 0.0
                )
            for i in range(len(sentences_1[1])):
                assert (
                    torch.norm(
                        sentences_1[1][i].get_embedding(embedding_names)
                        - sentences_2[0][i].get_embedding(embedding_names)
                    )
                    == 0.0
                )
        del embeddings
//  [python]
def test_keep_batch_order(self):
        embeddings = self.create_embedding_with_args(self.default_args)
        embedding_names = embeddings.get_names()

        sentences_1 = [Sentence("First sentence"), Sentence("This is second sentence")]
        sentences_2 = [Sentence("This is second sentence"), Sentence("First sentence")]

        embeddings.embed(sentences_1)
        embeddings.embed(sentences_2)

        assert sentences_1[0].to_original_text() == "First sentence"
        assert sentences_1[1].to_original_text() == "This is second sentence"

        if self.is_document_embedding:
            assert (
                torch.linalg.norm(
                    sentences_1[0].get_embedding(embedding_names) - sentences_2[1].get_embedding(embedding_names)
                )
                == 0.0
            )
            assert (
                torch.linalg.norm(
                    sentences_1[1].get_embedding(embedding_names) - sentences_2[0].get_embedding(embedding_names)
                )
                == 0.0
            )
        if self.is_token_embedding:
            for i in range(len(sentences_1[0])):
                assert (
                    torch.linalg.norm(
                        sentences_1[0][i].get_embedding(embedding_names)
                        - sentences_2[1][i].get_embedding(embedding_names)
                    )
                    == 0.0
                )
            for i in range(len(sentences_1[1])):
                assert (
                    torch.linalg.norm(
                        sentences_1[1][i].get_embedding(embedding_names)
                        - sentences_2[0][i].get_embedding(embedding_names)
                    )
                    == 0.0
                )
        del embeddings
//  [python]
def forward(self, M):
        for i in range(self.max_iter):
            self.L = self.svd_decom(1/self.mu, M - self.S + (1/self.mu)*self.Y)
            self.S = self.shrinkage(self.gamma/self.mu, M - self.L + (1/self.mu)*self.Y)
            Z = M - self.L - self.S
            self.Y = self.Y + self.mu * Z
            
            error = torch.norm(Z, p='fro')/self.M_norm
            
            if i % 10 == 0:
                print('iter:', i, 'error:', error, 'rank:', torch.matrix_rank(self.L).item())
                
        return self.L, self.S
//  [python]
def forward(self, M):
        for i in range(self.max_iter):
            self.L = self.svd_decom(1/self.mu, M - self.S + (1/self.mu)*self.Y)
            self.S = self.shrinkage(self.gamma/self.mu, M - self.L + (1/self.mu)*self.Y)
            Z = M - self.L - self.S
            self.Y = self.Y + self.mu * Z
            
            error = torch.linalg.norm(Z, ord='fro')/self.M_norm
            
            if i % 10 == 0:
                print('iter:', i, 'error:', error, 'rank:', torch.matrix_rank(self.L).item())
                
        return self.L, self.S
//  [python]
def __init__(self, M, gamma=None, mu=None, tol=None, max_iter=None):
        super(robust_PCA, self).__init__()
        self.num_pixel, self.band = M.shape
        self.M_norm = torch.norm(M, p='fro')
        self.L = torch.zeros([self.num_pixel, self.band])
        self.S = torch.zeros([self.num_pixel, self.band])
        self.Y = torch.zeros([self.num_pixel, self.band])
        
        ######  Define Hyper-parameteres #####
        if gamma:
            self.gamma = gamma
        else:
            self.gamma = 1 / np.sqrt(max(self.num_pixel, self.band))
            
        if mu:
            self.mu = mu
        else:
            self.mu = 10 * self.gamma
            
        if tol:
            self.tol = tol
        else:
            self.tol = 1e-6
        
        if max_iter:
            self.max_iter = max_iter
        else:
            self.max_iter = 1000
//  [python]
def __init__(self, M, gamma=None, mu=None, tol=None, max_iter=None):
        super(robust_PCA, self).__init__()
        self.num_pixel, self.band = M.shape
        self.M_norm = torch.linalg.norm(M, ord='fro')
        self.L = torch.zeros([self.num_pixel, self.band])
        self.S = torch.zeros([self.num_pixel, self.band])
        self.Y = torch.zeros([self.num_pixel, self.band])
        
        ######  Define Hyper-parameteres #####
        if gamma:
            self.gamma = gamma
        else:
            self.gamma = 1 / np.sqrt(max(self.num_pixel, self.band))
            
        if mu:
            self.mu = mu
        else:
            self.mu = 10 * self.gamma
            
        if tol:
            self.tol = tol
        else:
            self.tol = 1e-6
        
        if max_iter:
            self.max_iter = max_iter
        else:
            self.max_iter = 1000
//  [python]
def norm1(a):
        return torch.norm(a, p=1, dim=1)
//  [python]
def norm1(a):
        return torch.linalg.norm(a, ord=1, dim=1)
//  [python]
def momentum_update(old_value, new_value, momentum, debug=False):
    update = momentum * old_value + (1 - momentum) * new_value
    if debug:
        print(
            "# old prot: {:.3f} x |{:.3f}|, new val: {:.3f} x |{:.3f}|, result= |{:.3f}|".format(
                momentum,
                torch.norm(old_value, p=2),
                (1 - momentum),
                torch.norm(new_value, p=2),
                torch.norm(update, p=2),
            )
        )
    return update
//  [python]
def momentum_update(old_value, new_value, momentum, debug=False):
    update = momentum * old_value + (1 - momentum) * new_value
    if debug:
        print(
            "# old prot: {:.3f} x |{:.3f}|, new val: {:.3f} x |{:.3f}|, result= |{:.3f}|".format(
                momentum,
                torch.linalg.norm(old_value, p=2),
                (1 - momentum),
                torch.linalg.norm(new_value, p=2),
                torch.linalg.norm(update, p=2),
            )
        )
    return update
//  [python]
def flsw_temperature_scheduler_builder(beta, gamma, eps=1e-4, *args):
    '''
    adapted from arXiv:1911.07471
    '''

    def flsw_temperature_scheduler(logits_S, logits_T, base_temperature):
        v = logits_S.detach()
        t = logits_T.detach()
        with torch.no_grad():
            v = v / (torch.norm(v, dim=-1, keepdim=True) + eps)
            t = t / (torch.norm(t, dim=-1, keepdim=True) + eps)
            w = torch.pow((1 - (v * t).sum(dim=-1)), gamma)
            tau = base_temperature + (w.mean() - w) * beta
        return tau

    return flsw_temperature_scheduler
//  [python]
def flsw_temperature_scheduler_builder(beta, gamma, eps=1e-4, *args):
    '''
    adapted from arXiv:1911.07471
    '''

    def flsw_temperature_scheduler(logits_S, logits_T, base_temperature):
        v = logits_S.detach()
        t = logits_T.detach()
        with torch.no_grad():
            v = v / (torch.linalg.norm(v, dim=-1, keepdim=True) + eps)
            t = t / (torch.linalg.norm(t, dim=-1, keepdim=True) + eps)
            w = torch.pow((1 - (v * t).sum(dim=-1)), gamma)
            tau = base_temperature + (w.mean() - w) * beta
        return tau

    return flsw_temperature_scheduler
//  [python]
def mini_batch_neigh_recon(self, h1, h0, input_id,
                               neighbor_dict, id_mapping):
        """Computing the target neighbor distribution and 
        reconstructed neighbor distribution using mini_batch of the data
        and neighbor sampling.
        """
        gen_neighs, tar_neighs = [], []
        
        sampled_embeddings_list, mask_len_list = \
                                        self.sample_neighbors(input_id,
                                                              neighbor_dict,
                                                              id_mapping,
                                                              h0)
        for index, neighbor_embeddings in enumerate(sampled_embeddings_list):
            # Generating h^k_v, reparameterization trick
            # the center node embeddings start from first row
            # in the h1 embedding matrix
            mean = h1[index].repeat(self.sample_size, 1)
            mean = self.mlp_mean(mean)
            sigma = h1[index].repeat(self.sample_size, 1)
            sigma = self.mlp_sigma(sigma)
            std_z = self.m_minibatch.sample().to(self.device)
            var = mean + sigma.exp() * std_z
            nhij = self.mlp_gen(var)
            
            generated_neighbors = nhij
            sum_neighbor_norm = 0
            
            for _, generated_neighbor in enumerate(generated_neighbors):
                sum_neighbor_norm += \
                    torch.norm(generated_neighbor) / math.sqrt(self.out_dim)
            generated_neighbors = \
                torch.unsqueeze(generated_neighbors, dim=0).to(self.device)
            target_neighbors = \
                torch.unsqueeze(torch.FloatTensor(neighbor_embeddings), 
                                dim=0).to(self.device)
            
            gen_neighs.append(generated_neighbors)
            tar_neighs.append(target_neighbors)
        
        # the information needed for loss computation
        recon_info = [gen_neighs, tar_neighs, mask_len_list]

        return recon_info
//  [python]
def mini_batch_neigh_recon(self, h1, h0, input_id,
                               neighbor_dict, id_mapping):
        """Computing the target neighbor distribution and 
        reconstructed neighbor distribution using mini_batch of the data
        and neighbor sampling.
        """
        gen_neighs, tar_neighs = [], []
        
        sampled_embeddings_list, mask_len_list = \
                                        self.sample_neighbors(input_id,
                                                              neighbor_dict,
                                                              id_mapping,
                                                              h0)
        for index, neighbor_embeddings in enumerate(sampled_embeddings_list):
            # Generating h^k_v, reparameterization trick
            # the center node embeddings start from first row
            # in the h1 embedding matrix
            mean = h1[index].repeat(self.sample_size, 1)
            mean = self.mlp_mean(mean)
            sigma = h1[index].repeat(self.sample_size, 1)
            sigma = self.mlp_sigma(sigma)
            std_z = self.m_minibatch.sample().to(self.device)
            var = mean + sigma.exp() * std_z
            nhij = self.mlp_gen(var)
            
            generated_neighbors = nhij
            sum_neighbor_norm = 0
            
            for _, generated_neighbor in enumerate(generated_neighbors):
                sum_neighbor_norm += \
                    torch.linalg.norm(generated_neighbor) / math.sqrt(self.out_dim)
            generated_neighbors = \
                torch.unsqueeze(generated_neighbors, dim=0).to(self.device)
            target_neighbors = \
                torch.unsqueeze(torch.FloatTensor(neighbor_embeddings), 
                                dim=0).to(self.device)
            
            gen_neighs.append(generated_neighbors)
            tar_neighs.append(target_neighbors)
        
        # the information needed for loss computation
        recon_info = [gen_neighs, tar_neighs, mask_len_list]

        return recon_info
//  [python]
def compute_beamforming_vector(target_scm: torch.Tensor, noise_scm: torch.Tensor):
        noise_scm_t = noise_scm.permute(0, 3, 1, 2)
        noise_scm_t = condition_scm(noise_scm_t, 1e-06)
        e_val, e_vec = generalized_eigenvalue_decomposition(target_scm.permute(0, 3, 1, 2), noise_scm_t)
        bf_vect = e_vec[..., -1]
        bf_vect /= torch.norm(bf_vect, dim=-1, keepdim=True)
        bf_vect = bf_vect.squeeze(-1).transpose(-1, -2)
        return bf_vect
//  [python]
def compute_beamforming_vector(target_scm: torch.Tensor, noise_scm: torch.Tensor):
        noise_scm_t = noise_scm.permute(0, 3, 1, 2)
        noise_scm_t = condition_scm(noise_scm_t, 1e-06)
        e_val, e_vec = generalized_eigenvalue_decomposition(target_scm.permute(0, 3, 1, 2), noise_scm_t)
        bf_vect = e_vec[..., -1]
        bf_vect /= torch.linalg.norm(bf_vect, dim=-1, keepdim=True)
        bf_vect = bf_vect.squeeze(-1).transpose(-1, -2)
        return bf_vect
//  [python]
def batch_matrix_norm(matrix, norm_order=2):
    """Normalize a matrix according to `norm_order`

    Args:
        matrix (torch.Tensor): Expected shape [batch, *]
        norm_order (int): Norm order.

    Returns:
        torch.Tensor, normed matrix of shape [batch]
    """
    keep_batch = list(range(1, matrix.ndim))
    return torch.norm(matrix, p=norm_order, dim=keep_batch) ** norm_order
//  [python]
def batch_matrix_norm(matrix, norm_order=2):
    """Normalize a matrix according to `norm_order`

    Args:
        matrix (torch.Tensor): Expected shape [batch, *]
        norm_order (int): Norm order.

    Returns:
        torch.Tensor, normed matrix of shape [batch]
    """
    keep_batch = list(range(1, matrix.ndim))
    return torch.linalg.norm(matrix, ord=norm_order, dim=keep_batch) ** norm_order