jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def mini_ksvd(M, s, batch_maker, n_batches, n_iter, max_time=None):
    # Inspired by https://csaws.cs.technion.ac.il/~ronrubin/Publications/KSVD-OMP-v2.pdf
    assert n_iter >= 1
    k, _ = M.shape

    start = time.time()
    last_error = 10**10
    for iter in range(n_iter):
        error = 0
        new_Ms = torch.zeros_like(M)
        n = 0
        for j in tqdm(range(n_batches)):
            X = batch_maker(j)
            if X is None:
                break
            new_M = M.clone()
            # Sparse coding
            ids, m = omp(X, new_M, s)
            # Find new dictionary
            for j in range(k):
                mask = ids == j
                I = torch.any(mask, dim=1)
                if not torch.any(I):
                    continue
                E = X[I] - (m[I].unsqueeze(1) @ new_M[ids[I]]).squeeze(1) + torch.outer(m[mask], new_M[j])

                # Could use the cheaper alternative decomposition in Approximate K-SVD here
                U, Sigma, V = torch.svd_lowrank(E, q=1)
                new_M[j, :] = V[:, 0]
                m[mask] = Sigma[0] * U[:, 0]

            # "Gradient" update
            # M = (1-lr)*M + lr*new_M
            # M = new_M
            new_Ms += new_M
            n += 1

            SM = (m.unsqueeze(1) @ M[ids]).squeeze(1)  # SM = S @ M
            error += torch.linalg.norm(SM - X) ** 2

        M = new_Ms / n

        error = error**0.5
        if error < 1e-4:
            print("K-SVD: Stopping early because error is near 0.")
            break

        if error > last_error:
            print(f"K-SVD: Stopping early because error is growing. last_error={last_error.item()}.")
            break
        last_error = error

        if max_time is not None and time.time() - start > max_time:
            print("K-SVD: Stopping early because ran out of time.")
            break

    print(f"K-SVD: error at {iter=}:", error.item())

    return M
//  [python]
def preprocess(self, hidden_states):
        hidden_states = hidden_states.permute(0, 2, 1).contiguous()
        hidden_states = hidden_states.view(-1, hidden_states.shape[-1])

        if hidden_states.shape[-1] == self.codebook_width:
            prenorm = torch.norm(hidden_states - torch.mean(hidden_states)) / np.sqrt(np.prod(hidden_states.shape))
        elif hidden_states.shape[-1] == 2 * self.codebook_width:
            x1, x2 = hidden_states[..., : self.codebook_width], hidden_states[..., self.codebook_width :]
            prenorm = (torch.norm(x1 - torch.mean(x1)) / np.sqrt(np.prod(x1.shape))) + (
                torch.norm(x2 - torch.mean(x2)) / np.sqrt(np.prod(x2.shape))
            )

            # Normalise
            hidden_states = x1 + x2

        return hidden_states, prenorm
//  [python]
def preprocess(self, hidden_states):
        hidden_states = hidden_states.permute(0, 2, 1).contiguous()
        hidden_states = hidden_states.view(-1, hidden_states.shape[-1])

        if hidden_states.shape[-1] == self.codebook_width:
            prenorm = torch.linalg.norm(hidden_states - torch.mean(hidden_states)) / np.sqrt(np.prod(hidden_states.shape))
        elif hidden_states.shape[-1] == 2 * self.codebook_width:
            x1, x2 = hidden_states[..., : self.codebook_width], hidden_states[..., self.codebook_width :]
            prenorm = (torch.linalg.norm(x1 - torch.mean(x1)) / np.sqrt(np.prod(x1.shape))) + (
                torch.linalg.norm(x2 - torch.mean(x2)) / np.sqrt(np.prod(x2.shape))
            )

            # Normalise
            hidden_states = x1 + x2

        return hidden_states, prenorm
//  [python]
def forward(self, hidden_states, update_codebook=True):
        samples, _, seq_len = hidden_states.shape

        # Preprocess
        hidden_states, prenorm = self.preprocess(hidden_states)

        # Init codebook if not inited
        if update_codebook and not self.init:
            self.init_codebook(hidden_states)

        # Quantise and dequantise through bottleneck
        music_tokens, fit = self.quantise(hidden_states)
        dequantised_states = self.dequantise(music_tokens)

        # Update embeddings
        if update_codebook:
            update_metrics = self.update_codebook(hidden_states, music_tokens)
        else:
            update_metrics = {}

        # Loss
        commit_loss = torch.norm(dequantised_states.detach() - hidden_states) ** 2 / np.prod(hidden_states.shape)

        # Passthrough
        dequantised_states = hidden_states + (dequantised_states - hidden_states).detach()

        # Postprocess
        music_tokens, dequantised_states = self.postprocess(music_tokens, dequantised_states, (samples, seq_len))
        return music_tokens, dequantised_states, commit_loss, dict(fit=fit, pn=prenorm, **update_metrics)
//  [python]
def forward(self, hidden_states, update_codebook=True):
        samples, _, seq_len = hidden_states.shape

        # Preprocess
        hidden_states, prenorm = self.preprocess(hidden_states)

        # Init codebook if not inited
        if update_codebook and not self.init:
            self.init_codebook(hidden_states)

        # Quantise and dequantise through bottleneck
        music_tokens, fit = self.quantise(hidden_states)
        dequantised_states = self.dequantise(music_tokens)

        # Update embeddings
        if update_codebook:
            update_metrics = self.update_codebook(hidden_states, music_tokens)
        else:
            update_metrics = {}

        # Loss
        commit_loss = torch.linalg.norm(dequantised_states.detach() - hidden_states) ** 2 / np.prod(hidden_states.shape)

        # Passthrough
        dequantised_states = hidden_states + (dequantised_states - hidden_states).detach()

        # Postprocess
        music_tokens, dequantised_states = self.postprocess(music_tokens, dequantised_states, (samples, seq_len))
        return music_tokens, dequantised_states, commit_loss, dict(fit=fit, pn=prenorm, **update_metrics)
//  [python]
def RBF(x):
    
    d_total_t = torch.tensor(d_total, device='cuda').unsqueeze(1)
    
    w_t = torch.tensor(w,  device='cuda')

    x_l = x.unsqueeze(0).repeat(len(d_total_t), 1, 1) 

    R2 = torch.norm(d_total_t - x_l, dim=2)
    y = torch.mm(torch.exp(-gama*R2.T), w_t)
    return y
//  [python]
def RBF(x):
    
    d_total_t = torch.tensor(d_total, device='cuda').unsqueeze(1)
    
    w_t = torch.tensor(w,  device='cuda')

    x_l = x.unsqueeze(0).repeat(len(d_total_t), 1, 1) 

    R2 = torch.linalg.norm(d_total_t - x_l, dim=2)
    y = torch.mm(torch.exp(-gama*R2.T), w_t)
    return y
