jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_api():
    def f(x: torch.Tensor):
        return torch.sum(x)

    def fmulti(x: torch.Tensor):
        return torch.sum(x), torch.norm(x)

    with pytest.raises(ValueError):
        et.Problem("min", f)  # Missing solution_length

    prob_with_bounds_missing = et.Problem("min", f, solution_length=5)  # Missing (initial_)bounds
    with pytest.raises(RuntimeError):
        prob_with_bounds_missing.generate_batch(10)  # fails because `fill(...)` does not know the bounds

    with pytest.raises(RuntimeError):
        prob_with_bounds_missing.generate_values(10)  # fails because `fill(...)` does not know the bounds

    with pytest.raises(ValueError):
        # This should fail, because et.Problem cannot work with non-numeric dtypes
        # when solution_length and bounds are given
        et.Problem("min", f, solution_length=5, bounds=(-1.0, 1.0), dtype=object)

    prob = et.Problem("min", f, solution_length=5, bounds=(-1.0, 1.0))
    prob.ensure_numeric()
    prob.ensure_single_objective()

    with pytest.raises(ValueError):
        prob.ensure_unbounded()

    assert prob.objective_sense == "min"
    assert prob.senses == ["min"]

    assert prob.initial_lower_bounds is not None
    assert prob.initial_upper_bounds is not None
    testing.assert_allclose(prob.lower_bounds, prob.initial_lower_bounds, atol=0.00001)
    testing.assert_allclose(prob.upper_bounds, prob.initial_upper_bounds, atol=0.00001)

    prob2 = et.Problem("min", f, solution_length=5, initial_bounds=(-1.0, 1.0), bounds=(-5.0, 5.0))
    prob2.ensure_numeric()
    prob2.ensure_single_objective()

    with pytest.raises(ValueError):
        prob2.ensure_unbounded()

    assert prob2.objective_sense == "min"
    assert prob2.senses == ["min"]

    prob3 = et.Problem("max", f, solution_length=5, initial_bounds=(-1.0, 1.0))
    prob3.ensure_unbounded()
    prob3.ensure_numeric()
    prob3.ensure_single_objective()

    assert prob3.objective_sense == "max"
    assert prob3.senses == ["max"]

    prob4 = et.Problem(["min", "min"], fmulti, solution_length=5, initial_bounds=(-1.0, 1.0))
    prob4.ensure_unbounded()
    prob4.ensure_numeric()

    with pytest.raises(ValueError):
        prob4.ensure_single_objective()

    assert prob4.objective_sense == ["min", "min"]
    assert prob4.senses == ["min", "min"]
//  [python]
def test_api():
    def f(x: torch.Tensor):
        return torch.sum(x)

    def fmulti(x: torch.Tensor):
        return torch.sum(x), torch.linalg.norm(x)

    with pytest.raises(ValueError):
        et.Problem("min", f)  # Missing solution_length

    prob_with_bounds_missing = et.Problem("min", f, solution_length=5)  # Missing (initial_)bounds
    with pytest.raises(RuntimeError):
        prob_with_bounds_missing.generate_batch(10)  # fails because `fill(...)` does not know the bounds

    with pytest.raises(RuntimeError):
        prob_with_bounds_missing.generate_values(10)  # fails because `fill(...)` does not know the bounds

    with pytest.raises(ValueError):
        # This should fail, because et.Problem cannot work with non-numeric dtypes
        # when solution_length and bounds are given
        et.Problem("min", f, solution_length=5, bounds=(-1.0, 1.0), dtype=object)

    prob = et.Problem("min", f, solution_length=5, bounds=(-1.0, 1.0))
    prob.ensure_numeric()
    prob.ensure_single_objective()

    with pytest.raises(ValueError):
        prob.ensure_unbounded()

    assert prob.objective_sense == "min"
    assert prob.senses == ["min"]

    assert prob.initial_lower_bounds is not None
    assert prob.initial_upper_bounds is not None
    testing.assert_allclose(prob.lower_bounds, prob.initial_lower_bounds, atol=0.00001)
    testing.assert_allclose(prob.upper_bounds, prob.initial_upper_bounds, atol=0.00001)

    prob2 = et.Problem("min", f, solution_length=5, initial_bounds=(-1.0, 1.0), bounds=(-5.0, 5.0))
    prob2.ensure_numeric()
    prob2.ensure_single_objective()

    with pytest.raises(ValueError):
        prob2.ensure_unbounded()

    assert prob2.objective_sense == "min"
    assert prob2.senses == ["min"]

    prob3 = et.Problem("max", f, solution_length=5, initial_bounds=(-1.0, 1.0))
    prob3.ensure_unbounded()
    prob3.ensure_numeric()
    prob3.ensure_single_objective()

    assert prob3.objective_sense == "max"
    assert prob3.senses == ["max"]

    prob4 = et.Problem(["min", "min"], fmulti, solution_length=5, initial_bounds=(-1.0, 1.0))
    prob4.ensure_unbounded()
    prob4.ensure_numeric()

    with pytest.raises(ValueError):
        prob4.ensure_single_objective()

    assert prob4.objective_sense == ["min", "min"]
    assert prob4.senses == ["min", "min"]
//  [python]
def ee_pose_target(batch_size):
    torch.manual_seed(1)

    ee_pos_mid = torch.Tensor([0.6, 0.0, 0.5])
    ee_pos_range = torch.Tensor([0.1, 0.2, 0.2])
    ee_quat_mid = torch.Tensor([0.9383, 0.3442, -0.0072, -0.0318])
    ee_quat_range = torch.Tensor([0.5, 0.5, 0.5])

    ee_pos_dev = ee_pos_range * torch.randn(batch_size, 3)
    ee_rot_dev = ee_quat_range * torch.randn(batch_size, 3)

    ee_pos_target = ee_pos_mid + ee_pos_dev
    ee_quat_target = (
        th.SO3.unit_quaternion_to_SO3(
            ee_quat_mid / th.norm(ee_quat_mid)
        ).compose(th.SO3().exp_map(ee_rot_dev))
    ).to_quaternion()
    ee_pose_target = th.SE3(
        x_y_z_quaternion=torch.cat([ee_pos_target, ee_quat_target], dim=-1),
        name="ee_pose_target",
    )

    return ee_pose_target
