jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def dist(v1, v2, axis=None):
    """ Utility function to compute the Euclidean distance between two vectors.
    If only 1D vectors, an scalar is returned. If a 2D or 3D matrix is feed,
    the first dimension is interpreted as time and vector/matrix distance
    is computed along it."""

    if axis is None:
        if len(v1.shape) == 1: axis = 0        # normal 1D vectors
        elif len(v1.shape) == 2: axis = 1      # time x vector matrix, calculate along time
        elif len(v1.shape) == 3: axis = (1, 2)  # time x batch x vector matrix, calculate along time

    if isinstance(v1, torch.Tensor) and isinstance(v2, torch.Tensor):
        d = torch.linalg.norm(v1-v2, dim=axis, ord=2).detach()
    elif isinstance(v1, np.ndarray) and isinstance(v2, np.ndarray):
        d = np.linalg.norm(v1-v2, axis=axis)
    else:
        raise ValueError('Invalid types {},{} for dist function'.format(
            type(v1), type(v2)))

    return d
//  [python]
def compute_average_batch_angle(A, B):
    """
    Compute the average of the angles between the mini-batch samples of A and B.
    If the samples of the mini-batch have more than one dimension (minibatch
    dimension not included), the tensors will first be flattened
    Args:
        A (torch.Tensor):  A tensor with as first dimension the mini-batch
            dimension
        B (torch.Tensor): A tensor of the same shape as A
    Returns: The average angle between the two tensors in degrees.
    """

    A = A.flatten(1, -1)
    B = B.flatten(1, -1)
    if contains_nan(A):
        print('tensor A contains nans in activation angles:')
    if contains_nan(B):
        print('tensor B contains nans in activation angles:')
    inner_products = torch.sum(A*B, dim=1)
    A_norms = torch.norm(A, p=2, dim=1)
    B_norms = torch.norm(B, p=2, dim=1)
    cosines = inner_products/(A_norms*B_norms)
    if contains_nan(cosines):
        print('cosines contains nans in activation angles:')
    if torch.sum(A_norms == 0) > 0:
        print('A_norms contains zeros')
    if torch.sum(B_norms == 0) > 0:
        print('B_norms contains zeros')
    cosines = torch.min(cosines, torch.ones_like(cosines))
    angles = torch.acos(cosines)
    return 180/np.pi*torch.mean(angles)
//  [python]
def compute_average_batch_angle(A, B):
    """
    Compute the average of the angles between the mini-batch samples of A and B.
    If the samples of the mini-batch have more than one dimension (minibatch
    dimension not included), the tensors will first be flattened
    Args:
        A (torch.Tensor):  A tensor with as first dimension the mini-batch
            dimension
        B (torch.Tensor): A tensor of the same shape as A
    Returns: The average angle between the two tensors in degrees.
    """

    A = A.flatten(1, -1)
    B = B.flatten(1, -1)
    if contains_nan(A):
        print('tensor A contains nans in activation angles:')
    if contains_nan(B):
        print('tensor B contains nans in activation angles:')
    inner_products = torch.sum(A*B, dim=1)
    A_norms = torch.linalg.norm(A, p=2, dim=1)
    B_norms = torch.linalg.norm(B, p=2, dim=1)
    cosines = inner_products/(A_norms*B_norms)
    if contains_nan(cosines):
        print('cosines contains nans in activation angles:')
    if torch.sum(A_norms == 0) > 0:
        print('A_norms contains zeros')
    if torch.sum(B_norms == 0) > 0:
        print('B_norms contains zeros')
    cosines = torch.min(cosines, torch.ones_like(cosines))
    angles = torch.acos(cosines)
    return 180/np.pi*torch.mean(angles)
//  [python]
def fix_grad_norm_(parameters, fixed_norm, norm_type):
    """
    Inspired by implementation of clip_grad_norm_ in PyTorch

    Rescale the gradients such that the total norm of the gradients is
    equal to fixed_norm.
    Args:
        parameters (Iterable[torch.Tensor] or torch.Tensor): Iterable
            containing all the parameters that are optimized by the
            network
        fixed_norm (float or int): positive float indicating the total norm
            to which we want to rescale the gradients
        norm_type (float or int): type of the used p-norm. Can be ``'inf'`` for
            infinity norm.
    Returns:
        The total norm of the parameters before rescaling. This method
        changes the gradients of the parameters in-place.
    """

    if isinstance(parameters, torch.Tensor):
        parameters = [parameters]

    parameters = [p for p in parameters if p.grad is not None]
    assert fixed_norm > 0
    fixed_norm = float(fixed_norm)
    norm_type = float(norm_type)
    if len(parameters) == 0:
        return torch.tensor(0.)
    device = parameters[0].grad.device
    if norm_type == float('inf'):
        total_norm = max(p.grad.detach().abs().max().to(device) for p in parameters)
    else:
        total_norm = torch.norm(torch.stack([torch.norm(p.grad.detach(),
                                norm_type).to(device) for p in parameters]), norm_type)
    clip_coef = fixed_norm / (total_norm + 1e-6)
    for p in parameters:
        p.grad.detach().mul_(clip_coef.to(p.grad.device))
    return total_norm
//  [python]
def fix_grad_norm_(parameters, fixed_norm, norm_type):
    """
    Inspired by implementation of clip_grad_norm_ in PyTorch

    Rescale the gradients such that the total norm of the gradients is
    equal to fixed_norm.
    Args:
        parameters (Iterable[torch.Tensor] or torch.Tensor): Iterable
            containing all the parameters that are optimized by the
            network
        fixed_norm (float or int): positive float indicating the total norm
            to which we want to rescale the gradients
        norm_type (float or int): type of the used p-norm. Can be ``'inf'`` for
            infinity norm.
    Returns:
        The total norm of the parameters before rescaling. This method
        changes the gradients of the parameters in-place.
    """

    if isinstance(parameters, torch.Tensor):
        parameters = [parameters]

    parameters = [p for p in parameters if p.grad is not None]
    assert fixed_norm > 0
    fixed_norm = float(fixed_norm)
    norm_type = float(norm_type)
    if len(parameters) == 0:
        return torch.tensor(0.)
    device = parameters[0].grad.device
    if norm_type == float('inf'):
        total_norm = max(p.grad.detach().abs().max().to(device) for p in parameters)
    else:
        total_norm = torch.linalg.norm(torch.stack([torch.linalg.norm(p.grad.detach(),
                                norm_type).to(device) for p in parameters]), norm_type)
    clip_coef = fixed_norm / (total_norm + 1e-6)
    for p in parameters:
        p.grad.detach().mul_(clip_coef.to(p.grad.device))
    return total_norm