//  [python]
def test_pspace_ekfac_vs_kfac():
    """
    Check that EKFAC matrix is closer to block diag one in the
    sense of the Frobenius norm
    """
    eps = 1e-4
    for get_task in [get_conv1d_task, get_fullyconnect_task, get_conv_task]:
        loader, lc, parameters, model, function, n_output = get_task()
        model.train()
        generator = Jacobian(
            layer_collection=lc, model=model, function=function, n_output=n_output
        )

        M_kfac = PMatKFAC(generator=generator, examples=loader)
        M_ekfac = PMatEKFAC(generator=generator, examples=loader)
        M_blockdiag = PMatBlockDiag(generator=generator, examples=loader)

        # here KFAC and EKFAC should be the same
        for split in [True, False]:
            diff = M_kfac.get_dense_tensor(
                split_weight_bias=split
            ) - M_ekfac.get_dense_tensor(split_weight_bias=split)
            assert torch.norm(diff) < eps

        # now we compute the exact diagonal:
        M_ekfac.update_diag(loader)
        assert torch.norm(
            M_kfac.get_dense_tensor() - M_blockdiag.get_dense_tensor()
        ) > torch.norm(M_ekfac.get_dense_tensor() - M_blockdiag.get_dense_tensor())
//  [python]
def test_pspace_ekfac_vs_kfac():
    """
    Check that EKFAC matrix is closer to block diag one in the
    sense of the Frobenius norm
    """
    eps = 1e-4
    for get_task in [get_conv1d_task, get_fullyconnect_task, get_conv_task]:
        loader, lc, parameters, model, function, n_output = get_task()
        model.train()
        generator = Jacobian(
            layer_collection=lc, model=model, function=function, n_output=n_output
        )

        M_kfac = PMatKFAC(generator=generator, examples=loader)
        M_ekfac = PMatEKFAC(generator=generator, examples=loader)
        M_blockdiag = PMatBlockDiag(generator=generator, examples=loader)

        # here KFAC and EKFAC should be the same
        for split in [True, False]:
            diff = M_kfac.get_dense_tensor(
                split_weight_bias=split
            ) - M_ekfac.get_dense_tensor(split_weight_bias=split)
            assert torch.linalg.norm(diff) < eps

        # now we compute the exact diagonal:
        M_ekfac.update_diag(loader)
        assert torch.linalg.norm(
            M_kfac.get_dense_tensor() - M_blockdiag.get_dense_tensor()
        ) > torch.linalg.norm(M_ekfac.get_dense_tensor() - M_blockdiag.get_dense_tensor())
//  [python]
def GMRES(Avp,              # Linear operator
          b,                # RHS of the linear system
          x0=None,          # initial guess, tuple has the same shape as b
          max_iter=None,    # maximum number of iterations
          tol=1e-6,         # relative tolerance
          atol=1e-6,        # absolute tolerance
          track=False       # If True, keep a history of the relative residual error
          ):
    bnorm = torch.norm(b)
    _check_nan(b, 'RHS of the system is Nan')
    if max_iter == 0 or bnorm < 1e-8:
        return b, (0, 0)

    if max_iter is None:
        max_iter = b.shape[0]

    if x0 is None:
        x0 = torch.zeros_like(b)
        r0 = b
    else:
        r0 = b - Avp(x0)

    new_v, rnorm = _safe_normalize(r0)
    # initial guess residual
    beta = torch.zeros(max_iter + 1, device=b.device)
    beta[0] = rnorm

    err_history = []
    if track:
        err_history.append((rnorm / bnorm).item())

    V = []
    V.append(new_v)
    H = torch.zeros((max_iter + 1, max_iter + 1), device=b.device)
    cs = torch.zeros(max_iter, device=b.device)  # cosine values at each step
    ss = torch.zeros(max_iter, device=b.device)  # sine values at each step

    for j in range(max_iter):
        p = Avp(V[j])
        new_v = arnoldi(p, V, H, j + 1)  # Arnoldi iteration to get the j+1 th ba
        # sis
        V.append(new_v)

        H, cs, ss = apply_given_rotation(H, cs, ss, j)
        beta[j + 1] = ss[j] * beta[j]
        beta[j] = cs[j] * beta[j]
        residual = torch.abs(beta[j + 1])
        if track:
            err_history.append((residual / bnorm).item())
        if residual < tol * bnorm or residual < atol:
            # print(f'\nGMRES iterations: {j}')
            break
        # print(f'{j}the gmres iteration:  residual {residual/bnorm}')
    if j == max_iter - 1:
        print(f'\nMax number of iterations: {max_iter}')
    y, _ = torch.triangular_solve(beta[0:j + 1].unsqueeze(-1), H[0:j + 1, 0:j + 1])  # j x j
    V = torch.stack(V[:-1], dim=0)
    sol = x0 + V.T @ y.squeeze(-1)
    return sol, (j, err_history)
//  [python]
def GMRES(Avp,              # Linear operator
          b,                # RHS of the linear system
          x0=None,          # initial guess, tuple has the same shape as b
          max_iter=None,    # maximum number of iterations
          tol=1e-6,         # relative tolerance
          atol=1e-6,         # absolute tolerance
          track=False       # If True, keep a history of the relative residual error
          ):
    bnorm = torch.linalg.norm(b)
    _check_nan(b, 'RHS of the system is Nan')
    if max_iter == 0 or bnorm < 1e-8:
        return b, (0, 0)

    if max_iter is None:
        max_iter = b.shape[0]

    if x0 is None:
        x0 = torch.zeros_like(b)
        r0 = b
    else:
        r0 = b - Avp(x0)

    new_v, rnorm = _safe_normalize(r0)
    # initial guess residual
    beta = torch.zeros(max_iter + 1, device=b.device)
    beta[0] = rnorm

    err_history = []
    if track:
        err_history.append((rnorm / bnorm).item())

    V = []
    V.append(new_v)
    H = torch.zeros((max_iter + 1, max_iter + 1), device=b.device)
    cs = torch.zeros(max_iter, device=b.device)  # cosine values at each step
    ss = torch.zeros(max_iter, device=b.device)  # sine values at each step

    for j in range(max_iter):
        p = Avp(V[j])
        new_v = arnoldi(p, V, H, j + 1)  # Arnoldi iteration to get the j+1 th ba
        # sis
        V.append(new_v)

        H, cs, ss = apply_given_rotation(H, cs, ss, j)
        beta[j + 1] = ss[j] * beta[j]
        beta[j] = cs[j] * beta[j]
        residual = torch.abs(beta[j + 1])
        if track:
            err_history.append((residual / bnorm).item())
        if residual < tol * bnorm or residual < atol:
            # print(f'\nGMRES iterations: {j}')
            break
        # print(f'{j}the gmres iteration:  residual {residual/bnorm}')
    if j == max_iter - 1:
        print(f'\nMax number of iterations: {max_iter}')
    y, _ = torch.triangular_solve(beta[0:j + 1].unsqueeze(-1), H[0:j + 1, 0:j + 1])  # j x j
    V = torch.stack(V[:-1], dim=0)
    sol = x0 + V.T @ y.squeeze(-1)
    return sol, (j, err_history)