//  [python]
def ee_pose_target(batch_size):
    torch.manual_seed(1)

    ee_pos_mid = torch.Tensor([0.6, 0.0, 0.5])
    ee_pos_range = torch.Tensor([0.1, 0.2, 0.2])
    ee_quat_mid = torch.Tensor([0.9383, 0.3442, -0.0072, -0.0318])
    ee_quat_range = torch.Tensor([0.5, 0.5, 0.5])

    ee_pos_dev = ee_pos_range * torch.randn(batch_size, 3)
    ee_rot_dev = ee_quat_range * torch.randn(batch_size, 3)

    ee_pos_target = ee_pos_mid + ee_pos_dev
    ee_quat_target = (
        th.SO3.unit_quaternion_to_SO3(
            ee_quat_mid / torch.linalg.norm(ee_quat_mid)
        ).compose(th.SO3().exp_map(ee_rot_dev))
    ).to_quaternion()
    ee_pose_target = th.SE3(
        x_y_z_quaternion=torch.cat([ee_pos_target, ee_quat_target], dim=-1),
        name="ee_pose_target",
    )

    return ee_pose_target
//  [python]
def __init__(self, opt):
        super(AvatarNet, self).__init__()
        self.opt = opt

        self.random_style = opt.get('random_style', False)
        self.with_viewdirs = opt.get('with_viewdirs', True)

        # init canonical gausssian model
        self.max_sh_degree = 0
        self.cano_gaussian_model = GaussianModel(sh_degree = self.max_sh_degree)
        cano_smpl_map = cv.imread(config.opt['train']['data']['data_dir'] + '/smpl_pos_map/cano_smpl_pos_map.exr', cv.IMREAD_UNCHANGED)
        self.cano_smpl_map = torch.from_numpy(cano_smpl_map).to(torch.float32).to(config.device)
        self.cano_smpl_mask = torch.norm(self.cano_smpl_map, dim = -1) > 0.
        self.init_points = self.cano_smpl_map[self.cano_smpl_mask]
        self.lbs = torch.from_numpy(np.load(config.opt['train']['data']['data_dir'] + '/smpl_pos_map/init_pts_lbs.npy')).to(torch.float32).to(config.device)
        self.cano_gaussian_model.create_from_pcd(self.init_points, torch.rand_like(self.init_points), spatial_lr_scale = 2.5)

        self.color_net = DualStyleUNet(inp_size = 512, inp_ch = 3, out_ch = 3, out_size = 1024, style_dim = 512, n_mlp = 2)
        self.position_net = DualStyleUNet(inp_size = 512, inp_ch = 3, out_ch = 3, out_size = 1024, style_dim = 512, n_mlp = 2)
        self.other_net = DualStyleUNet(inp_size = 512, inp_ch = 3, out_ch = 8, out_size = 1024, style_dim = 512, n_mlp = 2)

        self.color_style = torch.ones([1, self.color_net.style_dim], dtype=torch.float32, device=config.device) / np.sqrt(self.color_net.style_dim)
        self.position_style = torch.ones([1, self.position_net.style_dim], dtype=torch.float32, device=config.device) / np.sqrt(self.position_net.style_dim)
        self.other_style = torch.ones([1, self.other_net.style_dim], dtype=torch.float32, device=config.device) / np.sqrt(self.other_net.style_dim)

        if self.with_viewdirs:
            cano_nml_map = cv.imread(config.opt['train']['data']['data_dir'] + '/smpl_pos_map/cano_smpl_nml_map.exr', cv.IMREAD_UNCHANGED)
            self.cano_nml_map = torch.from_numpy(cano_nml_map).to(torch.float32).to(config.device)
            self.cano_nmls = self.cano_nml_map[self.cano_smpl_mask]
            self.viewdir_net = nn.Sequential(
                nn.Conv2d(1, 64, 4, 2, 1),
                nn.LeakyReLU(0.2, inplace = True),
                nn.Conv2d(64, 128, 4, 2, 1)
            )
//  [python]
def __init__(self, opt):
        super(AvatarNet, self).__init__()
        self.opt = opt

        self.random_style = opt.get('random_style', False)
        self.with_viewdirs = opt.get('with_viewdirs', True)

        # init canonical gausssian model
        self.max_sh_degree = 0
        self.cano_gaussian_model = GaussianModel(sh_degree = self.max_sh_degree)
        cano_smpl_map = cv.imread(config.opt['train']['data']['data_dir'] + '/smpl_pos_map/cano_smpl_pos_map.exr', cv.IMREAD_UNCHANGED)
        self.cano_smpl_map = torch.from_numpy(cano_smpl_map).to(torch.float32).to(config.device)
        self.cano_smpl_mask = torch.linalg.norm(self.cano_smpl_map, dim = -1) > 0.
        self.init_points = self.cano_smpl_map[self.cano_smpl_mask]
        self.lbs = torch.from_numpy(np.load(config.opt['train']['data']['data_dir'] + '/smpl_pos_map/init_pts_lbs.npy')).to(torch.float32).to(config.device)
        self.cano_gaussian_model.create_from_pcd(self.init_points, torch.rand_like(self.init_points), spatial_lr_scale = 2.5)

        self.color_net = DualStyleUNet(inp_size = 512, inp_ch = 3, out_ch = 3, out_size = 1024, style_dim = 512, n_mlp = 2)
        self.position_net = DualStyleUNet(inp_size = 512, inp_ch = 3, out_ch = 3, out_size = 1024, style_dim = 512, n_mlp = 2)
        self.other_net = DualStyleUNet(inp_size = 512, inp_ch = 3, out_ch = 8, out_size = 1024, style_dim = 512, n_mlp = 2)

        self.color_style = torch.ones([1, self.color_net.style_dim], dtype=torch.float32, device=config.device) / np.sqrt(self.color_net.style_dim)
        self.position_style = torch.ones([1, self.position_net.style_dim], dtype=torch.float32, device=config.device) / np.sqrt(self.position_net.style_dim)
        self.other_style = torch.ones([1, self.other_net.style_dim], dtype=torch.float32, device=config.device) / np.sqrt(self.other_net.style_dim)

        if self.with_viewdirs:
            cano_nml_map = cv.imread(config.opt['train']['data']['data_dir'] + '/smpl_pos_map/cano_smpl_nml_map.exr', cv.IMREAD_UNCHANGED)
            self.cano_nml_map = torch.from_numpy(cano_nml_map).to(torch.float32).to(config.device)
            self.cano_nmls = self.cano_nml_map[self.cano_smpl_mask]
            self.viewdir_net = nn.Sequential(
                nn.Conv2d(1, 64, 4, 2, 1),
                nn.LeakyReLU(0.2, inplace = True),
                nn.Conv2d(64, 128, 4, 2, 1)
            )