//  [python]
def test_local_gain_control(self):
        x = torch.randn((10, 1, 256, 256), device=DEVICE)
        norm, direction = po.simul.non_linearities.local_gain_control(x)
        x_hat = po.simul.non_linearities.local_gain_release(norm, direction)
        assert torch.norm(x - x_hat) < 1e-4
//  [python]
def test_local_gain_control(self):
        x = torch.randn((10, 1, 256, 256), device=DEVICE)
        norm, direction = po.simul.non_linearities.local_gain_control(x)
        x_hat = po.simul.non_linearities.local_gain_release(norm, direction)
        assert torch.linalg.norm(x - x_hat) < 1e-4
//  [python]
def rot_from_2_vectors(v1, v2):
        """ Returns a Rotation matrix between vectors 'v1' and 'v2'    """
        v1 = v1/torch.norm(v1)
        v2 = v2/torch.norm(v2)
        v = torch.cross(v1, v2)
        cosang = v1.matmul(v2)
        sinang = torch.norm(v)
        Rot = TORCHIEKF.Id3 + TORCHIEKF.skew(v) + \
              TORCHIEKF.skew(v).mm(TORCHIEKF.skew(v))*(1-cosang)/(sinang**2)
        return Rot
//  [python]
def rot_from_2_vectors(v1, v2):
        """ Returns a Rotation matrix between vectors 'v1' and 'v2'    """
        v1 = v1/torch.linalg.norm(v1)
        v2 = v2/torch.linalg.norm(v2)
        v = torch.cross(v1, v2)
        cosang = v1.matmul(v2)
        sinang = torch.linalg.norm(v)
        Rot = TORCHIEKF.Id3 + TORCHIEKF.skew(v) + \
              TORCHIEKF.skew(v).mm(TORCHIEKF.skew(v))*(1-cosang)/(sinang**2)
        return Rot
//  [python]
def test_local_gain_control_dict(self, basic_stim):
        spr = po.simul.SteerablePyramidFreq(basic_stim.shape[-2:], height=5,
                                              order=1, is_complex=False, tight_frame=True).to(DEVICE)
        y = spr(basic_stim)
        energy, state = po.simul.non_linearities.local_gain_control_dict(y, residuals=True)
        y_hat = po.simul.non_linearities.local_gain_release_dict(energy, state, residuals=True)
        for key in y.keys():
            assert torch.norm(y[key] - y_hat[key]) < 1e-5
//  [python]
def test_local_gain_control_dict(self, basic_stim):
        spr = po.simul.SteerablePyramidFreq(basic_stim.shape[-2:], height=5,
                                              order=1, is_complex=False, tight_frame=True).to(DEVICE)
        y = spr(basic_stim)
        energy, state = po.simul.non_linearities.local_gain_control_dict(y, residuals=True)
        y_hat = po.simul.non_linearities.local_gain_release_dict(energy, state, residuals=True)
        for key in y.keys():
            assert torch.linalg.norm(y[key] - y_hat[key]) < 1e-5
//  [python]
def raw2outputs(raw, z_vals, rays_d, raw_noise_std=0, white_bkgd=False, pytest=False):
    """Transforms model's predictions to semantically meaningful values.
    Args:
        raw: [num_rays, num_samples along ray, 4]. Prediction from model.
        z_vals: [num_rays, num_samples along ray]. Integration time.
        rays_d: [num_rays, 3]. Direction of each ray.
    Returns:
        rgb_map: [num_rays, 3]. Estimated RGB color of a ray.
        disp_map: [num_rays]. Disparity map. Inverse of depth map.
        acc_map: [num_rays]. Sum of weights along each ray.
        weights: [num_rays, num_samples]. Weights assigned to each sampled color.
        depth_map: [num_rays]. Estimated distance to object.
    """
    #ipdb.set_trace()
    act_ff=nn.Softplus()

    raw2alpha = lambda raw, dists, act_fn=act_ff: 1.-torch.exp(-act_fn(raw)*dists)

    dists = z_vals[...,1:] - z_vals[...,:-1]
    dists = torch.cat([dists, torch.Tensor([1e10]).to(dists.device).expand(dists[...,:1].shape)], -1)  # [N_rays, N_samples]

    dists = dists * torch.norm(rays_d[...,None,:], dim=-1) 

    rgb = torch.sigmoid(raw[...,:3])  # [N_rays, N_samples, 3]
    noise = 0.
    if raw_noise_std > 0.:
        noise = torch.randn(raw[...,3].shape) * raw_noise_std

        # Overwrite randomly sampled data if pytest
        if pytest:
            np.random.seed(0)
            noise = np.random.rand(*list(raw[...,3].shape)) * raw_noise_std
            noise = torch.Tensor(noise)
    #ipdb.set_trace()
    alpha = raw2alpha(raw[...,3] + noise, dists)  # [N_rays, N_samples]
    
    #ipdb.set_trace()  
    weights = alpha * torch.cumprod(torch.cat([torch.ones((alpha.shape[0],alpha.shape[1], 1)).to(alpha.device), 1.-alpha + 1e-10], -1), -1)[:,:, :-1]
    rgb_map = torch.sum(weights[...,None] * rgb, -2)  # [N_rays, 3]

    depth_map = torch.sum(weights * z_vals, -1)
    disp_map = 1./torch.max(1e-10 * torch.ones_like(depth_map), depth_map / torch.sum(weights, -1))
    acc_map = torch.sum(weights, -1)

    if white_bkgd:
        rgb_map = rgb_map + (1.-acc_map[...,None])

    return rgb_map, disp_map, acc_map, weights, depth_map