//  [python]
def cosine_similarity(
    emb_a: torch.Tensor, emb_b: torch.Tensor, eps: Optional[torch.Tensor] = None
) -> torch.Tensor:
    """Calculate cosine similarities of the given two set of tensors.

    The output is an N by N matrix where N is the number of feature vectors.

    Args:
        emb_a (torch.Tensor): Matrix containing speaker representation vectors. (N x embedding_dim)
        emb_b (torch.Tensor): Matrix containing speaker representation vectors. (N x embedding_dim)
        eps (torch.Tensor): Small value to avoid division by zero.

    Returns:
        torch.Tensor: Matrix containing cosine similarities. (N x N)
    """
    if emb_a.shape[0] == 1 or emb_b.shape[0] == 1:
        raise ValueError(
            "Number of feature vectors should be greater than 1 but got"
            f" {emb_a.shape} and {emb_b.shape}"
        )

    if eps is None:
        eps = torch.tensor(3.5e-4).to(emb_a.device)
    else:
        eps = eps.to(emb_a.device)

    a_norm = emb_a / (torch.norm(emb_a, dim=1).unsqueeze(1) + eps)
    b_norm = emb_b / (torch.norm(emb_b, dim=1).unsqueeze(1) + eps)

    cosine_similarity = torch.mm(a_norm, b_norm.transpose(0, 1)).fill_diagonal_(1)

    return cosine_similarity
//  [python]
def cosine_similarity(
    emb_a: torch.Tensor, emb_b: torch.Tensor, eps: Optional[torch.Tensor] = None
) -> torch.Tensor:
    """Calculate cosine similarities of the given two set of tensors.

    The output is an N by N matrix where N is the number of feature vectors.

    Args:
        emb_a (torch.Tensor): Matrix containing speaker representation vectors. (N x embedding_dim)
        emb_b (torch.Tensor): Matrix containing speaker representation vectors. (N x embedding_dim)
        eps (torch.Tensor): Small value to avoid division by zero.

    Returns:
        torch.Tensor: Matrix containing cosine similarities. (N x N)
    """
    if emb_a.shape[0] == 1 or emb_b.shape[0] == 1:
        raise ValueError(
            "Number of feature vectors should be greater than 1 but got"
            f" {emb_a.shape} and {emb_b.shape}"
        )

    if eps is None:
        eps = torch.tensor(3.5e-4).to(emb_a.device)
    else:
        eps = eps.to(emb_a.device)

    a_norm = emb_a / (torch.linalg.norm(emb_a, dim=1, keepdims=True) + eps)
    b_norm = emb_b / (torch.linalg.norm(emb_b, dim=1, keepdims=True) + eps)

    cosine_similarity = torch.mm(a_norm, b_norm.transpose(0, 1)).fill_diagonal_(1)

    return cosine_similarity
//  [python]
def RBF(x):
    d_total_t = torch.from_numpy(d_total).unsqueeze(1).cuda()
    w_t = torch.from_numpy(w).cuda()
    x_l = x.unsqueeze(0).repeat(len(d_total_t), 1, 1) # 创立一个足够大的x矩阵
    # 得到R2大矩阵
    R2 = torch.norm(d_total_t - x_l, dim=2)
    y = torch.mm(torch.exp(-gama*R2.T), w_t)
    return y
//  [python]
def RBF(x):
    d_total_t = torch.from_numpy(d_total).unsqueeze(1).cuda()
    w_t = torch.from_numpy(w).cuda()
    x_l = x.unsqueeze(0).repeat(len(d_total_t), 1, 1) # 创立一个足够大的x矩阵
    # 得到R2大矩阵
    R2 = torch.linalg.norm(d_total_t - x_l, dim=2)
    y = torch.mm(torch.exp(-gama*R2.T), w_t)
    return y
//  [python]
def train_data(Nb, Nf):
    xu = np.hstack([np.random.rand(int(Nb/4),1)*2-1,  np.ones([int(Nb/4),1])])
    xd = np.hstack([np.random.rand(int(Nb/4),1)*2-1,  -np.ones([int(Nb/4),1])])
    xl = np.hstack([-np.ones([int(Nb/4),1]), np.random.rand(int(Nb/4),1)*2-1])
    xr = np.hstack([np.ones([int(Nb/4),1]), np.random.rand(int(Nb/4),1)*2-1])
    xcrack = np.hstack([-np.random.rand(int(Nb/4),1), np.zeros([int(Nb/4),1])])

    xc1 = np.array([[-1., 1.], [1., 1.], [1., -1.], [-1., -1.]]) 
    xc2 = np.array([[-1., 1.], [1., 1.], [1., -1.], [-1., -1.]])
    xc3 = np.array([[-1., 1.], [1., 1.], [1., -1.], [-1., -1.]])
    Xb = np.concatenate((xu, xd, xl, xr, xcrack, xc1, xc2, xc3)) 

    Xb = torch.tensor(Xb, device='cuda')
    
    Xf = torch.rand(Nf, 2)*2-1
    

    Xf1 = Xf[(Xf[:, 1]>0) & (torch.norm(Xf, dim=1)>=delta)] 
    Xf2 = Xf[(Xf[:, 1]<0) & (torch.norm(Xf, dim=1)>=delta)]  
    
    Xf1 = torch.tensor(Xf1, requires_grad=True, device='cuda')
    Xf2 = torch.tensor(Xf2,  requires_grad=True, device='cuda')
    
    return Xb, Xf1, Xf2