//  [python]
def _find_dq(
    target_norm: float,
    step_size: float,
    coords_x: torch.Tensor,
    grad_q: torch.Tensor,
    hess_q: torch.Tensor,
    ic: tico.ic.IC,
    rtol: float = 0.001,
) -> torch.Tensor:
    """

    Notes:
        Based on geomeTRIC, which is based on the iterative formula from Hebden (1973),
        equation 5.2.10 in "Practical methods of optimization" by Fletcher

    Returns:
        The found value of dq.
    """
    dq, dq_prime = _compute_dq(step_size, coords_x, grad_q, hess_q, ic)
    dq_norm = torch.norm(dq)

    if dq_norm < target_norm:
        return dq

    dq_norm_last, dq_norm_best, dq_best = 0.0, dq_norm, dq.clone()

    iteration = 0

    while True:
        step_size += (1 - dq_norm / target_norm) * (dq_norm / dq_prime)

        dq, dq_prime = _compute_dq(step_size, coords_x, grad_q, hess_q, ic)
        dq_norm = torch.norm(dq)

        if torch.abs((dq_norm - target_norm) / target_norm) < rtol:
            return dq
        if iteration > 10 and torch.abs(dq_norm_last - dq_norm) / dq_norm < rtol:
            return dq

        iteration += 1

        dq_norm_last = dq_norm

        if dq_norm < dq_norm_best:
            dq_norm_best, dq_best = dq_norm, dq.clone()

        # handle infinite loops... following geomeTRIC
        if iteration % 100 == 99:
            step_size += random.random() * iteration / 100
        if iteration % 1000 == 999:
            return dq_best
//  [python]
def _find_dq(
    target_norm: float,
    step_size: float,
    coords_x: torch.Tensor,
    grad_q: torch.Tensor,
    hess_q: torch.Tensor,
    ic: tico.ic.IC,
    rtol: float = 0.001,
) -> torch.Tensor:
    """

    Notes:
        Based on geomeTRIC, which is based on the iterative formula from Hebden (1973),
        equation 5.2.10 in "Practical methods of optimization" by Fletcher

    Returns:
        The found value of dq.
    """
    dq, dq_prime = _compute_dq(step_size, coords_x, grad_q, hess_q, ic)
    dq_norm = torch.linalg.norm(dq)

    if dq_norm < target_norm:
        return dq

    dq_norm_last, dq_norm_best, dq_best = 0.0, dq_norm, dq.clone()

    iteration = 0

    while True:
        step_size += (1 - dq_norm / target_norm) * (dq_norm / dq_prime)

        dq, dq_prime = _compute_dq(step_size, coords_x, grad_q, hess_q, ic)
        dq_norm = torch.linalg.norm(dq)

        if torch.abs((dq_norm - target_norm) / target_norm) < rtol:
            return dq
        if iteration > 10 and torch.abs(dq_norm_last - dq_norm) / dq_norm < rtol:
            return dq

        iteration += 1

        dq_norm_last = dq_norm

        if dq_norm < dq_norm_best:
            dq_norm_best, dq_best = dq_norm, dq.clone()

        # handle infinite loops... following geomeTRIC
        if iteration % 100 == 99:
            step_size += random.random() * iteration / 100
        if iteration % 1000 == 999:
            return dq_best
//  [python]
def forward_one_pass(self, items):
        total_loss = 0
        batch_losses = {}

        """ random sampling """
        if 'nerf_random' in items:
            items.update(items['nerf_random'])
            render_output = self.network.render(items, depth_guided_sampling = self.opt['train']['depth_guided_sampling'])

            # color loss
            if 'rgb_map' in render_output:
                color_loss = torch.nn.L1Loss()(render_output['rgb_map'], items['color_gt'])
                total_loss += self.loss_weight['color'] * color_loss
                batch_losses.update({
                    'color_loss_random': color_loss.item()
                })

            # mask loss
            if 'acc_map' in render_output:
                mask_loss = torch.nn.L1Loss()(render_output['acc_map'], items['mask_gt'])
                total_loss += self.loss_weight['mask'] * mask_loss
                batch_losses.update({
                    'mask_loss_random': mask_loss.item()
                })

            # eikonal loss
            if 'normal' in render_output:
                eikonal_loss = ((torch.norm(render_output['normal'], dim = -1) - 1.) ** 2).mean()
                total_loss += self.loss_weight['eikonal'] * eikonal_loss
                batch_losses.update({
                    'eikonal_loss': eikonal_loss.item()
                })

        self.zero_grad()
        total_loss.backward()
        self.step()

        return total_loss, batch_losses
//  [python]
def forward_one_pass(self, items):
        total_loss = 0
        batch_losses = {}

        """ random sampling """
        if 'nerf_random' in items:
            items.update(items['nerf_random'])
            render_output = self.network.render(items, depth_guided_sampling = self.opt['train']['depth_guided_sampling'])

            # color loss
            if 'rgb_map' in render_output:
                color_loss = torch.nn.L1Loss()(render_output['rgb_map'], items['color_gt'])
                total_loss += self.loss_weight['color'] * color_loss
                batch_losses.update({
                    'color_loss_random': color_loss.item()
                })

            # mask loss
            if 'acc_map' in render_output:
                mask_loss = torch.nn.L1Loss()(render_output['acc_map'], items['mask_gt'])
                total_loss += self.loss_weight['mask'] * mask_loss
                batch_losses.update({
                    'mask_loss_random': mask_loss.item()
                })

            # eikonal loss
            if 'normal' in render_output:
                eikonal_loss = ((torch.linalg.norm(render_output['normal'], dim = -1) - 1.) ** 2).mean()
                total_loss += self.loss_weight['eikonal'] * eikonal_loss
                batch_losses.update({
                    'eikonal_loss': eikonal_loss.item()
                })

        self.zero_grad()
        total_loss.backward()
        self.step()

        return total_loss, batch_losses