//  [python]
def raw2outputs(raw, z_vals, rays_d, raw_noise_std=0, white_bkgd=False, pytest=False):
    """Transforms model's predictions to semantically meaningful values.
    Args:
        raw: [num_rays, num_samples along ray, 4]. Prediction from model.
        z_vals: [num_rays, num_samples along ray]. Integration time.
        rays_d: [num_rays, 3]. Direction of each ray.
    Returns:
        rgb_map: [num_rays, 3]. Estimated RGB color of a ray.
        disp_map: [num_rays]. Disparity map. Inverse of depth map.
        acc_map: [num_rays]. Sum of weights along each ray.
        weights: [num_rays, num_samples]. Weights assigned to each sampled color.
        depth_map: [num_rays]. Estimated distance to object.
    """
    #ipdb.set_trace()
    act_ff=nn.Softplus()

    raw2alpha = lambda raw, dists, act_fn=act_ff: 1.-torch.exp(-act_fn(raw)*dists)

    dists = z_vals[...,1:] - z_vals[...,:-1]
    dists = torch.cat([dists, torch.Tensor([1e10]).to(dists.device).expand(dists[...,:1].shape)], -1)  # [N_rays, N_samples]

    dists = dists * torch.linalg.norm(rays_d[...,None,:], dim=-1) 

    rgb = torch.sigmoid(raw[...,:3])  # [N_rays, N_samples, 3]
    noise = 0.
    if raw_noise_std > 0.:
        noise = torch.randn(raw[...,3].shape) * raw_noise_std

        # Overwrite randomly sampled data if pytest
        if pytest:
            np.random.seed(0)
            noise = np.random.rand(*list(raw[...,3].shape)) * raw_noise_std
            noise = torch.Tensor(noise)
    #ipdb.set_trace()
    alpha = raw2alpha(raw[...,3] + noise, dists)  # [N_rays, N_samples]
    
    #ipdb.set_trace()  
    weights = alpha * torch.cumprod(torch.cat([torch.ones((alpha.shape[0],alpha.shape[1], 1)).to(alpha.device), 1.-alpha + 1e-10], -1), -1)[:,:, :-1]
    rgb_map = torch.sum(weights[...,None] * rgb, -2)  # [N_rays, 3]

    depth_map = torch.sum(weights * z_vals, -1)
    disp_map = 1./torch.max(1e-10 * torch.ones_like(depth_map), depth_map / torch.sum(weights, -1))
    acc_map = torch.sum(weights, -1)

    if white_bkgd:
        rgb_map = rgb_map + (1.-acc_map[...,None])

    return rgb_map, disp_map, acc_map, weights, depth_map
//  [python]
def so3left_jacobian(phi):
        angle = torch.norm(phi)

        # Near |phi|==0, use first order Taylor expansion
        if isclose(angle, 0.):
            skew_phi = torch.Tensor([[0, -phi[2], phi[1]],
                          [phi[2], 0, -phi[0]],
                          [-phi[1], phi[0], 0]]).double()
            return TORCHIEKF.Id3 + 0.5 * skew_phi

        axis = phi / angle
        skew_axis = torch.Tensor([[0, -axis[2], axis[1]],
                          [axis[2], 0, -axis[0]],
                          [-axis[1], axis[0], 0]]).double()
        s = torch.sin(angle)
        c = torch.cos(angle)

        return (s / angle) * TORCHIEKF.Id3 + (1 - s / angle) * TORCHIEKF.outer(axis, axis)\
               + ((1 - c) / angle) * skew_axis
//  [python]
def so3left_jacobian(phi):
        angle = torch.linalg.norm(phi)

        # Near |phi|==0, use first order Taylor expansion
        if isclose(angle, 0.):
            skew_phi = torch.Tensor([[0, -phi[2], phi[1]],
                          [phi[2], 0, -phi[0]],
                          [-phi[1], phi[0], 0]]).double()
            return TORCHIEKF.Id3 + 0.5 * skew_phi

        axis = phi / angle
        skew_axis = torch.Tensor([[0, -axis[2], axis[1]],
                          [axis[2], 0, -axis[0]],
                          [-axis[1], axis[0], 0]]).double()
        s = torch.sin(angle)
        c = torch.cos(angle)

        return (s / angle) * TORCHIEKF.Id3 + (1 - s / angle) * TORCHIEKF.outer(axis, axis)\
               + ((1 - c) / angle) * skew_axis
//  [python]
def sen3exp(xi):
        phi = xi[:3]
        angle = torch.norm(phi)

        # Near |phi|==0, use first order Taylor expansion
        if isclose(angle, 0.):
            skew_phi = torch.Tensor([[0, -phi[2], phi[1]],
                          [phi[2], 0, -phi[0]],
                          [-phi[1], phi[0], 0]]).double()
            J = TORCHIEKF.Id3 + 0.5 * skew_phi
            Rot = TORCHIEKF.Id3 + skew_phi
        else:
            axis = phi / angle
            skew_axis = torch.Tensor([[0, -axis[2], axis[1]],
                              [axis[2], 0, -axis[0]],
                              [-axis[1], axis[0], 0]]).double()
            s = torch.sin(angle)
            c = torch.cos(angle)

            J = (s / angle) * TORCHIEKF.Id3 + (1 - s / angle) * TORCHIEKF.outer(axis, axis)\
                   + ((1 - c) / angle) * skew_axis
            Rot = c * TORCHIEKF.Id3 + (1 - c) * TORCHIEKF.outer(axis, axis) \
                 + s * skew_axis

        x = J.mm(xi[3:].view(-1, 3).t())
        return Rot, x