//  [python]
def train_data(Nb, Nf):
    xu = np.hstack([np.random.rand(int(Nb/4),1)*2-1,  np.ones([int(Nb/4),1])])
    xd = np.hstack([np.random.rand(int(Nb/4),1)*2-1,  -np.ones([int(Nb/4),1])])
    xl = np.hstack([-np.ones([int(Nb/4),1]), np.random.rand(int(Nb/4),1)*2-1])
    xr = np.hstack([np.ones([int(Nb/4),1]), np.random.rand(int(Nb/4),1)*2-1])
    xcrack = np.hstack([-np.random.rand(int(Nb/4),1), np.zeros([int(Nb/4),1])])

    xc1 = np.array([[-1., 1.], [1., 1.], [1., -1.], [-1., -1.]]) 
    xc2 = np.array([[-1., 1.], [1., 1.], [1., -1.], [-1., -1.]])
    xc3 = np.array([[-1., 1.], [1., 1.], [1., -1.], [-1., -1.]])
    Xb = np.concatenate((xu, xd, xl, xr, xcrack, xc1, xc2, xc3)) 

    Xb = torch.tensor(Xb, device='cuda')
    
    Xf = torch.rand(Nf, 2)*2-1
    

    Xf1 = Xf[(Xf[:, 1]>0) & (torch.linalg.norm(Xf, dim=1)>=delta)] 
    Xf2 = Xf[(Xf[:, 1]<0) & (torch.linalg.norm(Xf, dim=1)>=delta)]  
    
    Xf1 = torch.tensor(Xf1, requires_grad=True, device='cuda')
    Xf2 = torch.tensor(Xf2,  requires_grad=True, device='cuda')
    
    return Xb, Xf1, Xf2
//  [python]
def calculate_angle(self, v1, v2):
        unit_v1 = v1 / torch.norm(v1, dim=-1, keepdim=True)
        unit_v2 = v2 / torch.norm(v2, dim=-1, keepdim=True)
        dot_product = torch.sum(unit_v1 * unit_v2, dim=-1)
        angle = torch.acos(dot_product)
        return angle
//  [python]
def calculate_angle(self, v1, v2):
        unit_v1 = v1 / torch.linalg.norm(v1, dim=-1, keepdim=True)
        unit_v2 = v2 / torch.linalg.norm(v2, dim=-1, keepdim=True)
        dot_product = torch.sum(unit_v1 * unit_v2, dim=-1)
        angle = torch.acos(dot_product)
        return angle
//  [python]
def normalize_tensor(tensor, min_norm=1e-5, feat_dim=-1):
    ''' Normalizes the tensor.

    Args:
        tensor (tensor): tensor
        min_norm (float): minimum norm for numerical stability
        feat_dim (int): feature dimension in tensor (default: -1)
    '''
    norm_tensor = torch.clamp(torch.norm(tensor, dim=feat_dim, keepdim=True),
                              min=min_norm)
    normed_tensor = tensor / norm_tensor
    return normed_tensor
//  [python]
def normalize_tensor(tensor, min_norm=1e-5, feat_dim=-1):
    ''' Normalizes the tensor.

    Args:
        tensor (tensor): tensor
        min_norm (float): minimum norm for numerical stability
        feat_dim (int): feature dimension in tensor (default: -1)
    '''
    norm_tensor = torch.clamp(torch.linalg.norm(tensor, dim=feat_dim, keepdim=True),
                              min=min_norm)
    normed_tensor = tensor / norm_tensor
    return normed_tensor
//  [python]
def test_pspace_ekfac_vs_direct():
    """
    Check EKFAC basis operations against direct computation using
    get_dense_tensor
    """
    for get_task in [get_conv1d_task, get_fullyconnect_task, get_conv_task]:
        loader, lc, parameters, model, function, n_output = get_task()
        model.train()

        generator = Jacobian(
            layer_collection=lc, model=model, function=function, n_output=n_output
        )

        M_ekfac = PMatEKFAC(generator=generator, examples=loader)
        v = random_pvector(lc, device=device)

        # the second time we will have called update_diag
        for i in range(2):
            vTMv_direct = torch.dot(
                torch.mv(M_ekfac.get_dense_tensor(), v.get_flat_representation()),
                v.get_flat_representation(),
            )
            vTMv_ekfac = M_ekfac.vTMv(v)
            check_ratio(vTMv_direct, vTMv_ekfac)

            trace_ekfac = M_ekfac.trace()
            trace_direct = torch.trace(M_ekfac.get_dense_tensor())
            check_ratio(trace_direct, trace_ekfac)

            frob_ekfac = M_ekfac.frobenius_norm()
            frob_direct = torch.norm(M_ekfac.get_dense_tensor())
            check_ratio(frob_direct, frob_ekfac)

            mv_direct = torch.mv(
                M_ekfac.get_dense_tensor(), v.get_flat_representation()
            )
            mv_ekfac = M_ekfac.mv(v)
            check_tensors(mv_direct, mv_ekfac.get_flat_representation())

            # Test inverse
            regul = 1e-5
            M_inv = M_ekfac.inverse(regul=regul)
            v_back = M_inv.mv(mv_ekfac + regul * v)
            check_tensors(v.get_flat_representation(), v_back.get_flat_representation())

            # Test solve
            v_back = M_ekfac.solve(mv_ekfac + regul * v, regul=regul)
            check_tensors(v.get_flat_representation(), v_back.get_flat_representation())

            # Test rmul
            M_mul = 1.23 * M_ekfac
            check_tensors(1.23 * M_ekfac.get_dense_tensor(), M_mul.get_dense_tensor())

            M_ekfac.update_diag(loader)
//  [python]
def test_pspace_ekfac_vs_direct():
    """
    Check EKFAC basis operations against direct computation using
    get_dense_tensor
    """
    for get_task in [get_conv1d_task, get_fullyconnect_task, get_conv_task]:
        loader, lc, parameters, model, function, n_output = get_task()
        model.train()

        generator = Jacobian(
            layer_collection=lc, model=model, function=function, n_output=n_output
        )

        M_ekfac = PMatEKFAC(generator=generator, examples=loader)
        v = random_pvector(lc, device=device)

        # the second time we will have called update_diag
        for i in range(2):
            vTMv_direct = torch.dot(
                torch.mv(M_ekfac.get_dense_tensor(), v.get_flat_representation()),
                v.get_flat_representation(),
            )
            vTMv_ekfac = M_ekfac.vTMv(v)
            check_ratio(vTMv_direct, vTMv_ekfac)

            trace_ekfac = M_ekfac.trace()
            trace_direct = torch.trace(M_ekfac.get_dense_tensor())
            check_ratio(trace_direct, trace_ekfac)

            frob_ekfac = M_ekfac.frobenius_norm()
            frob_direct = torch.linalg.norm(M_ekfac.get_dense_tensor())
            check_ratio(frob_direct, frob_ekfac)

            mv_direct = torch.mv(
                M_ekfac.get_dense_tensor(), v.get_flat_representation()
            )
            mv_ekfac = M_ekfac.mv(v)
            check_tensors(mv_direct, mv_ekfac.get_flat_representation())

            # Test inverse
            regul = 1e-5
            M_inv = M_ekfac.inverse(regul=regul)
            v_back = M_inv.mv(mv_ekfac + regul * v)
            check_tensors(v.get_flat_representation(), v_back.get_flat_representation())

            # Test solve
            v_back = M_ekfac.solve(mv_ekfac + regul * v, regul=regul)
            check_tensors(v.get_flat_representation(), v_back.get_flat_representation())

            # Test rmul
            M_mul = 1.23 * M_ekfac
            check_tensors(1.23 * M_ekfac.get_dense_tensor(), M_mul.get_dense_tensor())

            M_ekfac.update_diag(loader)