//  [python]
def _line_search(
    trust: float,
    step_size: float,
    coords_x: torch.Tensor,
    grad_q: torch.Tensor,
    hess_q: torch.Tensor,
    dq: torch.Tensor,
    ic: tico.ic.IC,
    max_trials: int = 4,
) -> torch.Tensor:
    """Perform a line search to find a step size that matches the trust radius."""
    assert max_trials >= 1

    @functools.lru_cache
    def evaluate_dx(trial: float) -> tuple[torch.Tensor, bool]:
        dq_trial = _find_dq(trial, step_size, coords_x, grad_q, hess_q, ic)

        coords_x_new, _, q_converged = ic.dq_to_x(coords_x, dq_trial)
        dx = tico.utils.compute_rmsd(coords_x_new, coords_x)[0]

        return dx, q_converged

    def evaluate(trial: float, target_dx: float) -> tuple[float, bool]:
        if trial == 0.0:
            return -trust, False

        dx, q_converged = evaluate_dx(trial)
        return (dx - target_dx), q_converged

    dq_norm = torch.norm(dq).detach().cpu().item()

    is_trial_valid = True  # could we map dq to new cartesian coordinates?
    target = trust

    for i in range(max_trials):
        dq_norm, trials, converged = tico.utils.brent(
            evaluate, 0.0, dq_norm, target, cvg=0.1, args=(target,)
        )

        best_trials = sorted(
            ((trial, delta) for trial, delta, _ in trials if delta < 0),
            key=lambda x: x[1],
        )
        is_trial_valid = trials[-1][-1]

        if i == 0:
            if not converged and len(best_trials) > 0:
                dq_norm = best_trials[0][0]
                break
            elif not is_trial_valid:
                pass
            else:
                break
        elif is_trial_valid:
            break

        target *= 0.5

    if not is_trial_valid:  # could not find a valid step that lets x converge...
        raise RuntimeError("line search failed to find a valid step")

    return _find_dq(dq_norm, step_size, coords_x, grad_q, hess_q, ic)
//  [python]
def _line_search(
    trust: float,
    step_size: float,
    coords_x: torch.Tensor,
    grad_q: torch.Tensor,
    hess_q: torch.Tensor,
    dq: torch.Tensor,
    ic: tico.ic.IC,
    max_trials: int = 4,
) -> torch.Tensor:
    """Perform a line search to find a step size that matches the trust radius."""
    assert max_trials >= 1

    @functools.lru_cache
    def evaluate_dx(trial: float) -> tuple[torch.Tensor, bool]:
        dq_trial = _find_dq(trial, step_size, coords_x, grad_q, hess_q, ic)

        coords_x_new, _, q_converged = ic.dq_to_x(coords_x, dq_trial)
        dx = tico.utils.compute_rmsd(coords_x_new, coords_x)[0]

        return dx, q_converged

    def evaluate(trial: float, target_dx: float) -> tuple[float, bool]:
        if trial == 0.0:
            return -trust, False

        dx, q_converged = evaluate_dx(trial)
        return (dx - target_dx), q_converged

    dq_norm = torch.linalg.norm(dq).detach().cpu().item()

    is_trial_valid = True  # could we map dq to new cartesian coordinates?
    target = trust

    for i in range(max_trials):
        dq_norm, trials, converged = tico.utils.brent(
            evaluate, 0.0, dq_norm, target, cvg=0.1, args=(target,)
        )

        best_trials = sorted(
            ((trial, delta) for trial, delta, _ in trials if delta < 0),
            key=lambda x: x[1],
        )
        is_trial_valid = trials[-1][-1]

        if i == 0:
            if not converged and len(best_trials) > 0:
                dq_norm = best_trials[0][0]
                break
            elif not is_trial_valid:
                pass
            else:
                break
        elif is_trial_valid:
            break

        target *= 0.5

    if not is_trial_valid:  # could not find a valid step that lets x converge...
        raise RuntimeError("line search failed to find a valid step")

    return _find_dq(dq_norm, step_size, coords_x, grad_q, hess_q, ic)
//  [python]
def lamb(params: List[Tensor],
         grads: List[Tensor],
         exp_avgs: List[Tensor],
         exp_avg_sqs: List[Tensor],
         state_steps: List[Tensor],
         beta1: float,
         beta2: float,
         lr: float,
         weight_decay: float,
         eps: float):

  if not all(isinstance(t, torch.Tensor) for t in state_steps):
    raise RuntimeError(
        'API has changed, `state_steps` argument must contain a list of' +
        ' singleton tensors')

  for i, param in enumerate(params):
    grad = grads[i]
    exp_avg = exp_avgs[i]
    exp_avg_sq = exp_avg_sqs[i]
    step_t = state_steps[i]

    # Update step.
    step_t += 1

    # Decay the first and second moment running average coefficient.
    exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
    exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)

    step = step_t.item()

    bias_correction1 = 1 - beta1**step
    bias_correction2 = 1 - beta2**step

    bias_correction2_sqrt = math.sqrt(bias_correction2)
    denom = (exp_avg_sq.sqrt() / bias_correction2_sqrt).add_(eps)

    update = exp_avg / denom
    update.div_(bias_correction1)
    update.add_(weight_decay * param)

    # Scale updates by trust ratio.
    param_norm = param.norm()
    update_norm = update.norm()

    # Set trust_ratio to 1 in case where parameters would never be updated.
    if param_norm == 0. or update_norm == 0.:
      trust_ratio = 1.
    else:
      trust_ratio = param_norm / update_norm

    param.add_(update, alpha=-lr * trust_ratio)
//  [python]
def lamb(params: List[Tensor],
         grads: List[Tensor],
         exp_avgs: List[Tensor],
         exp_avg_sqs: List[Tensor],
         state_steps: List[Tensor],
         beta1: float,
         beta2: float,
         lr: float,
         weight_decay: float,
         eps: float):

  if not all(isinstance(t, torch.Tensor) for t in state_steps):
    raise RuntimeError(
        'API has changed, `state_steps` argument must contain a list of' +
        ' singleton tensors')

  for i, param in enumerate(params):
    grad = grads[i]
    exp_avg = exp_avgs[i]
    exp_avg_sq = exp_avg_sqs[i]
    step_t = state_steps[i]

    # Update step.
    step_t += 1

    # Decay the first and second moment running average coefficient.
    exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
    exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)

    step = step_t.item()

    bias_correction1 = 1 - beta1**step
    bias_correction2 = 1 - beta2**step

    bias_correction2_sqrt = math.sqrt(bias_correction2)
    denom = (exp_avg_sq.sqrt() / bias_correction2_sqrt).add_(eps)

    update = exp_avg / denom
    update.div_(bias_correction1)
    update.add_(weight_decay * param)

    # Scale updates by trust ratio.
    param_norm = torch.linalg.norm(param)
    update_norm = torch.linalg.norm(update)

    # Set trust_ratio to 1 in case where parameters would never be updated.
    if param_norm == 0. or update_norm == 0.:
      trust_ratio = 1.
    else:
      trust_ratio = param_norm / update_norm

    param.add_(update, alpha=-lr * trust_ratio)