//  [python]
def sen3exp(xi):
        phi = xi[:3]
        angle = torch.linalg.norm(phi)

        # Near |phi|==0, use first order Taylor expansion
        if isclose(angle, 0.):
            skew_phi = torch.Tensor([[0, -phi[2], phi[1]],
                          [phi[2], 0, -phi[0]],
                          [-phi[1], phi[0], 0]]).double()
            J = TORCHIEKF.Id3 + 0.5 * skew_phi
            Rot = TORCHIEKF.Id3 + skew_phi
        else:
            axis = phi / angle
            skew_axis = torch.Tensor([[0, -axis[2], axis[1]],
                              [axis[2], 0, -axis[0]],
                              [-axis[1], axis[0], 0]]).double()
            s = torch.sin(angle)
            c = torch.cos(angle)

            J = (s / angle) * TORCHIEKF.Id3 + (1 - s / angle) * TORCHIEKF.outer(axis, axis)\
                   + ((1 - c) / angle) * skew_axis
            Rot = c * TORCHIEKF.Id3 + (1 - c) * TORCHIEKF.outer(axis, axis) \
                 + s * skew_axis

        x = J.mm(xi[3:].view(-1, 3).t())
        return Rot, x
//  [python]
def render(chunk=1024*32, rays=None, c2w=None, ndc=True,label=None,
                  near=0., far=1.,
                  use_viewdirs=False, c2w_staticcam=None,
                  **kwargs):
    """Render rays
    Args:
      H: int. Height of image in pixels.
      W: int. Width of image in pixels.
      focal: float. Focal length of pinhole camera.
      chunk: int. Maximum number of rays to process simultaneously. Used to
        control maximum memory usage. Does not affect final results.
      rays: array of shape [2, batch_size, 3]. Ray origin and direction for
        each example in batch.
      c2w: array of shape [3, 4]. Camera-to-world transformation matrix.
      ndc: bool. If True, represent ray origin, direction in NDC coordinates.
      near: float or array of shape [batch_size]. Nearest distance for a ray.
      far: float or array of shape [batch_size]. Farthest distance for a ray.
      use_viewdirs: bool. If True, use viewing direction of a point in space in model.
      c2w_staticcam: array of shape [3, 4]. If not None, use this transformation matrix for 
       camera while using other c2w argument for viewing directions.
    Returns:
      rgb_map: [batch_size, 3]. Predicted RGB values for rays.
      disp_map: [batch_size]. Disparity map. Inverse of depth.
      acc_map: [batch_size]. Accumulated opacity (alpha) along a ray.
      extras: dict with everything returned by render_rays().
    """
    #ipdb.set_trace()

    rays_o, rays_d = rays[:,0,...], rays[:,1,...]

    
    viewdirs = rays_d

    viewdirs = viewdirs / torch.norm(viewdirs, dim=-1, keepdim=True)
    viewdirs = torch.reshape(viewdirs, [rays_d.shape[0],-1,3]).float()

    sh = rays_d.shape # [..., 3]

    # Create ray batch
    rays_o = torch.reshape(rays_o, [sh[0],-1,3]).float()
    rays_d = torch.reshape(rays_d, [sh[0],-1,3]).float()
    #ipdb.set_trace()

    near, far = near[:,None,:] * torch.ones_like(rays_d[...,:1]), far[:,None,:] * torch.ones_like(rays_d[...,:1])
    rays = torch.cat([rays_o, rays_d, near, far], -1)
    if use_viewdirs:
        rays = torch.cat([rays, viewdirs], -1)
    #ipdb.set_trace()
    # Render and reshape
    all_ret = batchify_rays(rays, label,chunk,**kwargs)
    for k in all_ret:
        k_sh = list(sh[:-1]) + list(all_ret[k].shape[2:])
        all_ret[k] = torch.reshape(all_ret[k], k_sh)

    k_extract = ['rgb_map', 'disp_map', 'acc_map']
    ret_list = [all_ret[k] for k in k_extract]
    ret_dict = {k : all_ret[k] for k in all_ret if k not in k_extract}
    return ret_list + [ret_dict]
//  [python]
def render(chunk=1024*32, rays=None, c2w=None, ndc=True,label=None,
                  near=0., far=1.,
                  use_viewdirs=False, c2w_staticcam=None,
                  **kwargs):
    """Render rays
    Args:
      H: int. Height of image in pixels.
      W: int. Width of image in pixels.
      focal: float. Focal length of pinhole camera.
      chunk: int. Maximum number of rays to process simultaneously. Used to
        control maximum memory usage. Does not affect final results.
      rays: array of shape [2, batch_size, 3]. Ray origin and direction for
        each example in batch.
      c2w: array of shape [3, 4]. Camera-to-world transformation matrix.
      ndc: bool. If True, represent ray origin, direction in NDC coordinates.
      near: float or array of shape [batch_size]. Nearest distance for a ray.
      far: float or array of shape [batch_size]. Farthest distance for a ray.
      use_viewdirs: bool. If True, use viewing direction of a point in space in model.
      c2w_staticcam: array of shape [3, 4]. If not None, use this transformation matrix for 
       camera while using other c2w argument for viewing directions.
    Returns:
      rgb_map: [batch_size, 3]. Predicted RGB values for rays.
      disp_map: [batch_size]. Disparity map. Inverse of depth.
      acc_map: [batch_size]. Accumulated opacity (alpha) along a ray.
      extras: dict with everything returned by render_rays().
    """
    #ipdb.set_trace()

    rays_o, rays_d = rays[:,0,...], rays[:,1,...]

    
    viewdirs = rays_d

    viewdirs = viewdirs / torch.linalg.norm(viewdirs, dim=-1, keepdim=True)
    viewdirs = torch.reshape(viewdirs, [rays_d.shape[0],-1,3]).float()

    sh = rays_d.shape # [..., 3]

    # Create ray batch
    rays_o = torch.reshape(rays_o, [sh[0],-1,3]).float()
    rays_d = torch.reshape(rays_d, [sh[0],-1,3]).float()
    #ipdb.set_trace()

    near, far = near[:,None,:] * torch.ones_like(rays_d[...,:1]), far[:,None,:] * torch.ones_like(rays_d[...,:1])
    rays = torch.cat([rays_o, rays_d, near, far], -1)
    if use_viewdirs:
        rays = torch.cat([rays, viewdirs], -1)
    #ipdb.set_trace()
    # Render and reshape
    all_ret = batchify_rays(rays, label,chunk,**kwargs)
    for k in all_ret:
        k_sh = list(sh[:-1]) + list(all_ret[k].shape[2:])
        all_ret[k] = torch.reshape(all_ret[k], k_sh)

    k_extract = ['rgb_map', 'disp_map', 'acc_map']
    ret_list = [all_ret[k] for k in k_extract]
    ret_dict = {k : all_ret[k] for k in all_ret if k not in k_extract}
    return ret_list + [ret_dict]