//  [python]
def evaluate_fine(self, output_dict, data_dict):
        transform = data_dict['transform']
        ref_corr_points = output_dict['ref_corr_points']
        src_corr_points = output_dict['src_corr_points']
        src_corr_points = apply_transform(src_corr_points, transform)
        corr_distances = torch.norm(ref_corr_points - src_corr_points, dim=1)
        precision = torch.lt(corr_distances, self.acceptance_radius).float().mean()
        return precision
//  [python]
def evaluate_fine(self, output_dict, data_dict):
        transform = data_dict['transform']
        ref_corr_points = output_dict['ref_corr_points']
        src_corr_points = output_dict['src_corr_points']
        src_corr_points = apply_transform(src_corr_points, transform)
        corr_distances = torch.linalg.norm(ref_corr_points - src_corr_points, dim=1)
        precision = torch.lt(corr_distances, self.acceptance_radius).float().mean()
        return precision
//  [python]
def reduction_ops(self):
        a = torch.randn(4)
        b = torch.randn(4)
        c = torch.tensor(0.5)
        return len(
            torch.argmax(a),
            torch.argmin(a),
            torch.amax(a),
            torch.amin(a),
            torch.aminmax(a),
            torch.all(a),
            torch.any(a),
            torch.max(a),
            a.max(a),
            torch.max(a, 0),
            torch.min(a),
            a.min(a),
            torch.min(a, 0),
            torch.dist(a, b),
            torch.logsumexp(a, 0),
            torch.mean(a),
            torch.mean(a, 0),
            torch.nanmean(a),
            torch.median(a),
            torch.nanmedian(a),
            torch.mode(a),
            torch.norm(a),
            a.norm(2),
            torch.norm(a, dim=0),
            torch.norm(c, torch.tensor(2)),
            torch.nansum(a),
            torch.prod(a),
            torch.quantile(a, torch.tensor([0.25, 0.5, 0.75])),
            torch.quantile(a, 0.5),
            torch.nanquantile(a, torch.tensor([0.25, 0.5, 0.75])),
            torch.std(a),
            torch.std_mean(a),
            torch.sum(a),
            torch.unique(a),
            torch.unique_consecutive(a),
            torch.var(a),
            torch.var_mean(a),
            torch.count_nonzero(a),
        )
//  [python]
def reduction_ops(self):
        a = torch.randn(4)
        b = torch.randn(4)
        c = torch.tensor(0.5)
        return len(
            torch.argmax(a),
            torch.argmin(a),
            torch.amax(a),
            torch.amin(a),
            torch.aminmax(a),
            torch.all(a),
            torch.any(a),
            torch.max(a),
            a.max(a),
            torch.max(a, 0),
            torch.min(a),
            a.min(a),
            torch.min(a, 0),
            torch.dist(a, b),
            torch.logsumexp(a, 0),
            torch.mean(a),
            torch.mean(a, 0),
            torch.nanmean(a),
            torch.median(a),
            torch.nanmedian(a),
            torch.mode(a),
            torch.linalg.norm(a),
            a.norm(2),
            torch.linalg.norm(a, dim=0),
            torch.linalg.norm(c, torch.tensor(2)),
            torch.nansum(a),
            torch.prod(a),
            torch.quantile(a, torch.tensor([0.25, 0.5, 0.75])),
            torch.quantile(a, 0.5),
            torch.nanquantile(a, torch.tensor([0.25, 0.5, 0.75])),
            torch.std(a),
            torch.std_mean(a),
            torch.sum(a),
            torch.unique(a),
            torch.unique_consecutive(a),
            torch.var(a),
            torch.var_mean(a),
            torch.count_nonzero(a),
        )
//  [python]
def minimum_distortion_l2_phase(Y: pt.Tensor, ref: pt.Tensor) -> pt.Tensor:
    """
    This function computes the frequency-domain filter that minimizes
    the squared error to a reference signal. This is commonly used
    to solve the scale ambiguity in BSS.

    Parameters
    ----------
    Y: torch.Tensor (..., n_channels, n_frequencies, n_frames)
        The STFT data to project back on the reference signal
    ref: torch.Tensor (..., n_frequencies, n_frames)
        The reference signal
    """
    Y = Y.transpose(-3, -2)  # put Y in order (..., n_freq, n_chan, n_frames)

    num = pt.view_as_real(pt.mean(ref[..., None, :].conj() * Y, dim=-1))
    # only correct the phase
    c = pt.view_as_complex(num / pt.clamp(pt.norm(num, dim=-1, keepdim=True), min=1e-5))

    return (Y * c[..., :, :, None].conj()).transpose(-3, -2)
//  [python]
def minimum_distortion_l2_phase(Y: pt.Tensor, ref: pt.Tensor) -> pt.Tensor:
    """
    This function computes the frequency-domain filter that minimizes
    the squared error to a reference signal. This is commonly used
    to solve the scale ambiguity in BSS.

    Parameters
    ----------
    Y: torch.Tensor (..., n_channels, n_frequencies, n_frames)
        The STFT data to project back on the reference signal
    ref: torch.Tensor (..., n_frequencies, n_frames)
        The reference signal
    """
    Y = Y.transpose(-3, -2)  # put Y in order (..., n_freq, n_chan, n_frames)

    num = pt.view_as_real(pt.mean(ref[..., None, :].conj() * Y, dim=-1))
    # only correct the phase
    c = pt.view_as_complex(num / pt.clamp(pt.linalg.norm(num, dim=-1, keepdim=True), min=1e-5))

    return (Y * c[..., :, :, None].conj()).transpose(-3, -2)
//  [python]
def get_plane_distance(self, r):
        """Signed distance of r from each perpendicular bisector plane from 0->faces."""
        bgrid_mag = torch.norm(self.faces, dim=-1)
        bgrid_hat = self.faces / bgrid_mag[:,
                                None]  # outward normal of the perpendicular bisector planes
        return r @ bgrid_hat.T - 0.5 * bgrid_mag
//  [python]
def get_plane_distance(self, r):
        """Signed distance of r from each perpendicular bisector plane from 0->faces."""
        bgrid_mag = torch.linalg.norm(self.faces, axis=-1)
        bgrid_hat = self.faces / bgrid_mag[:,
                                None]  # outward normal of the perpendicular bisector planes
        return r @ bgrid_hat.T - 0.5 * bgrid_mag