//  [python]
def is_close_gt(anchor, gt, stride_lst, max_dist=2.0, alpha=2.):
    """Calculate distance ratio of box1 and box2 in batch for larger stride
        anchors dist/stride to promote the survive of large distance match
    Args:
        anchor (Tensor): box with the shape [L, 2]
        gt (Tensor): box with the shape [N, M2, 4]
    Return:
        dist (Tensor): dist ratio between box1 and box2 with the shape [N, M1, M2]
    """
    center1 = anchor.unsqueeze(0)
    center2 = (gt[..., :2] + gt[..., -2:]) / 2.
    center1 = center1.unsqueeze(1)  # [N, M1, 2] -> [N, 1, M1, 2]
    center2 = center2.unsqueeze(2)  # [N, M2, 2] -> [N, M2, 1, 2]

    stride = torch.cat([
        torch.full([x], 32 / pow(2, idx)) for idx, x in enumerate(stride_lst)
    ]).unsqueeze(0).unsqueeze(0)
    dist = torch.norm(center1 - center2, p=2, dim=-1) / stride
    dist_ratio = dist
    dist_ratio[dist < max_dist] = 1.
    dist_ratio[dist >= max_dist] = 0.
    return dist_ratio
//  [python]
def is_close_gt(anchor, gt, stride_lst, max_dist=2.0, alpha=2.):
    """Calculate distance ratio of box1 and box2 in batch for larger stride
        anchors dist/stride to promote the survive of large distance match
    Args:
        anchor (Tensor): box with the shape [L, 2]
        gt (Tensor): box with the shape [N, M2, 4]
    Return:
        dist (Tensor): dist ratio between box1 and box2 with the shape [N, M1, M2]
    """
    center1 = anchor.unsqueeze(0)
    center2 = (gt[..., :2] + gt[..., -2:]) / 2.
    center1 = center1.unsqueeze(1)  # [N, M1, 2] -> [N, 1, M1, 2]
    center2 = center2.unsqueeze(2)  # [N, M2, 2] -> [N, M2, 1, 2]

    stride = torch.cat([
        torch.full([x], 32 / pow(2, idx)) for idx, x in enumerate(stride_lst)
    ]).unsqueeze(0).unsqueeze(0)
    dist = torch.linalg.norm(center1 - center2, p=2, axis=-1) / stride
    dist_ratio = dist
    dist_ratio[dist < max_dist] = 1.
    dist_ratio[dist >= max_dist] = 0.
    return dist_ratio
//  [python]
def _optimize(x, y, penalty, alpha):
        B, C = x.shape[0], x.shape[1]
        x = x.detach()
        y = y.clone()
        y.requires_grad = True
        opt = torch.optim.LBFGS([y], lr=1.0, max_iter=100, max_eval=None,
            tolerance_grad=1e-05, tolerance_change=1e-09, history_size=5, line_search_fn=None)

        def reevaluate():
            opt.zero_grad()
            z = torch.norm((y.view(B, C, 1) - x.view(B, C, -1)), dim=1)
            f = penalty.phi(z, alpha).sum()
            f.backward()
            return f

        # optimize to convergence (unlike step for other optimizers)
        opt.step(reevaluate)

        return y
//  [python]
def _optimize(x, y, penalty, alpha):
        B, C = x.shape[0], x.shape[1]
        x = x.detach()
        y = y.clone()
        y.requires_grad = True
        opt = torch.optim.LBFGS([y], lr=1.0, max_iter=100, max_eval=None,
            tolerance_grad=1e-05, tolerance_change=1e-09, history_size=5, line_search_fn=None)

        def reevaluate():
            opt.zero_grad()
            z = torch.linalg.norm((y.view(B, C, 1) - x.view(B, C, -1)), dim=1)
            f = penalty.phi(z, alpha).sum()
            f.backward()
            return f

        # optimize to convergence (unlike step for other optimizers)
        opt.step(reevaluate)

        return y