//  [python]
def calculate_adaptive_weight(self, nll_loss, g_loss, last_layer=None):
        if last_layer is not None:
            nll_grads = torch.autograd.grad(nll_loss, last_layer, retain_graph=True)[0]
            g_grads = torch.autograd.grad(g_loss, last_layer, retain_graph=True)[0]
        else:
            nll_grads = torch.autograd.grad(nll_loss, self.last_layer[0], retain_graph=True)[0]
            g_grads = torch.autograd.grad(g_loss, self.last_layer[0], retain_graph=True)[0]

        d_weight = torch.norm(nll_grads) / (torch.norm(g_grads) + 1e-4)
        d_weight = torch.clamp(d_weight, 0.0, 1e4).detach()
        d_weight = d_weight * self.discriminator_weight
        return d_weight
//  [python]
def calculate_adaptive_weight(self, nll_loss, g_loss, last_layer=None):
        if last_layer is not None:
            nll_grads = torch.autograd.grad(nll_loss, last_layer, retain_graph=True)[0]
            g_grads = torch.autograd.grad(g_loss, last_layer, retain_graph=True)[0]
        else:
            nll_grads = torch.autograd.grad(nll_loss, self.last_layer[0], retain_graph=True)[0]
            g_grads = torch.autograd.grad(g_loss, self.last_layer[0], retain_graph=True)[0]

        d_weight = torch.linalg.norm(nll_grads) / (torch.linalg.norm(g_grads) + 1e-4)
        d_weight = torch.clamp(d_weight, 0.0, 1e4).detach()
        d_weight = d_weight * self.discriminator_weight
        return d_weight
//  [python]
def test_transformer_vs_pytorch():
    torch.manual_seed(1234)
    # Check that in the non causal setting, we get the same result as
    # PyTorch Transformer encoder.
    for custom in [False, True]:
        tr = StreamingTransformer(
            16, 4, 2,
            causal=False, custom=custom, dropout=0., positional_scale=0.)
        layer = torch.nn.TransformerEncoderLayer(16, 4, dropout=0., batch_first=True)
        tr_ref = torch.nn.TransformerEncoder(layer, 2)
        tr.load_state_dict(tr_ref.state_dict())

        x = torch.randn(4, 20, 16)
        y = tr(x)
        y2 = tr_ref(x)
        delta = torch.norm(y2 - y) / torch.norm(y)
        assert delta < 1e-6, delta
//  [python]
def test_transformer_vs_pytorch():
    torch.manual_seed(1234)
    # Check that in the non causal setting, we get the same result as
    # PyTorch Transformer encoder.
    for custom in [False, True]:
        tr = StreamingTransformer(
            16, 4, 2,
            causal=False, custom=custom, dropout=0., positional_scale=0.)
        layer = torch.nn.TransformerEncoderLayer(16, 4, dropout=0., batch_first=True)
        tr_ref = torch.nn.TransformerEncoder(layer, 2)
        tr.load_state_dict(tr_ref.state_dict())

        x = torch.randn(4, 20, 16)
        y = tr(x)
        y2 = tr_ref(x)
        delta = torch.linalg.norm(y2 - y) / torch.linalg.norm(y)
        assert delta < 1e-6, delta
//  [python]
def test_transformer_causal_streaming():
    torch.manual_seed(1234)

    for context, custom in product([None, 10], [False, True]):
        # Test that causality and receptive fields are properly handled.
        # looking at the gradients
        tr = StreamingTransformer(
            16, 4, 1 if context else 2,
            causal=True, past_context=context, custom=custom,
            dropout=0.)
        steps = 20
        for k in [0, 10, 15, 19]:
            x = torch.randn(4, steps, 16, requires_grad=True)
            y = tr(x)
            y[:, k].abs().sum().backward()
            if k + 1 < steps:
                assert torch.allclose(x.grad[:, k + 1:], torch.tensor(0.)), x.grad[:, k + 1:].norm()
            assert not torch.allclose(x.grad[:, :k + 1], torch.tensor(0.)), x.grad[:, :k + 1].norm()
            if context is not None and k > context:
                limit = k - context - 1
                assert torch.allclose(x.grad[:, :limit],
                                      torch.tensor(0.)), x.grad[:, :limit].norm()

        # Now check that streaming gives the same result at batch eval.
        x = torch.randn(4, steps, 16)
        y = tr(x)
        ys = []
        with tr.streaming():
            for k in range(steps):
                chunk = x[:, k:k + 1, :]
                ys.append(tr(chunk))
        y_stream = torch.cat(ys, dim=1)
        delta = torch.norm(y_stream - y) / torch.norm(y)
        assert delta < 1e-6, delta
//  [python]
def test_transformer_causal_streaming():
    torch.manual_seed(1234)

    for context, custom in product([None, 10], [False, True]):
        # Test that causality and receptive fields are properly handled.
        # looking at the gradients
        tr = StreamingTransformer(
            16, 4, 1 if context else 2,
            causal=True, past_context=context, custom=custom,
            dropout=0.)
        steps = 20
        for k in [0, 10, 15, 19]:
            x = torch.randn(4, steps, 16, requires_grad=True)
            y = tr(x)
            y[:, k].abs().sum().backward()
            if k + 1 < steps:
                assert torch.allclose(x.grad[:, k + 1:], torch.tensor(0.)), x.grad[:, k + 1:].norm()
            assert not torch.allclose(x.grad[:, :k + 1], torch.tensor(0.)), x.grad[:, :k + 1].norm()
            if context is not None and k > context:
                limit = k - context - 1
                assert torch.allclose(x.grad[:, :limit],
                                      torch.tensor(0.)), x.grad[:, :limit].norm()

        # Now check that streaming gives the same result at batch eval.
        x = torch.randn(4, steps, 16)
        y = tr(x)
        ys = []
        with tr.streaming():
            for k in range(steps):
                chunk = x[:, k:k + 1, :]
                ys.append(tr(chunk))
        y_stream = torch.cat(ys, dim=1)
        delta = torch.linalg.norm(y_stream - y) / torch.linalg.norm(y)
        assert delta < 1e-6, delta