//  [python]
def weld_points(points: torch.Tensor, tol: float) -> torch.Tensor:
    """Return unique points based on distance < tol from N x d tensor."""
    distances = torch.norm(points[:, None] - points[None, :], dim=-1)
    lowest_equivalent_index = torch.argmax(torch.where(distances < tol, 1, 0), dim=0)
    _, unique_indices = torch.unique(lowest_equivalent_index, return_inverse=True)
    return points[unique_indices]
//  [python]
def weld_points(points: torch.Tensor, tol: float) -> torch.Tensor:
    """Return unique points based on distance < tol from N x d tensor."""
    distances = torch.linalg.norm(points[:, None] - points[None, :], dim=-1)
    lowest_equivalent_index = torch.argmax(torch.where(distances < tol, 1, 0), dim=0)
    _, unique_indices = torch.unique(lowest_equivalent_index, return_inverse=True)
    return points[unique_indices]
//  [python]
def __init__(
        self,
        rot_mats: Optional[torch.Tensor] = None,
        quats: Optional[torch.Tensor] = None,
        normalize_quats: bool = True,
    ):
        """
        Args:
            rot_mats:
                A [*, 3, 3] rotation matrix tensor. Mutually exclusive with quats
            quats:
                A [*, 4] quaternion. Mutually exclusive with rot_mats. If normalize_quats is not True, must be a unit
                quaternion
            normalize_quats:
                If quats is specified, whether to normalize quats
        """
        if (rot_mats is None and quats is None) or (rot_mats is not None and quats is not None):
            raise ValueError("Exactly one input argument must be specified")

        if (rot_mats is not None and rot_mats.shape[-2:] != (3, 3)) or (quats is not None and quats.shape[-1] != 4):
            raise ValueError("Incorrectly shaped rotation matrix or quaternion")

        # Force full-precision
        if quats is not None:
            quats = quats.to(dtype=torch.float32)
        if rot_mats is not None:
            rot_mats = rot_mats.to(dtype=torch.float32)

        if quats is not None and normalize_quats:
            quats = quats / torch.norm(quats, dim=-1, keepdim=True)

        self._rot_mats = rot_mats
        self._quats = quats
//  [python]
def __init__(
        self,
        rot_mats: Optional[torch.Tensor] = None,
        quats: Optional[torch.Tensor] = None,
        normalize_quats: bool = True,
    ):
        """
        Args:
            rot_mats:
                A [*, 3, 3] rotation matrix tensor. Mutually exclusive with quats
            quats:
                A [*, 4] quaternion. Mutually exclusive with rot_mats. If normalize_quats is not True, must be a unit
                quaternion
            normalize_quats:
                If quats is specified, whether to normalize quats
        """
        if (rot_mats is None and quats is None) or (rot_mats is not None and quats is not None):
            raise ValueError("Exactly one input argument must be specified")

        if (rot_mats is not None and rot_mats.shape[-2:] != (3, 3)) or (quats is not None and quats.shape[-1] != 4):
            raise ValueError("Incorrectly shaped rotation matrix or quaternion")

        # Force full-precision
        if quats is not None:
            quats = quats.to(dtype=torch.float32)
        if rot_mats is not None:
            rot_mats = rot_mats.to(dtype=torch.float32)

        if quats is not None and normalize_quats:
            quats = quats / torch.linalg.norm(quats, dim=-1, keepdim=True)

        self._rot_mats = rot_mats
        self._quats = quats
//  [python]
def evaluate_registration(self, output_dict, data_dict):
        transform = data_dict['transform']
        est_transform = output_dict['estimated_transform']
        src_points = output_dict['src_points']

        rre, rte = isotropic_transform_error(transform, est_transform)

        realignment_transform = torch.matmul(torch.inverse(transform), est_transform)
        realigned_src_points_f = apply_transform(src_points, realignment_transform)
        rmse = torch.norm(realigned_src_points_f - src_points, dim=1).mean()
        recall = torch.lt(rmse, self.acceptance_rmse).float()

        return rre, rte, rmse, recall
//  [python]
def evaluate_registration(self, output_dict, data_dict):
        transform = data_dict['transform']
        est_transform = output_dict['estimated_transform']
        src_points = output_dict['src_points']

        rre, rte = isotropic_transform_error(transform, est_transform)

        realignment_transform = torch.matmul(torch.inverse(transform), est_transform)
        realigned_src_points_f = apply_transform(src_points, realignment_transform)
        rmse = torch.linalg.norm(realigned_src_points_f - src_points, dim=1).mean()
        recall = torch.lt(rmse, self.acceptance_rmse).float()

        return rre, rte, rmse, recall
//  [python]
def __init__(self,
        rot_mats: Optional[torch.Tensor] = None,
        quats: Optional[torch.Tensor] = None,
        normalize_quats: bool = True,
    ):
        """
            Args:
                rot_mats:
                    A [*, 3, 3] rotation matrix tensor. Mutually exclusive with
                    quats
                quats:
                    A [*, 4] quaternion. Mutually exclusive with rot_mats. If
                    normalize_quats is not True, must be a unit quaternion
                normalize_quats:
                    If quats is specified, whether to normalize quats
        """
        if((rot_mats is None and quats is None) or 
            (rot_mats is not None and quats is not None)):
            raise ValueError("Exactly one input argument must be specified")

        if((rot_mats is not None and rot_mats.shape[-2:] != (3, 3)) or 
            (quats is not None and quats.shape[-1] != 4)):
            raise ValueError(
                "Incorrectly shaped rotation matrix or quaternion"
            )

        # Force full-precision
        if(quats is not None):
            quats = quats.to(dtype=torch.float32)
        if(rot_mats is not None):
            rot_mats = rot_mats.to(dtype=torch.float32)

        if(quats is not None and normalize_quats):
            quats = quats / torch.norm(quats, dim=-1, keepdim=True)

        self._rot_mats = rot_mats
        self._quats = quats
//  [python]
def __init__(self,
        rot_mats: Optional[torch.Tensor] = None,
        quats: Optional[torch.Tensor] = None,
        normalize_quats: bool = True,
    ):
        """
            Args:
                rot_mats:
                    A [*, 3, 3] rotation matrix tensor. Mutually exclusive with
                    quats
                quats:
                    A [*, 4] quaternion. Mutually exclusive with rot_mats. If
                    normalize_quats is not True, must be a unit quaternion
                normalize_quats:
                    If quats is specified, whether to normalize quats
        """
        if((rot_mats is None and quats is None) or 
            (rot_mats is not None and quats is not None)):
            raise ValueError("Exactly one input argument must be specified")

        if((rot_mats is not None and rot_mats.shape[-2:] != (3, 3)) or 
            (quats is not None and quats.shape[-1] != 4)):
            raise ValueError(
                "Incorrectly shaped rotation matrix or quaternion"
            )

        # Force full-precision
        if(quats is not None):
            quats = quats.to(dtype=torch.float32)
        if(rot_mats is not None):
            rot_mats = rot_mats.to(dtype=torch.float32)

        if(quats is not None and normalize_quats):
            quats = quats / torch.linalg.norm(quats, dim=-1, keepdim=True)

        self._rot_mats = rot_mats
        self._quats = quats