//  [python]
def gen_rays_at(self, img_idx, resolution_level=1):
        """
        Generate rays at world space from one camera.
        """
        l = resolution_level
        tx = torch.linspace(0, self.W - 1, self.W // l)
        ty = torch.linspace(0, self.H - 1, self.H // l)
        pixels_x, pixels_y = torch.meshgrid(tx, ty)
        p = torch.stack([pixels_x, pixels_y, torch.ones_like(pixels_y)], dim=-1) # W, H, 3
        # Camera
        if self.camera_trainable:
            intrinsic_inv, pose = self.dynamic_paras_to_mat(img_idx)
        else:
            if self.is_monocular:
                intrinsic_inv = self.intrinsics_all_inv[0]
            else:
                intrinsic_inv = self.intrinsics_all_inv[img_idx]
            pose = self.poses_all[img_idx]
        p = torch.matmul(intrinsic_inv[None, None, :3, :3], p[:, :, :, None]).squeeze() # W, H, 3
        rays_v = p / torch.norm(p, p=2, dim=-1, keepdim=True) # W, H, 3
        rays_v = torch.matmul(pose[None, None, :3, :3], rays_v[:, :, :, None]).squeeze() # W, H, 3
        rays_o = pose[None, None, :3, 3].expand(rays_v.shape) # W, H, 3
        return rays_o.transpose(0, 1), rays_v.transpose(0, 1)
//  [python]
def gen_rays_at(self, img_idx, resolution_level=1):
        """
        Generate rays at world space from one camera.
        """
        l = resolution_level
        tx = torch.linspace(0, self.W - 1, self.W // l)
        ty = torch.linspace(0, self.H - 1, self.H // l)
        pixels_x, pixels_y = torch.meshgrid(tx, ty)
        p = torch.stack([pixels_x, pixels_y, torch.ones_like(pixels_y)], dim=-1) # W, H, 3
        # Camera
        if self.camera_trainable:
            intrinsic_inv, pose = self.dynamic_paras_to_mat(img_idx)
        else:
            if self.is_monocular:
                intrinsic_inv = self.intrinsics_all_inv[0]
            else:
                intrinsic_inv = self.intrinsics_all_inv[img_idx]
            pose = self.poses_all[img_idx]
        p = torch.matmul(intrinsic_inv[None, None, :3, :3], p[:, :, :, None]).squeeze() # W, H, 3
        rays_v = p / torch.linalg.norm(p, ord=2, dim=-1, keepdim=True) # W, H, 3
        rays_v = torch.matmul(pose[None, None, :3, :3], rays_v[:, :, :, None]).squeeze() # W, H, 3
        rays_o = pose[None, None, :3, 3].expand(rays_v.shape) # W, H, 3
        return rays_o.transpose(0, 1), rays_v.transpose(0, 1)
//  [python]
def gen_random_rays_at(self, img_idx, batch_size):
        """
        Generate random rays at world space from one camera.
        """
        pixels_x = torch.randint(low=0, high=self.W, size=[batch_size])
        pixels_y = torch.randint(low=0, high=self.H, size=[batch_size])
        color = self.images[img_idx].to(self.device)[(pixels_y, pixels_x)] # batch_size, 3
        mask = self.masks[img_idx].to(self.device)[(pixels_y, pixels_x)] # batch_size, 3
        p = torch.stack([pixels_x, pixels_y, torch.ones_like(pixels_y)], dim=-1).to(self.dtype) # batch_size, 3
        # Camera
        if self.camera_trainable:
            intrinsic_inv, pose = self.dynamic_paras_to_mat(img_idx)
        else:
            if self.is_monocular:
                intrinsic_inv = self.intrinsics_all_inv[0]
            else:
                intrinsic_inv = self.intrinsics_all_inv[img_idx]
            pose = self.poses_all[img_idx]
        p = torch.matmul(intrinsic_inv[None, :3, :3], p[:, :, None]).squeeze() # batch_size, 3
        rays_v = p / torch.norm(p, p=2, dim=-1, keepdim=True) # batch_size, 3
        rays_v = torch.matmul(pose[None, :3, :3], rays_v[:, :, None]).squeeze() # batch_size, 3
        rays_o = pose[None, :3, 3].expand(rays_v.shape) # batch_size, 3
        return torch.cat([rays_o, rays_v, color, mask[:, :1]], dim=-1)
//  [python]
def gen_random_rays_at(self, img_idx, batch_size):
        """
        Generate random rays at world space from one camera.
        """
        pixels_x = torch.randint(low=0, high=self.W, size=[batch_size])
        pixels_y = torch.randint(low=0, high=self.H, size=[batch_size])
        color = self.images[img_idx].to(self.device)[(pixels_y, pixels_x)] # batch_size, 3
        mask = self.masks[img_idx].to(self.device)[(pixels_y, pixels_x)] # batch_size, 3
        p = torch.stack([pixels_x, pixels_y, torch.ones_like(pixels_y)], dim=-1).to(self.dtype) # batch_size, 3
        # Camera
        if self.camera_trainable:
            intrinsic_inv, pose = self.dynamic_paras_to_mat(img_idx)
        else:
            if self.is_monocular:
                intrinsic_inv = self.intrinsics_all_inv[0]
            else:
                intrinsic_inv = self.intrinsics_all_inv[img_idx]
            pose = self.poses_all[img_idx]
        p = torch.matmul(intrinsic_inv[None, :3, :3], p[:, :, None]).squeeze() # batch_size, 3
        rays_v = p / torch.linalg.norm(p, ord=2, dim=-1, keepdim=True) # batch_size, 3
        rays_v = torch.matmul(pose[None, :3, :3], rays_v[:, :, None]).squeeze() # batch_size, 3
        rays_o = pose[None, :3, 3].expand(rays_v.shape) # batch_size, 3
        return torch.cat([rays_o, rays_v, color, mask[:, :1]], dim=-1)
//  [python]
def forward(self, x):
        x = x.view(x.shape[0], -1, self.param_channels)
        mag = torch.norm(x, dim=-1)
        mag = self.act(mag * self.inner_fac) * self.outer_fac
        x = torch.nn.functional.normalize(x, dim=-1) * mag[..., None]
        return x
//  [python]
def forward(self, x):
        x = x.view(x.shape[0], -1, self.param_channels)
        mag = torch.linalg.norm(x, dim=-1)
        mag = self.act(mag * self.inner_fac) * self.outer_fac
        x = torch.nn.functional.normalize(x, dim=-1) * mag[..., None]
        return x
//  [python]
def gen_rays_between(self, idx_0, idx_1, ratio, resolution_level=1):
        """
        Interpolate pose between two cameras.
        """
        l = resolution_level
        tx = torch.linspace(0, self.W - 1, self.W // l)
        ty = torch.linspace(0, self.H - 1, self.H // l)
        pixels_x, pixels_y = torch.meshgrid(tx, ty)
        p = torch.stack([pixels_x, pixels_y, torch.ones_like(pixels_y)], dim=-1) # W, H, 3
        p = torch.matmul(self.intrinsics_all_inv[0, None, None, :3, :3], p[:, :, :, None]).squeeze() # W, H, 3
        rays_v = p / torch.norm(p, p=2, dim=-1, keepdim=True) # W, H, 3
        trans = self.pose_all[idx_0, :3, 3] * (1.0 - ratio) + self.pose_all[idx_1, :3, 3] * ratio
        pose_0 = self.pose_all[idx_0].detach().cpu().numpy()
        pose_1 = self.pose_all[idx_1].detach().cpu().numpy()
        pose_0 = np.linalg.inv(pose_0)
        pose_1 = np.linalg.inv(pose_1)
        rot_0 = pose_0[:3, :3]
        rot_1 = pose_1[:3, :3]
        rots = Rot.from_matrix(np.stack([rot_0, rot_1]))
        key_times = [0, 1]
        slerp = Slerp(key_times, rots)
        rot = slerp(ratio)
        pose = np.diag([1.0, 1.0, 1.0, 1.0])
        pose = pose.astype(np.float32)
        pose[:3, :3] = rot.as_matrix()
        pose[:3, 3] = ((1.0 - ratio) * pose_0 + ratio * pose_1)[:3, 3]
        pose = np.linalg.inv(pose)
        rot = torch.from_numpy(pose[:3, :3]).cuda()
        trans = torch.from_numpy(pose[:3, 3]).cuda()
        rays_v = torch.matmul(rot[None, None, :3, :3], rays_v[:, :, :, None]).squeeze() # W, H, 3
        rays_o = trans[None, None, :3].expand(rays_v.shape) # W, H, 3
        return rays_o.transpose(0, 1), rays_v.transpose(0, 1)
//  [python]
def gen_rays_between(self, idx_0, idx_1, ratio, resolution_level=1):
        """
        Interpolate pose between two cameras.
        """
        l = resolution_level
        tx = torch.linspace(0, self.W - 1, self.W // l)
        ty = torch.linspace(0, self.H - 1, self.H // l)
        pixels_x, pixels_y = torch.meshgrid(tx, ty)
        p = torch.stack([pixels_x, pixels_y, torch.ones_like(pixels_y)], dim=-1) # W, H, 3
        p = torch.matmul(self.intrinsics_all_inv[0, None, None, :3, :3], p[:, :, :, None]).squeeze() # W, H, 3
        rays_v = p / torch.linalg.norm(p, ord=2, dim=-1, keepdim=True) # W, H, 3
        trans = self.pose_all[idx_0, :3, 3] * (1.0 - ratio) + self.pose_all[idx_1, :3, 3] * ratio
        pose_0 = self.pose_all[idx_0].detach().cpu().numpy()
        pose_1 = self.pose_all[idx_1].detach().cpu().numpy()
        pose_0 = np.linalg.inv(pose_0)
        pose_1 = np.linalg.inv(pose_1)
        rot_0 = pose_0[:3, :3]
        rot_1 = pose_1[:3, :3]
        rots = Rot.from_matrix(np.stack([rot_0, rot_1]))
        key_times = [0, 1]
        slerp = Slerp(key_times, rots)
        rot = slerp(ratio)
        pose = np.diag([1.0, 1.0, 1.0, 1.0])
        pose = pose.astype(np.float32)
        pose[:3, :3] = rot.as_matrix()
        pose[:3, 3] = ((1.0 - ratio) * pose_0 + ratio * pose_1)[:3, 3]
        pose = np.linalg.inv(pose)
        rot = torch.from_numpy(pose[:3, :3]).cuda()
        trans = torch.from_numpy(pose[:3, 3]).cuda()
        rays_v = torch.matmul(rot[None, None, :3, :3], rays_v[:, :, :, None]).squeeze() # W, H, 3
        rays_o = trans[None, None, :3].expand(rays_v.shape) # W, H, 3
        return rays_o.transpose(0, 1), rays_v.transpose(0, 1)
//  [python]
def forward(ctx, x, penalty, alpha=1.0, restarts=0, hess_reg=1.0e-16):
        assert len(x.shape) >= 3
        assert alpha > 0.0
        assert restarts >= 0

        B, C = x.shape[0], x.shape[1]
        with torch.no_grad():
            if penalty.is_convex:
                y = torch.mean(x.view(B, C, -1), dim=2)
                if penalty is not Quadratic:
                    y = RobustVectorPool2dFcn._optimize(x, y, penalty, alpha)
            else:
                y_mean = torch.mean(x.view(B, C, -1), dim=2)
                y = RobustVectorPool2dFcn._optimize(x, y_mean, penalty, alpha)
                f = penalty.phi(torch.norm((y.view(B, C, 1) - x.view(B, C, -1)), dim=1), alpha).view(B, -1).sum(-1)

                y_median, _ = torch.median(x.view(B, C, -1), dim=2)
                y_median = RobustVectorPool2dFcn._optimize(x, y_median, penalty, alpha)
                f_median = penalty.phi(torch.norm((y_median.view(B, C, 1) - x.view(B, C, -1)), dim=1), alpha).view(B, -1).sum(-1)
                y = torch.where((f <= f_median).view(B, 1), y, y_median)
                f = torch.minimum(f, f_median)

                if restarts > 0:
                    # choose points evenly spaced in batch (can be randomly chosen but that affects reproducibility)
                    for i in torch.linspace(start=0, end=x.view(B, C, -1).shape[2] - 1, steps=restarts):
                        y_init = x.view(B, C, -1)[:, :, int(i.item())]
                        y_final = RobustVectorPool2dFcn._optimize(x, y_init, penalty, alpha)
                        f_final = penalty.phi(torch.norm((y_final.view(B, C, 1) - x.view(B, C, -1)), dim=1), alpha).view(B, -1).sum(-1)
                        y = torch.where((f < f_final).view(B, 1), y, y_final)
                        f = torch.minimum(f, f_final)

        ctx.save_for_backward(x, y)
        ctx.penalty = penalty
        ctx.alpha = alpha
        ctx.hess_reg = hess_reg

        return y
//  [python]
def forward(ctx, x, penalty, alpha=1.0, restarts=0, hess_reg=1.0e-16):
        assert len(x.shape) >= 3
        assert alpha > 0.0
        assert restarts >= 0

        B, C = x.shape[0], x.shape[1]
        with torch.no_grad():
            if penalty.is_convex:
                y = torch.mean(x.view(B, C, -1), dim=2)
                if penalty is not Quadratic:
                    y = RobustVectorPool2dFcn._optimize(x, y, penalty, alpha)
            else:
                y_mean = torch.mean(x.view(B, C, -1), dim=2)
                y = RobustVectorPool2dFcn._optimize(x, y_mean, penalty, alpha)
                f = penalty.phi(torch.linalg.norm((y.view(B, C, 1) - x.view(B, C, -1)), dim=1), alpha).view(B, -1).sum(-1)

                y_median, _ = torch.median(x.view(B, C, -1), dim=2)
                y_median = RobustVectorPool2dFcn._optimize(x, y_median, penalty, alpha)
                f_median = penalty.phi(torch.linalg.norm((y_median.view(B, C, 1) - x.view(B, C, -1)), dim=1), alpha).view(B, -1).sum(-1)
                y = torch.where((f <= f_median).view(B, 1), y, y_median)
                f = torch.minimum(f, f_median)

                if restarts > 0:
                    # choose points evenly spaced in batch (can be randomly chosen but that affects reproducibility)
                    for i in torch.linspace(start=0, end=x.view(B, C, -1).shape[2] - 1, steps=restarts):
                        y_init = x.view(B, C, -1)[:, :, int(i.item())]
                        y_final = RobustVectorPool2dFcn._optimize(x, y_init, penalty, alpha)
                        f_final = penalty.phi(torch.linalg.norm((y_final.view(B, C, 1) - x.view(B, C, -1)), dim=1), alpha).view(B, -1).sum(-1)
                        y = torch.where((f < f_final).view(B, 1), y, y_final)
                        f = torch.minimum(f, f_final)

        ctx.save_for_backward(x, y)
        ctx.penalty = penalty
        ctx.alpha = alpha
        ctx.hess_reg = hess_reg

        return y
//  [python]
def score_assignment_step(attn, v):
        """
        Token Score Assignment Step.
        :param attn: attention matrix
        :param v: values
        :return: sorted significance scores and their corresponding indices
        """

        B, H, _, _ = attn.shape
        C = v.shape[3] * H
        v_norm = torch.norm(
            v.transpose(1, 2).reshape(B, attn.shape[2], C), 2, 2
        )  # value norm of size [B x T]
        significance_score = attn[:, :, 0].sum(
            dim=1
        )  # attention weights of CLS token of size [B x T]
        significance_score = significance_score * v_norm  # [B x T]
        significance_score = significance_score[:, 1:]  # [B x T-1]

        significance_score = significance_score / significance_score.sum(
            dim=1, keepdim=True
        )  # [B x T-1]
        
        return significance_score
//  [python]
def score_assignment_step(attn, v):
        """
        Token Score Assignment Step.
        :param attn: attention matrix
        :param v: values
        :return: sorted significance scores and their corresponding indices
        """

        B, H, _, _ = attn.shape
        C = v.shape[3] * H
        v_norm = torch.linalg.norm(
            v.transpose(1, 2).reshape(B, attn.shape[2], C), ord=2, dim=2
        )  # value norm of size [B x T]
        significance_score = attn[:, :, 0].sum(
            dim=1
        )  # attention weights of CLS token of size [B x T]
        significance_score = significance_score * v_norm  # [B x T]
        significance_score = significance_score[:, 1:]  # [B x T-1]

        significance_score = significance_score / significance_score.sum(
            dim=1, keepdim=True
        )  # [B x T-1]
        
        return significance_score
//  [python]
def encode_speaker(
            self,
            profile: torch.Tensor,
            profile_lengths: torch.Tensor,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        with autocast(False):
            if profile.shape[1] < self.max_spk_num:
                profile = F.pad(profile, [0, 0, 0, self.max_spk_num-profile.shape[1], 0, 0], "constant", 0.0)
            profile_mask = (torch.norm(profile, p=2, dim=2, keepdim=True) > 0).float()
            profile = F.normalize(profile, dim=2)
            if self.speaker_encoder is not None:
                profile = self.speaker_encoder(profile, profile_lengths)[0]
                return profile * profile_mask, profile_lengths
            else:
                return profile, profile_lengths
//  [python]
def encode_speaker(
            self,
            profile: torch.Tensor,
            profile_lengths: torch.Tensor,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        with autocast(False):
            if profile.shape[1] < self.max_spk_num:
                profile = F.pad(profile, [0, 0, 0, self.max_spk_num-profile.shape[1], 0, 0], "constant", 0.0)
            profile_mask = (torch.linalg.norm(profile, ord=2, dim=2, keepdim=True) > 0).float()
            profile = F.normalize(profile, dim=2)
            if self.speaker_encoder is not None:
                profile = self.speaker_encoder(profile, profile_lengths)[0]
                return profile * profile_mask, profile_lengths
            else:
                return profile, profile_lengths
//  [python]
def get_det_sensitivities(self, 
                              theta: torch.Tensor, 
                              ego_x: torch.Tensor, # Just the current ego state (self.state_dim, )
                              det_mus: torch.Tensor,
                              extra_info: Dict[str, torch.Tensor]) -> torch.Tensor:

        def detection_reward_term(det_mus_vec: torch.Tensor, ego_x: torch.Tensor, 
                                  theta: torch.Tensor, extra_info: Dict[str, Any]):
            det_mus = det_mus_vec.reshape(extra_info['det_mus_shape'])

            ego_obs_delta = ego_x[..., :2] - det_mus
            return -theta[3] * utils.pt_rbf(ego_obs_delta, scale=2).sum(-1)

        extra_info['det_mus_shape'] = det_mus.shape[1:]

        det_sens_mags = torch.zeros((det_mus.shape[0], det_mus.shape[-1]))
        for det_idx in range(det_mus.shape[0]):
            g_mus = AF.jacobian(partial(detection_reward_term, 
                                        ego_x=ego_x,
                                        theta=theta, 
                                        extra_info=extra_info), 
                                det_mus[det_idx].flatten(),
                                create_graph=False,
                                vectorize=True)

            det_sens_mags[det_idx] = g_mus
        return torch.norm(det_sens_mags, dim=-1)
//  [python]
def get_det_sensitivities(self, 
                              theta: torch.Tensor, 
                              ego_x: torch.Tensor, # Just the current ego state (self.state_dim, )
                              det_mus: torch.Tensor,
                              extra_info: Dict[str, torch.Tensor]) -> torch.Tensor:

        def detection_reward_term(det_mus_vec: torch.Tensor, ego_x: torch.Tensor, 
                                  theta: torch.Tensor, extra_info: Dict[str, Any]):
            det_mus = det_mus_vec.reshape(extra_info['det_mus_shape'])

            ego_obs_delta = ego_x[..., :2] - det_mus
            return -theta[3] * utils.pt_rbf(ego_obs_delta, scale=2).sum(-1)

        extra_info['det_mus_shape'] = det_mus.shape[1:]

        det_sens_mags = torch.zeros((det_mus.shape[0], det_mus.shape[-1]))
        for det_idx in range(det_mus.shape[0]):
            g_mus = AF.jacobian(partial(detection_reward_term, 
                                        ego_x=ego_x,
                                        theta=theta, 
                                        extra_info=extra_info), 
                                det_mus[det_idx].flatten(),
                                create_graph=False,
                                vectorize=True)

            det_sens_mags[det_idx] = g_mus
        return torch.linalg.norm(det_sens_mags, dim=-1)
//  [python]
def speaker_discrimination_loss(
            self,
            profile: torch.Tensor,
            profile_lengths: torch.Tensor
    ) -> torch.Tensor:
        profile_mask = (torch.norm(profile, p=2, dim=2, keepdim=True) > 0).float()  # (B, N, 1)
        mask = torch.matmul(profile_mask, profile_mask.transpose(1, 2))  # (B, N, N)
        mask = mask * (1.0 - torch.eye(self.max_spk_num).unsqueeze(0).to(mask))

        eps = 1e-12
        coding_norm = torch.norm(
            profile * profile_mask + (1 - profile_mask) * eps,
            dim=2, keepdim=True
        ) * profile_mask
        # profile: Batch, N, dim
        cos_theta = F.cosine_similarity(profile.unsqueeze(2), profile.unsqueeze(1), dim=-1, eps=eps) * mask
        cos_theta = torch.clip(cos_theta, -1 + eps, 1 - eps)
        loss = (F.relu(mask * coding_norm * (cos_theta - 0.0))).sum() / mask.sum()

        return loss
//  [python]
def speaker_discrimination_loss(
            self,
            profile: torch.Tensor,
            profile_lengths: torch.Tensor
    ) -> torch.Tensor:
        profile_mask = (torch.linalg.norm(profile, ord=2, dim=2, keepdim=True) > 0).float()  # (B, N, 1)
        mask = torch.matmul(profile_mask, profile_mask.transpose(1, 2))  # (B, N, N)
        mask = mask * (1.0 - torch.eye(self.max_spk_num).unsqueeze(0).to(mask))

        eps = 1e-12
        coding_norm = torch.linalg.norm(
            profile * profile_mask + (1 - profile_mask) * eps,
            dim=2, keepdim=True
        ) * profile_mask
        # profile: Batch, N, dim
        cos_theta = F.cosine_similarity(profile.unsqueeze(2), profile.unsqueeze(1), dim=-1, eps=eps) * mask
        cos_theta = torch.clip(cos_theta, -1 + eps, 1 - eps)
        loss = (F.relu(mask * coding_norm * (cos_theta - 0.0))).sum() / mask.sum()

        return loss