//  [python]
def test_streaming_memory_efficient():
    for backend in ['torch']:
        torch.manual_seed(1234)
        set_efficient_attention_backend(backend)
        tr = StreamingTransformer(16, 4, 2, causal=True, dropout=0., custom=True)
        tr_mem_efficient = StreamingTransformer(
            16, 4, 2, dropout=0., memory_efficient=True, causal=True)
        tr.load_state_dict(tr_mem_efficient.state_dict())
        tr.eval()
        tr_mem_efficient.eval()
        steps = 12
        x = torch.randn(3, steps, 16)

        ref = tr(x)

        with tr_mem_efficient.streaming():
            outs = []
            # frame_sizes = [2] + [1] * (steps - 2)
            frame_sizes = [1] * steps

            for frame_size in frame_sizes:
                frame = x[:, :frame_size]
                x = x[:, frame_size:]
                outs.append(tr_mem_efficient(frame))

        out = torch.cat(outs, dim=1)
        delta = torch.norm(out - ref) / torch.norm(out)
        assert delta < 1e-6, delta
//  [python]
def test_streaming_memory_efficient():
    for backend in ['torch']:
        torch.manual_seed(1234)
        set_efficient_attention_backend(backend)
        tr = StreamingTransformer(16, 4, 2, causal=True, dropout=0., custom=True)
        tr_mem_efficient = StreamingTransformer(
            16, 4, 2, dropout=0., memory_efficient=True, causal=True)
        tr.load_state_dict(tr_mem_efficient.state_dict())
        tr.eval()
        tr_mem_efficient.eval()
        steps = 12
        x = torch.randn(3, steps, 16)

        ref = tr(x)

        with tr_mem_efficient.streaming():
            outs = []
            # frame_sizes = [2] + [1] * (steps - 2)
            frame_sizes = [1] * steps

            for frame_size in frame_sizes:
                frame = x[:, :frame_size]
                x = x[:, frame_size:]
                outs.append(tr_mem_efficient(frame))

        out = torch.cat(outs, dim=1)
        delta = torch.linalg.norm(out - ref) / torch.linalg.norm(out)
        assert delta < 1e-6, delta
//  [python]
def translation_error(t1, t2):
    """
    Torch batch implementation of the rotation error between the estimated and the ground truth rotatiom matrix. 
    Rotation error is defined as r_e = \arccos(\frac{Trace(\mathbf{R}_{ij}^{T}\mathbf{R}_{ij}^{\mathrm{GT}) - 1}{2})

    Args: 
        t1 (torch tensor): Estimated translation vectors [b,3,1]
        t2 (torch tensor): Ground truth translation vectors [b,3,1]

    Returns:
        te (torch tensor): translation error in meters [b,1]

    """
    return torch.norm(t1-t2, dim=(1, 2))
//  [python]
def translation_error(t1, t2):
    """
    Torch batch implementation of the rotation error between the estimated and the ground truth rotatiom matrix. 
    Rotation error is defined as r_e = \arccos(\frac{Trace(\mathbf{R}_{ij}^{T}\mathbf{R}_{ij}^{\mathrm{GT}) - 1}{2})

    Args: 
        t1 (torch tensor): Estimated translation vectors [b,3,1]
        t2 (torch tensor): Ground truth translation vectors [b,3,1]

    Returns:
        te (torch tensor): translation error in meters [b,1]

    """
    return torch.linalg.norm(t1-t2, dim=(1, 2))
//  [python]
def forward(self, utterances, hidden_init=None):
        """
        Computes the embeddings of a batch of utterance spectrograms.
        
        :param utterances: batch of mel-scale filterbanks of same duration as a tensor of shape 
        (batch_size, n_frames, n_channels) 
        :param hidden_init: initial hidden state of the LSTM as a tensor of shape (num_layers, 
        batch_size, hidden_size). Will default to a tensor of zeros if None.
        :return: the embeddings as a tensor of shape (batch_size, embedding_size)
        """
        out, (hidden, cell) = self.lstm(utterances, hidden_init)
        embeds_raw = self.relu(self.linear(hidden[-1]))
        embeds = embeds_raw / torch.norm(embeds_raw, dim=1, keepdim=True)
        return embeds
//  [python]
def forward(self, utterances, hidden_init=None):
        """
        Computes the embeddings of a batch of utterance spectrograms.
        
        :param utterances: batch of mel-scale filterbanks of same duration as a tensor of shape 
        (batch_size, n_frames, n_channels) 
        :param hidden_init: initial hidden state of the LSTM as a tensor of shape (num_layers, 
        batch_size, hidden_size). Will default to a tensor of zeros if None.
        :return: the embeddings as a tensor of shape (batch_size, embedding_size)
        """
        out, (hidden, cell) = self.lstm(utterances, hidden_init)
        embeds_raw = self.relu(self.linear(hidden[-1]))
        embeds = embeds_raw / torch.linalg.norm(embeds_raw, dim=1, keepdim=True)
        return embeds