//  [python]
def axis_angle_from_quat(quat, eps=1.0e-6):
    """Convert tensor of quaternions to tensor of axis-angles."""
    # Reference: https://github.com/facebookresearch/pytorch3d/blob/bee31c48d3d36a8ea268f9835663c52ff4a476ec/pytorch3d/transforms/rotation_conversions.py#L516-L544

    mag = torch.norm(quat[:, 1:4], dim=1)
    half_angle = torch.atan2(mag, quat[:, 0])
    angle = 2.0 * half_angle
    sin_half_angle_over_angle = torch.where(
        torch.abs(angle) > eps, torch.sin(half_angle) / angle, 1 / 2 - angle**2.0 / 48
    )
    axis_angle = quat[:, 1:4] / sin_half_angle_over_angle.unsqueeze(-1)

    return axis_angle
//  [python]
def axis_angle_from_quat(quat, eps=1.0e-6):
    """Convert tensor of quaternions to tensor of axis-angles."""
    # Reference: https://github.com/facebookresearch/pytorch3d/blob/bee31c48d3d36a8ea268f9835663c52ff4a476ec/pytorch3d/transforms/rotation_conversions.py#L516-L544

    mag = torch.linalg.norm(quat[:, 1:4], dim=1)
    half_angle = torch.atan2(mag, quat[:, 0])
    angle = 2.0 * half_angle
    sin_half_angle_over_angle = torch.where(
        torch.abs(angle) > eps, torch.sin(half_angle) / angle, 1 / 2 - angle**2.0 / 48
    )
    axis_angle = quat[:, 1:4] / sin_half_angle_over_angle.unsqueeze(-1)

    return axis_angle
//  [python]
def check_unit_quaternion(quaternion: torch.Tensor):
    def _impl(t_: torch.Tensor):
        QUANTERNION_EPS = LIE_PARAMS.get_eps("so3", "quat", t_.dtype)

        if t_.dtype != torch.float64:
            t_ = t_.double()

        if (t_.norm(dim=-1) - 1).abs().max().item() >= QUANTERNION_EPS:
            raise ValueError("Not unit quaternions.")

    if quaternion.shape[-1] != 4:
        raise ValueError(shape_err_msg("Quaternions", "(..., 4)", quaternion.shape))

    checks_base(quaternion, _impl)
//  [python]
def check_unit_quaternion(quaternion: torch.Tensor):
    def _impl(t_: torch.Tensor):
        QUANTERNION_EPS = LIE_PARAMS.get_eps("so3", "quat", t_.dtype)

        if t_.dtype != torch.float64:
            t_ = t_.double()

        if (torch.linalg.norm(t_, dim=-1) - 1).abs().max().item() >= QUANTERNION_EPS:
            raise ValueError("Not unit quaternions.")

    if quaternion.shape[-1] != 4:
        raise ValueError(shape_err_msg("Quaternions", "(..., 4)", quaternion.shape))

    checks_base(quaternion, _impl)
//  [python]
def generate_rays(width=800, height=800):
    def normalize(x):
        return x / torch.norm(x, dim=-1, keepdim=True)

    m_eye = torch.tensor((0.0, 1.0, 0.0), dtype=torch.float32)
    m_lookat = torch.tensor((0, 0, 0), dtype=torch.float32)
    m_up = torch.tensor((0, 0, 1), dtype=torch.float32)
    m_fovY = 45.0
    m_aspectRatio = width / height

    W = m_lookat - m_eye
    wlen = torch.norm(W, dim=-1)
    U = normalize(torch.cross(W, m_up))
    V = normalize(torch.cross(U, W))

    vlen = wlen * math.tan(0.5 * m_fovY * math.pi / 180.0)
    V *= vlen
    ulen = vlen * m_aspectRatio
    U *= ulen

    d = torch.stack(
        tuple(
            reversed(
                torch.meshgrid(
                    torch.linspace(0, 1, width),
                    torch.linspace(0, 1, height),
                    indexing="ij",
                )
            )
        ),
        -1,
    )
    d = 2.0 * d - 1.0
    d = d.view(-1, 2)
    ray_directions = normalize(d[:, :1] * U[None] + d[:, 1:] * V[None] + W[None])
    ray_origins = torch.repeat_interleave(m_eye[None], len(ray_directions), 0)
    return ray_origins, ray_directions
//  [python]
def generate_rays(width=800, height=800):
    def normalize(x):
        return x / torch.linalg.norm(x, dim=-1, keepdim=True)

    m_eye = torch.tensor((0.0, 1.0, 0.0), dtype=torch.float32)
    m_lookat = torch.tensor((0, 0, 0), dtype=torch.float32)
    m_up = torch.tensor((0, 0, 1), dtype=torch.float32)
    m_fovY = 45.0
    m_aspectRatio = width / height

    W = m_lookat - m_eye
    wlen = torch.linalg.norm(W, dim=-1)
    U = normalize(torch.linalg.cross(W, m_up))
    V = normalize(torch.linalg.cross(U, W))

    vlen = wlen * math.tan(0.5 * m_fovY * math.pi / 180.0)
    V *= vlen
    ulen = vlen * m_aspectRatio
    U *= ulen

    d = torch.stack(
        tuple(
            reversed(
                torch.meshgrid(
                    torch.linspace(0, 1, width),
                    torch.linspace(0, 1, height),
                    indexing="ij",
                )
            )
        ),
        -1,
    )
    d = 2.0 * d - 1.0
    d = d.view(-1, 2)
    ray_directions = normalize(d[:, :1] * U[None] + d[:, 1:] * V[None] + W[None])
    ray_origins = torch.repeat_interleave(m_eye[None], len(ray_directions), 0)
    return ray_origins, ray_directions