//  [python]
def pairwise_distances(a, b, p=2, inv=False, num_samples=500):
    if not inv:
        tmp = a; a = b; b= tmp

    a = torch.tensor(a[None, :, :]).cuda()
    b = torch.tensor(b[None, :, :]).cuda()
    num_batches = a.shape[1] // num_samples

    dists = []
    for i in range(num_batches):
        dist = torch.norm((a[:,i*num_samples : (i+1)*num_samples, None, :] - b[:, None, :, :]),p=2,dim=3)
        dist, _ = torch.min(dist, 2)
        dist = dist.reshape(-1)
        dists.append(dist)

    if a.shape[1] % num_samples > 0:
        dist = torch.norm((a[:,-1 * (a.shape[1] % num_samples):, None, :] - b[:, None, :, :]),p=2,dim=3)
        dist, _ = torch.min(dist, 2)
        dist = dist.reshape(-1)
        dists.append(dist)

    dist= torch.cat(dists).mean().cpu()
    return dist
//  [python]
def pairwise_distances(a, b, p=2, inv=False, num_samples=500):
    if not inv:
        tmp = a; a = b; b= tmp

    a = torch.tensor(a[None, :, :]).cuda()
    b = torch.tensor(b[None, :, :]).cuda()
    num_batches = a.shape[1] // num_samples

    dists = []
    for i in range(num_batches):
        dist = torch.linalg.norm((a[:,i*num_samples : (i+1)*num_samples, None, :] - b[:, None, :, :]),p=2,dim=3)
        dist, _ = torch.min(dist, 2)
        dist = dist.reshape(-1)
        dists.append(dist)

    if a.shape[1] % num_samples > 0:
        dist = torch.linalg.norm((a[:,-1 * (a.shape[1] % num_samples):, None, :] - b[:, None, :, :]),p=2,dim=3)
        dist, _ = torch.min(dist, 2)
        dist = dist.reshape(-1)
        dists.append(dist)

    dist= torch.cat(dists).mean().cpu()
    return dist
//  [python]
def add_implant_residual(r,keyp,ikeyp,body_support_0, setpoint = 0.0135, scaling = 10.):
    #     # stack on first dim?
    #  [c_hip,c_ass,c_mid,c_nose,c_tip,c_impl]
    c_impl = body_support_0[5][...,0]

    keyp_implant = (ikeyp == 0)

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

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

    keyp_implant = (ikeyp == 0)

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

        r_implant = scaling * torch.abs(keypoint_distance - setpoint)
        #         print(r_implant)
        r = torch.cat([r,r_implant],dim=1)
    return r
//  [python]
def similarity_matrix(self, embeds):
        """
        Computes the similarity matrix according the section 2.1 of GE2E.

        :param embeds: the embeddings as a tensor of shape (speakers_per_batch, 
        utterances_per_speaker, embedding_size)
        :return: the similarity matrix as a tensor of shape (speakers_per_batch,
        utterances_per_speaker, speakers_per_batch)
        """
        speakers_per_batch, utterances_per_speaker = embeds.shape[:2]
        centroids_incl = torch.mean(embeds, dim=1, keepdim=True)
        centroids_incl = centroids_incl.clone() / torch.norm(centroids_incl, dim=2, keepdim=True)
        centroids_excl = torch.sum(embeds, dim=1, keepdim=True) - embeds
        centroids_excl /= utterances_per_speaker - 1
        centroids_excl = centroids_excl.clone() / torch.norm(centroids_excl, dim=2, keepdim=True)
        sim_matrix = torch.zeros(speakers_per_batch, utterances_per_speaker, speakers_per_batch)
        mask_matrix = 1 - np.eye(speakers_per_batch, dtype=np.int)
        for j in range(speakers_per_batch):
            mask = np.where(mask_matrix[j])[0]
            sim_matrix[mask, :, j] = (embeds[mask] * centroids_incl[j]).sum(dim=2)
            sim_matrix[j, :, j] = (embeds[j] * centroids_excl[j]).sum(dim=1)
        sim_matrix = sim_matrix * self.similarity_weight + self.similarity_bias
        return sim_matrix
//  [python]
def similarity_matrix(self, embeds):
        """
        Computes the similarity matrix according the section 2.1 of GE2E.

        :param embeds: the embeddings as a tensor of shape (speakers_per_batch, 
        utterances_per_speaker, embedding_size)
        :return: the similarity matrix as a tensor of shape (speakers_per_batch,
        utterances_per_speaker, speakers_per_batch)
        """
        speakers_per_batch, utterances_per_speaker = embeds.shape[:2]
        centroids_incl = torch.mean(embeds, dim=1, keepdim=True)
        centroids_incl = centroids_incl.clone() / torch.linalg.norm(centroids_incl, dim=2, keepdim=True)
        centroids_excl = torch.sum(embeds, dim=1, keepdim=True) - embeds
        centroids_excl /= utterances_per_speaker - 1
        centroids_excl = centroids_excl.clone() / torch.linalg.norm(centroids_excl, dim=2, keepdim=True)
        sim_matrix = torch.zeros(speakers_per_batch, utterances_per_speaker, speakers_per_batch)
        mask_matrix = 1 - np.eye(speakers_per_batch, dtype=np.int)
        for j in range(speakers_per_batch):
            mask = np.where(mask_matrix[j])[0]
            sim_matrix[mask, :, j] = (embeds[mask] * centroids_incl[j]).sum(dim=2)
            sim_matrix[j, :, j] = (embeds[j] * centroids_excl[j]).sum(dim=1)
        sim_matrix = sim_matrix * self.similarity_weight + self.similarity_bias
        return sim_matrix
//  [python]
def add_ass_residual(r,keyp,ikeyp,body_support_0,body_support_1,which_keyp = 3, setpoint = 0.0, scaling = 10.):
    #     # stack on first dim?
    #  [c_hip,c_ass,c_mid,c_nose,c_tip,c_impl]
    c_impl = torch.cat(( body_support_0[1][...,0], body_support_1[1][...,0]))

    keyp_implant = (ikeyp == which_keyp)

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

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

        r = torch.cat([r,r_implant],dim=1)
    return r
//  [python]
def add_ass_residual(r,keyp,ikeyp,body_support_0,body_support_1,which_keyp = 3, setpoint = 0.0, scaling = 10.):
    c_impl = torch.cat(( body_support_0[1][...,0], body_support_1[1][...,0]))

    keyp_implant = (ikeyp == which_keyp)

    if sum(keyp_implant) > 0:
        target_keyp = keyp[keyp_implant,:]
        keypoint_distance = torch.linalg.norm( c_impl[:,np.newaxis,:] - target_keyp[np.newaxis,:,:] ,dim=2)
        keypoint_distance = torch.min(keypoint_distance[:r.shape[0],:], keypoint_distance[r.shape[0]:,:])

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

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

    keyp_implant = (ikeyp == which_keyp)

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

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

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

    keyp_implant = (ikeyp == which_keyp)

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

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

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

    keyp_implant = (ikeyp == which_keyp)

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

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

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

    keyp_implant = (ikeyp == which_keyp)

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

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

        r = torch.cat([r,r_implant],dim=1)
    return r
//  [python]
def update(self):
        """Set and update iteration variables.
        """
        if self.ivars['istep'] == 0:
            X_norm = float(torch.norm(self.X))
            iX_norm = X_norm ** -1
            A_norm = float(torch.norm(_utils.matmul(self.A, self.X))) * iX_norm
            B_norm = float(torch.norm(_utils.matmul(self.B, self.X))) * iX_norm
            self.fvars['X_norm'] = X_norm
            self.fvars['A_norm'] = A_norm
            self.fvars['B_norm'] = B_norm
            self.ivars['iterations_left'] = self.iparams['niter']
            self.ivars['converged_count'] = 0
            self.ivars['converged_end'] = 0

        if self.method == 'ortho':
            self._update_ortho()
        else:
            self._update_basic()

        self.ivars['iterations_left'] = self.ivars['iterations_left'] - 1
        self.ivars['istep'] = self.ivars['istep'] + 1
//  [python]
def update(self):
        """Set and update iteration variables.
        """
        if self.ivars['istep'] == 0:
            X_norm = float(torch.linalg.norm(self.X))
            iX_norm = X_norm ** -1
            A_norm = float(torch.linalg.norm(_utils.matmul(self.A, self.X))) * iX_norm
            B_norm = float(torch.linalg.norm(_utils.matmul(self.B, self.X))) * iX_norm
            self.fvars['X_norm'] = X_norm
            self.fvars['A_norm'] = A_norm
            self.fvars['B_norm'] = B_norm
            self.ivars['iterations_left'] = self.iparams['niter']
            self.ivars['converged_count'] = 0
            self.ivars['converged_end'] = 0

        if self.method == 'ortho':
            self._update_ortho()
        else:
            self._update_basic()

        self.ivars['iterations_left'] = self.ivars['iterations_left'] - 1
        self.ivars['istep'] = self.ivars['istep'] + 1
//  [python]
def add_cosine_b(
    a: Tensor | LiftFlag[MergeSpace.BASE],
    b: Tensor | LiftFlag[MergeSpace.BASE],
    *,
    alpha: Hyper,
    **kwargs,
) -> Tensor | LiftFlag[MergeSpace.BASE]:
    similarity = torch.nn.functional.cosine_similarity(a, b, dim=0)
    dot_product = torch.sum(a * b)
    magnitude_similarity = dot_product / (torch.norm(a) * torch.norm(b))
    combined_similarity = (similarity + magnitude_similarity) / 2.0
    return add_cosine_generic(a, b, alpha, combined_similarity)
//  [python]
def add_cosine_b(
    a: Tensor | LiftFlag[MergeSpace.BASE],
    b: Tensor | LiftFlag[MergeSpace.BASE],
    *,
    alpha: Hyper,
    **kwargs,
) -> Tensor | LiftFlag[MergeSpace.BASE]:
    similarity = torch.nn.functional.cosine_similarity(a, b, dim=0)
    dot_product = torch.sum(a * b)
    magnitude_similarity = dot_product / (torch.linalg.norm(a) * torch.linalg.norm(b))
    combined_similarity = (similarity + magnitude_similarity) / 2.0
    return add_cosine_generic(a, b, alpha, combined_similarity)
//  [python]
def update_converged_count(self):
        """Determine the number of converged eigenpairs using backward stable
        convergence criterion, see discussion in Sec 4.3 of [DuerschEtal2018].

        Users may redefine this method for custom convergence criteria.
        """
        # (...) -> int
        prev_count = self.ivars['converged_count']
        tol = self.fparams['tol']
        A_norm = self.fvars['A_norm']
        B_norm = self.fvars['B_norm']
        E, X, R = self.E, self.X, self.R
        rerr = torch.norm(R, 2, (0, )) * (torch.norm(X, 2, (0, )) * (A_norm + E[:X.shape[-1]] * B_norm)) ** -1
        converged = rerr < tol
        count = 0
        for b in converged:
            if not b:
                # ignore convergence of following pairs to ensure
                # strict ordering of eigenpairs
                break
            count += 1
        assert count >= prev_count, 'the number of converged eigenpairs ' \
            '(was {}, got {}) cannot decrease'.format(prev_count, count)
        self.ivars['converged_count'] = count
        self.tvars['rerr'] = rerr
        return count
//  [python]
def update_converged_count(self):
        """Determine the number of converged eigenpairs using backward stable
        convergence criterion, see discussion in Sec 4.3 of [DuerschEtal2018].

        Users may redefine this method for custom convergence criteria.
        """
        # (...) -> int
        prev_count = self.ivars['converged_count']
        tol = self.fparams['tol']
        A_norm = self.fvars['A_norm']
        B_norm = self.fvars['B_norm']
        E, X, R = self.E, self.X, self.R
        rerr = torch.linalg.norm(R, 2, (0, )) * (torch.linalg.norm(X, 2, (0, )) * (A_norm + E[:X.shape[-1]] * B_norm)) ** -1
        converged = rerr < tol
        count = 0
        for b in converged:
            if not b:
                # ignore convergence of following pairs to ensure
                # strict ordering of eigenpairs
                break
            count += 1
        assert count >= prev_count, 'the number of converged eigenpairs ' \
            '(was {}, got {}) cannot decrease'.format(prev_count, count)
        self.ivars['converged_count'] = count
        self.tvars['rerr'] = rerr
        return count