//  [python]
def get_grid(basis: torch.Tensor) -> tuple[torch.Tensor, ...]:
    """Return grid containing linear combinations of basis vectors sufficient to cover Wigner-Seitz cell"""
    a, b, c = basis[:, 0], basis[:, 1], basis[:, 2]
    # Find maximum radius of sphere within parallelopiped from basis lattice vectors
    Rmax = np.max([t.norm(_vec) for _vec in
                   [a, b, c, a + b, b + c, a + c, a + b + c]])
    # This sphere should be within shape_min in each direction
    # Therefore shape_min >= Rmax * (R^-1 or G^-1)
    shape_min = (Rmax * (t.inverse(basis).norm(dim=0)))
    shape_min = torch.ceil(shape_min)
    iGridX = np.arange(-shape_min[0], shape_min[0] + 1)
    iGridY = np.arange(-shape_min[1], shape_min[1] + 1)
    iGridZ = np.arange(-shape_min[2], shape_min[2] + 1)
    # Reduce grid for any orthogonal basis vector
    if np.dot(a, b) == 0 == np.dot(a, c):
        iGridX = np.arange(-1., 2.)
    if np.dot(b, a) == 0 == np.dot(b, c):
        iGridY = np.arange(-1., 2.)
    if np.dot(c, a) == 0 == np.dot(c, b):
        iGridZ = np.arange(-1., 2.)
    iGrid = np.array([x.flatten() for x in np.meshgrid(iGridX, iGridY, iGridZ)]).T
    iGrid = iGrid[np.where((iGrid ** 2).sum(axis=1) > 0)[0]]  # eliminate origin
    return torch.tensor(iGrid), torch.tensor(iGrid) @ basis.T
//  [python]
def get_grid(basis: torch.Tensor) -> tuple[torch.Tensor, ...]:
    """Return grid containing linear combinations of basis vectors sufficient to cover Wigner-Seitz cell"""
    a, b, c = basis[:, 0], basis[:, 1], basis[:, 2]
    # Find maximum radius of sphere within parallelopiped from basis lattice vectors
    Rmax = np.max([torch.linalg.norm(_vec) for _vec in
                   [a, b, c, a + b, b + c, a + c, a + b + c]])
    # This sphere should be within shape_min in each direction
    # Therefore shape_min >= Rmax * (R^-1 or G^-1)
    shape_min = (Rmax * (torch.linalg.inv(basis).norm(dim=0)))
    shape_min = torch.ceil(shape_min)
    iGridX = np.arange(-shape_min[0], shape_min[0] + 1)
    iGridY = np.arange(-shape_min[1], shape_min[1] + 1)
    iGridZ = np.arange(-shape_min[2], shape_min[2] + 1)
    # Reduce grid for any orthogonal basis vector
    if np.dot(a, b) == 0 == np.dot(a, c):
        iGridX = np.arange(-1., 2.)
    if np.dot(b, a) == 0 == np.dot(b, c):
        iGridY = np.arange(-1., 2.)
    if np.dot(c, a) == 0 == np.dot(c, b):
        iGridZ = np.arange(-1., 2.)
    iGrid = np.array([x.flatten() for x in np.meshgrid(iGridX, iGridY, iGridZ)]).T
    iGrid = iGrid[np.where((iGrid ** 2).sum(axis=1) > 0)[0]]  # eliminate origin
    return torch.tensor(iGrid), torch.tensor(iGrid) @ basis.T
//  [python]
def forward(self):
        vecX = self.theta0 / torch.norm(self.theta0)
        vecY = self.theta1 / torch.norm(self.theta1)
        self.vecX = vecX
        self.vecY = vecY
        loss = hf_torch_norm_square((self.matB @ vecY) @ vecX)
        if self.tag_loss_yBx: #do not help
            loss = loss + hf_torch_norm_square((self.matB @ vecX) @ vecY)
        return loss
//  [python]
def forward(self):
        vecX = self.theta0 / torch.linalg.norm(self.theta0)
        vecY = self.theta1 / torch.linalg.norm(self.theta1)
        self.vecX = vecX
        self.vecY = vecY
        loss = hf_torch_norm_square((self.matB @ vecY) @ vecX)
        if self.tag_loss_yBx: #do not help
            loss = loss + hf_torch_norm_square((self.matB @ vecX) @ vecY)
        return loss
//  [python]
def get_weight_map(
    rays,
    jitter_rays,
    cfg,
    weights=None,
    softmax=True
):
    ray_dim = rays.shape[-1] // 2

    # Angles
    angles = torch.acos(
        torch.clip(
            dot(rays[..., ray_dim:], jitter_rays[..., ray_dim:]),
            -1 + 1e-8, 1 - 1e-8
        )
    ).detach()

    # Distances
    dists = torch.norm(
        rays[..., :ray_dim] - jitter_rays[..., :ray_dim],
        dim=-1
    ).detach()

    # Weights
    if weights is None:
        weights = torch.zeros_like(angles)

    if softmax:
        weights = torch.nn.functional.softmax(
            0.5 * -(torch.square(angles / cfg.angle_std) + torch.square(dists / cfg.dist_std)) + weights, dim=0
        )[..., None]
    else:
        #print("Angle:", angles.max(), angles.mean(), cfg.angle_std)
        #print("Dist:", dists.max(), dists.mean(), cfg.dist_std)

        weights = torch.exp(
            0.5 * -(torch.square(angles / cfg.angle_std) + torch.square(dists / cfg.dist_std)) + weights
        )[..., None]

    # Normalization constant
    constant = np.power(2 * np.pi * cfg.angle_std * cfg.angle_std, -1.0 / 2.0) \
        * np.power(2 * np.pi * cfg.dist_std * cfg.dist_std, -1.0 / 2.0)

    return weights / constant
//  [python]
def get_weight_map(
    rays,
    jitter_rays,
    cfg,
    weights=None,
    softmax=True
):
    ray_dim = rays.shape[-1] // 2

    # Angles
    angles = torch.acos(
        torch.clip(
            dot(rays[..., ray_dim:], jitter_rays[..., ray_dim:]),
            -1 + 1e-8, 1 - 1e-8
        )
    ).detach()

    # Distances
    dists = torch.linalg.norm(
        rays[..., :ray_dim] - jitter_rays[..., :ray_dim],
        dim=-1
    ).detach()

    # Weights
    if weights is None:
        weights = torch.zeros_like(angles)

    if softmax:
        weights = torch.nn.functional.softmax(
            0.5 * -(torch.square(angles / cfg.angle_std) + torch.square(dists / cfg.dist_std)) + weights, dim=0
        )[..., None]
    else:
        #print("Angle:", angles.max(), angles.mean(), cfg.angle_std)
        #print("Dist:", dists.max(), dists.mean(), cfg.dist_std)

        weights = torch.exp(
            0.5 * -(torch.square(angles / cfg.angle_std) + torch.square(dists / cfg.dist_std)) + weights
        )[..., None]

    # Normalization constant
    constant = np.power(2 * np.pi * cfg.angle_std * cfg.angle_std, -1.0 / 2.0) \
        * np.power(2 * np.pi * cfg.dist_std * cfg.dist_std, -1.0 / 2.0)

    return weights / constant