jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def compute_reward(self, state: Tensor, action: Tensor) -> float:
    agent_atom = state if self.state_only else torch.cat([state, action], dim=1)
    agent_atom = self.data_scale * (agent_atom + self.data_offset)  # Normalise the agent atom
    weight, cost = 1 / self.time_horizon - 1e-6, 0  # Note: subtracting eps from initial agent atom weight for numerical stability
    dists = torch.linalg.norm(self.expert_atoms - agent_atom, dim=1)
    while weight > 0:
      closest_expert_idx = dists.argmin().item()  # Find closest expert atom
      expert_weight = self.expert_weights[closest_expert_idx].item()
      # Update costs and weights
      if weight >= expert_weight:
        cost += expert_weight * dists[closest_expert_idx].item()
        weight -= expert_weight
        self.expert_atoms, self.expert_weights, dists = _delete_row(self.expert_atoms, closest_expert_idx), _delete_row(self.expert_weights, closest_expert_idx), _delete_row(dists, closest_expert_idx)  # Remove the expert atom
      else:
        cost += weight * dists[closest_expert_idx].item()
        self.expert_weights[closest_expert_idx] -= weight
        weight = 0
    return self.reward_scale * exp(-self.reward_bandwidth * cost)
//  [python]
def _project_impl(self, pt: Point) -> Point:
        # Assume that pt.shape == (m, p). The first operation is to center it:
        pt = pt - torch.mean(pt, dim=0)
        # If score_method is 'angular' then scale doesn't matter; normalize the point
        if self._score_method == "angular":
            pt = pt / torch.norm(pt, ord="fro")
        # Note – whitening is only done in neural_data_to_point since whitening is only idempotent if alpha==0.
        return pt
//  [python]
def _project_impl(self, pt: Point) -> Point:
        # Assume that pt.shape == (m, p). The first operation is to center it:
        pt = pt - torch.mean(pt, dim=0)
        # If score_method is 'angular' then scale doesn't matter; normalize the point
        if self._score_method == "angular":
            pt = pt / norm(pt, ord="fro")
        # Note – whitening is only done in neural_data_to_point since whitening is only idempotent if alpha==0.
        return pt
//  [python]
def _length_impl(self, pt_a: Point, pt_b: Point) -> Scalar:
        if self._score_method == "euclidean":
            diff_ab = pt_a - pt_b
            return torch.norm(diff_ab)
        elif self._score_method == "angular":
            cos_ab = torch.sum(pt_a * pt_b) / torch.sqrt(torch.sum(pt_a * pt_a) * torch.sum(pt_b * pt_b))
            return torch.arccos(torch.clip(cos_ab, -1.0, +1.0))
//  [python]
def _length_impl(self, pt_a: Point, pt_b: Point) -> Scalar:
        if self._score_method == "euclidean":
            diff_ab = pt_a - pt_b
            return norm(diff_ab, ord="fro")
        elif self._score_method == "angular":
            cos_ab = torch.sum(pt_a * pt_b) / torch.sqrt(torch.sum(pt_a * pt_a) * torch.sum(pt_b * pt_b))
            return torch.arccos(torch.clip(cos_ab, -1.0, +1.0))
//  [python]
def normalize(x) -> TensorType[...]:
    """Returns a normalized vector."""
    return x / torch.norm(x)
//  [python]
def normalize(x) -> TensorType[...]:
    """Returns a normalized vector."""
    return x / torch.linalg.norm(x)
//  [python]
def _contains_impl(self, pt: Point, atol: float = 1e-6) -> bool:
        # Test shape
        if pt.shape != (self.m, self.p):
            return False
        # Test centered
        if not torch.allclose(torch.mean(pt, dim=0), pt.new_zeros(pt.shape[1:]), atol=atol):
            return False
        # Test unit norm (if angular)
        if self._score_method == "angular" and not torch.isclose(norm(pt, ord="fro"), pt.new_ones((1,))):
            return False
        if self._alpha == 0.:
            # Test whitened (if alpha==0)
            whitened_pt = _whiten(pt, 0.)
            if self._score_method == "angular":
                if not torch.allclose(pt, whitened_pt / torch.norm(whitened_pt, ord="fro"), atol=atol):
                    return False
            elif self._score_method == "euclidean":
                if not torch.allclose(pt, whitened_pt, atol=atol):
                    return False
        return True
//  [python]
def _contains_impl(self, pt: Point, atol: float = 1e-6) -> bool:
        # Test shape
        if pt.shape != (self.m, self.p):
            return False
        # Test centered
        if not torch.allclose(torch.mean(pt, dim=0), pt.new_zeros(pt.shape[1:]), atol=atol):
            return False
        # Test unit norm (if angular)
        if self._score_method == "angular" and not torch.isclose(norm(pt, ord="fro"), pt.new_ones((1,))):
            return False
        if self._alpha == 0.:
            # Test whitened (if alpha==0)
            whitened_pt = _whiten(pt, 0.)
            if self._score_method == "angular":
                if not torch.allclose(pt, whitened_pt / torch.linalg.norm(whitened_pt, ord="fro"), atol=atol):
                    return False
            elif self._score_method == "euclidean":
                if not torch.allclose(pt, whitened_pt, atol=atol):
                    return False
        return True
//  [python]
def neural_data_to_point(self, x: NeuralData) -> Point:
        """Convert size (m,d) neural data to a size (m,p) matrix. If d>p then we keep the top p dimensions by PCA. If
        d<p we pad with zeros.

        Essentially, this takes care of 'translation' and 'scaling' invariances, up to some regularization.
        """
        if x.shape[0] != self.m:
            raise ValueError(f"Expected x to be size ({self.m}, ?) but is size {x.shape}")

        # Flatten all but first dimension. TODO optional conv reshape into (m*h*w,c) as done by Williams et al?
        x = torch.reshape(x, (self.m, -1))

        # Center columns to handle translation-invariance
        x = x - torch.mean(x, dim=0, keepdim=True)

        # Pad or truncate to p dimensions. Important that this happens before whitening because whitening destroys
        # principal components (makes them all equivalent).
        d = prod(x.shape) // self.m
        if d > self.p:
            x = _dim_reduce(x, self.p)
        elif d < self.p:
            x = _pad_zeros(x, self.p)

        # Rescale and (partially) whiten to handle scale-invariance, using self._alpha.
        x = _whiten(x, self._alpha)

        # In case of 'angular' metric only, make all point clouds unit-frobenius-norm. Then, points live on the
        # hypesphere of centered and unit-norm m-by-p matrices.
        if self._score_method == "angular":
            x = x / torch.norm(x, ord="fro")

        return x
//  [python]
def neural_data_to_point(self, x: NeuralData) -> Point:
        """Convert size (m,d) neural data to a size (m,p) matrix. If d>p then we keep the top p dimensions by PCA. If
        d<p we pad with zeros.

        Essentially, this takes care of 'translation' and 'scaling' invariances, up to some regularization.
        """
        if x.shape[0] != self.m:
            raise ValueError(f"Expected x to be size ({self.m}, ?) but is size {x.shape}")

        # Flatten all but first dimension. TODO optional conv reshape into (m*h*w,c) as done by Williams et al?
        x = torch.reshape(x, (self.m, -1))

        # Center columns to handle translation-invariance
        x = x - torch.mean(x, dim=0, keepdim=True)

        # Pad or truncate to p dimensions. Important that this happens before whitening because whitening destroys
        # principal components (makes them all equivalent).
        d = prod(x.shape) // self.m
        if d > self.p:
            x = _dim_reduce(x, self.p)
        elif d < self.p:
            x = _pad_zeros(x, self.p)

        # Rescale and (partially) whiten to handle scale-invariance, using self._alpha.
        x = _whiten(x, self._alpha)

        # In case of 'angular' metric only, make all point clouds unit-frobenius-norm. Then, points live on the
        # hypesphere of centered and unit-norm m-by-p matrices.
        if self._score_method == "angular":
            x = x / norm(x, ord="fro")

        return x
//  [python]
def norm(self, arr: array_like) -> float:
        return float(torch.norm(torch.view_as_real(arr)))
//  [python]
def norm(self, arr: array_like) -> float:
        return float(torch.linalg.norm(torch.view_as_real(arr)))
//  [python]
def robustify_photo(r, dIt_dT, invalid_mask, photo_sigma):
    info_sqrt = 1.0 / photo_sigma
    whitened_r = r * info_sqrt
    weight = robust.huber(whitened_r)
    weight[invalid_mask[...], :] = 0.0

    total_err = torch.sum(weight * torch.square(whitened_r))
    num_valid = invalid_mask.shape[-1] - torch.count_nonzero(invalid_mask, dim=-1)
    mean_sq_err = total_err / num_valid

    J_W = weight[:, :, :, None] * dIt_dT
    grad = torch.sum(J_W * r[..., None], dim=(1, 2))
    H = torch.einsum("bnck,bncl->bkl", J_W, dIt_dT)

    grad_norm = torch.norm(grad)

    return H, grad, total_err, mean_sq_err, grad_norm
//  [python]
def robustify_photo(r, dIt_dT, invalid_mask, photo_sigma):
    info_sqrt = 1.0 / photo_sigma
    whitened_r = r * info_sqrt
    weight = robust.huber(whitened_r)
    weight[invalid_mask[...], :] = 0.0

    total_err = torch.sum(weight * torch.square(whitened_r))
    num_valid = invalid_mask.shape[-1] - torch.count_nonzero(invalid_mask, dim=-1)
    mean_sq_err = total_err / num_valid

    J_W = weight[:, :, :, None] * dIt_dT
    grad = torch.sum(J_W * r[..., None], dim=(1, 2))
    H = torch.einsum("bnck,bncl->bkl", J_W, dIt_dT)

    grad_norm = torch.linalg.norm(grad)

    return H, grad, total_err, mean_sq_err, grad_norm
//  [python]
def pose_error_torch(R, t, Tgt, reduce=None):
    """Compute angular, scale and euclidean error of translation vector (metric). Compute angular rotation error."""

    Rgt = Tgt[:, :3, :3]                  # [B, 3, 3]
    tgt = Tgt[:, :3, 3:].transpose(1, 2)  # [B, 1, 3]

    scale_t = torch.norm(t, dim=-1)
    scale_tgt = torch.norm(tgt, dim=-1)

    cosine = (t @ tgt.transpose(1, 2)).squeeze(-1) / (scale_t * scale_tgt + 1e-9)
    cosine = torch.clip(cosine, -1.0, 1.0)    # handle numerical errors
    t_ang_err = torch.rad2deg(torch.acos(cosine))
    t_ang_err = torch.minimum(t_ang_err, 180 - t_ang_err)

    t_scale_err = scale_t / scale_tgt
    t_scale_err_sym = torch.maximum(scale_t / scale_tgt, scale_tgt / scale_t)
    t_euclidean_err = torch.norm(t - tgt, dim=-1)

    residual = R.transpose(1, 2) @ Rgt
    trace = torch.diagonal(residual, dim1=-2, dim2=-1).sum(-1)
    cosine = (trace - 1) / 2
    cosine = torch.clip(cosine, -1., 1.)  # handle numerical errors
    R_err = torch.rad2deg(torch.acos(cosine))

    if reduce is None:
        def fn(x): return x
    elif reduce == 'mean':
        fn = torch.mean
    elif reduce == 'median':
        fn = torch.median

    t_ang_err = fn(t_ang_err)
    t_scale_err = fn(t_scale_err)
    t_euclidean_err = fn(t_euclidean_err)
    R_err = fn(R_err)

    errors = {'t_err_ang': t_ang_err,
              't_err_scale': t_scale_err,
              't_err_scale_sym': t_scale_err_sym,
              't_err_euc': t_euclidean_err,
              'R_err': R_err}
    return errors
//  [python]
def pose_error_torch(R, t, Tgt, reduce=None):
    """Compute angular, scale and euclidean error of translation vector (metric). Compute angular rotation error."""

    Rgt = Tgt[:, :3, :3]                  # [B, 3, 3]
    tgt = Tgt[:, :3, 3:].transpose(1, 2)  # [B, 1, 3]

    scale_t = torch.linalg.norm(t, dim=-1)
    scale_tgt = torch.linalg.norm(tgt, dim=-1)

    cosine = (t @ tgt.transpose(1, 2)).squeeze(-1) / (scale_t * scale_tgt + 1e-9)
    cosine = torch.clip(cosine, -1.0, 1.0)    # handle numerical errors
    t_ang_err = torch.rad2deg(torch.acos(cosine))
    t_ang_err = torch.minimum(t_ang_err, 180 - t_ang_err)

    t_scale_err = scale_t / scale_tgt
    t_scale_err_sym = torch.maximum(scale_t / scale_tgt, scale_tgt / scale_t)
    t_euclidean_err = torch.linalg.norm(t - tgt, dim=-1)

    residual = R.transpose(1, 2) @ Rgt
    trace = torch.diagonal(residual, dim1=-2, dim2=-1).sum(-1)
    cosine = (trace - 1) / 2
    cosine = torch.clip(cosine, -1., 1.)  # handle numerical errors
    R_err = torch.rad2deg(torch.acos(cosine))

    if reduce is None:
        def fn(x): return x
    elif reduce == 'mean':
        fn = torch.mean
    elif reduce == 'median':
        fn = torch.median

    t_ang_err = fn(t_ang_err)
    t_scale_err = fn(t_scale_err)
    t_euclidean_err = fn(t_euclidean_err)
    R_err = fn(R_err)

    errors = {'t_err_ang': t_ang_err,
              't_err_scale': t_scale_err,
              't_err_scale_sym': t_scale_err_sym,
              't_err_euc': t_euclidean_err,
              'R_err': R_err}
    return errors
//  [python]
def _compute_v_times_Wo_norm(
    attn_value: torch.Tensor,  # thd
    W_O: torch.Tensor,  # hdo
) -> torch.Tensor:
    """
    This function computes the norm of the product of attention values and the output weight matrix
    (W_O in the transformer-circuits convention, or c_proj in many internal settings).
    The attention values are represented by the tensor attn_value and the output weight matrix by W_O.
    The function returns a tensor representing the norm of the product.
    """

    num_attended_to_sequence_tokens, nheads, d_head = attn_value.shape
    assert W_O.shape[:2] == (nheads, d_head)  # third dim is d_model
    v_times_Wo = torch.einsum("thd,hdo->tho", attn_value, W_O)
    return torch.norm(v_times_Wo, dim=-1)
//  [python]
def _compute_v_times_Wo_norm(
    attn_value: torch.Tensor,  # thd
    W_O: torch.Tensor,  # hdo
) -> torch.Tensor:
    """
    This function computes the norm of the product of attention values and the output weight matrix
    (W_O in the transformer-circuits convention, or c_proj in many internal settings).
    The attention values are represented by the tensor attn_value and the output weight matrix by W_O.
    The function returns a tensor representing the norm of the product.
    """

    num_attended_to_sequence_tokens, nheads, d_head = attn_value.shape
    assert W_O.shape[:2] == (nheads, d_head)  # third dim is d_model
    v_times_Wo = torch.einsum("thd,hdo->tho", attn_value, W_O)
    return torch.linalg.norm(v_times_Wo, dim=-1)
//  [python]
def print_params_min_max_norm(optimizer, iteration):
    """Print min, max, and norm of all parameters."""
    index = 0
    rank = torch.distributed.get_rank()
    string = 'iteration, rank, index, tensor-model-parallel, min, max, norm\n'
    optimizer_ = optimizer.optimizer
    for param_group in optimizer_.param_groups:
        for param in param_group['params']:
            index += 1
            min_ = param.data.min()
            max_ = param.data.max()
            norm = param.data.norm()
            string += '{:7d}, {:4d}, {:4d}, {:2d}, '.format(
                iteration, rank, index, int(param.tensor_model_parallel))
            string += '{:.6E}, {:.6E}, {:.6E}\n'.format(min_, max_, norm)
    print(string, flush=True)
//  [python]
def print_params_min_max_norm(optimizer, iteration):
    """Print min, max, and norm of all parameters."""
    index = 0
    rank = torch.distributed.get_rank()
    string = 'iteration, rank, index, tensor-model-parallel, min, max, norm\n'
    optimizer_ = optimizer.optimizer
    for param_group in optimizer_.param_groups:
        for param in param_group['params']:
            index += 1
            min_ = param.data.min()
            max_ = param.data.max()
            norm = torch.linalg.norm(param.data)
            string += '{:7d}, {:4d}, {:4d}, {:2d}, '.format(
                iteration, rank, index, int(param.tensor_model_parallel))
            string += '{:.6E}, {:.6E}, {:.6E}\n'.format(min_, max_, norm)
    print(string, flush=True)
//  [python]
def plan_col_loss(
    pred,
    target,
    agent_fut_preds,
    x_dis_thresh=1.5,
    y_dis_thresh=3.0,
    dis_thresh=3.0
):
    """Planning ego-agent collsion constraint.

    Args:
        pred (torch.Tensor): ego_fut_preds, [B, fut_ts, 2].
        target (torch.Tensor): agent_preds, [B, num_agent, 2].
        agent_fut_preds (Tensor): [B, num_agent, fut_ts, 2].
        weight (torch.Tensor): [B, fut_ts, 2].
        x_dis_thresh (float, optional): distance threshold between ego and other agents in x-axis.
        y_dis_thresh (float, optional): distance threshold between ego and other agents in y-axis.
        dis_thresh (float, optional): distance threshold to filter distant agents.

    Returns:
        torch.Tensor: Calculated loss [B, fut_mode, fut_ts, 2]
    """
    pred = pred.cumsum(dim=-2)
    agent_fut_preds = agent_fut_preds.cumsum(dim=-2)
    target = target[:, :, None, :] + agent_fut_preds
    # filter distant agents from ego vehicle
    dist = torch.norm(pred[:, None, :, :] - target, dim=-1)
    dist_mask = dist > dis_thresh
    target[dist_mask] = 1e6

    # [B, num_agent, fut_ts]
    x_dist = torch.abs(pred[:, None, :, 0] - target[..., 0])
    y_dist = torch.abs(pred[:, None, :, 1] - target[..., 1])
    x_min_idxs = torch.argmin(x_dist, dim=1).tolist()
    y_min_idxs = torch.argmin(y_dist, dim=1).tolist()
    batch_idxs = [[i] for i in range(y_dist.shape[0])]
    ts_idxs = [[i for i in range(y_dist.shape[-1])] for j in range(y_dist.shape[0])]

    # [B, fut_ts]
    x_min_dist = x_dist[batch_idxs, x_min_idxs, ts_idxs]
    y_min_dist = y_dist[batch_idxs, y_min_idxs, ts_idxs]
    x_loss = x_min_dist
    safe_idx = x_loss > x_dis_thresh
    unsafe_idx = x_loss <= x_dis_thresh
    x_loss[safe_idx] = 0
    x_loss[unsafe_idx] = x_dis_thresh - x_loss[unsafe_idx]
    y_loss = y_min_dist
    safe_idx = y_loss > y_dis_thresh
    unsafe_idx = y_loss <= y_dis_thresh
    y_loss[safe_idx] = 0
    y_loss[unsafe_idx] = y_dis_thresh - y_loss[unsafe_idx]
    loss = torch.cat([x_loss.unsqueeze(-1), y_loss.unsqueeze(-1)], dim=-1)

    return loss
//  [python]
def plan_col_loss(
    pred,
    target,
    agent_fut_preds,
    x_dis_thresh=1.5,
    y_dis_thresh=3.0,
    dis_thresh=3.0
):
    """Planning ego-agent collsion constraint.

    Args:
        pred (torch.Tensor): ego_fut_preds, [B, fut_ts, 2].
        target (torch.Tensor): agent_preds, [B, num_agent, 2].
        agent_fut_preds (Tensor): [B, num_agent, fut_ts, 2].
        weight (torch.Tensor): [B, fut_ts, 2].
        x_dis_thresh (float, optional): distance threshold between ego and other agents in x-axis.
        y_dis_thresh (float, optional): distance threshold between ego and other agents in y-axis.
        dis_thresh (float, optional): distance threshold to filter distant agents.

    Returns:
        torch.Tensor: Calculated loss [B, fut_mode, fut_ts, 2]
    """
    pred = pred.cumsum(dim=-2)
    agent_fut_preds = agent_fut_preds.cumsum(dim=-2)
    target = target[:, :, None, :] + agent_fut_preds
    # filter distant agents from ego vehicle
    dist = torch.linalg.norm(pred[:, None, :, :] - target, dim=-1)
    dist_mask = dist > dis_thresh
    target[dist_mask] = 1e6

    # [B, num_agent, fut_ts]
    x_dist = torch.abs(pred[:, None, :, 0] - target[..., 0])
    y_dist = torch.abs(pred[:, None, :, 1] - target[..., 1])
    x_min_idxs = torch.argmin(x_dist, dim=1).tolist()
    y_min_idxs = torch.argmin(y_dist, dim=1).tolist()
    batch_idxs = [[i] for i in range(y_dist.shape[0])]
    ts_idxs = [[i for i in range(y_dist.shape[-1])] for j in range(y_dist.shape[0])]

    # [B, fut_ts]
    x_min_dist = x_dist[batch_idxs, x_min_idxs, ts_idxs]
    y_min_dist = y_dist[batch_idxs, y_min_idxs, ts_idxs]
    x_loss = x_min_dist
    safe_idx = x_loss > x_dis_thresh
    unsafe_idx = x_loss <= x_dis_thresh
    x_loss[safe_idx] = 0
    x_loss[unsafe_idx] = x_dis_thresh - x_loss[unsafe_idx]
    y_loss = y_min_dist
    safe_idx = y_loss > y_dis_thresh
    unsafe_idx = y_loss <= y_dis_thresh
    y_loss[safe_idx] = 0
    y_loss[unsafe_idx] = y_dis_thresh - y_loss[unsafe_idx]
    loss = torch.cat([x_loss.unsqueeze(-1), y_loss.unsqueeze(-1)], dim=-1)

    return loss
//  [python]
def generate_mask_init_pcl(self):
        init_pcl_mask = torch.norm(self.init_pcl, dim=-1) < 5
        self.init_pcl_mask = init_pcl_mask
        # init_pcl[~init_pcl_mask] = 0
        self.init_pcl_reduced = self.init_pcl[self.init_pcl_mask]
//  [python]
def generate_mask_init_pcl(self):
        init_pcl_mask = torch.linalg.norm(self.init_pcl, dim=-1) < 5
        self.init_pcl_mask = init_pcl_mask
        # init_pcl[~init_pcl_mask] = 0
        self.init_pcl_reduced = self.init_pcl[self.init_pcl_mask]
//  [python]
def plan_map_bound_loss(pred, target, dis_thresh=1.0):
    """Planning map bound constraint (L1 distance).

    Args:
        pred (torch.Tensor): ego_fut_preds, [B, fut_ts, 2].
        target (torch.Tensor): lane_bound_preds, [B, num_vec, num_pts, 2].
        weight (torch.Tensor): [B, fut_ts]

    Returns:
        torch.Tensor: Calculated loss [B, fut_ts]
    """
    pred = pred.cumsum(dim=-2)
    ego_traj_starts = pred[:, :-1, :]
    ego_traj_ends = pred
    B, T, _ = ego_traj_ends.size()
    padding_zeros = torch.zeros((B, 1, 2), dtype=pred.dtype, device=pred.device)  # initial position
    ego_traj_starts = torch.cat((padding_zeros, ego_traj_starts), dim=1)
    _, V, P, _ = target.size()
    ego_traj_expanded = ego_traj_ends.unsqueeze(2).unsqueeze(3)  # [B, T, 1, 1, 2]
    maps_expanded = target.unsqueeze(1)  # [1, 1, M, P, 2]
    dist = torch.norm(ego_traj_expanded - maps_expanded, dim=-1)  # [B, T, M, P]
    dist = dist.min(dim=-1, keepdim=False)[0]
    min_inst_idxs = torch.argmin(dist, dim=-1).tolist()
    batch_idxs = [[i] for i in range(dist.shape[0])]
    ts_idxs = [[i for i in range(dist.shape[1])] for j in range(dist.shape[0])]
    bd_target = target.unsqueeze(1).repeat(1, pred.shape[1], 1, 1, 1)
    min_bd_insts = bd_target[batch_idxs, ts_idxs, min_inst_idxs]  # [B, T, P, 2]
    bd_inst_starts = min_bd_insts[:, :, :-1, :].flatten(0, 2)
    bd_inst_ends = min_bd_insts[:, :, 1:, :].flatten(0, 2)
    ego_traj_starts = ego_traj_starts.unsqueeze(2).repeat(1, 1, P-1, 1).flatten(0, 2)
    ego_traj_ends = ego_traj_ends.unsqueeze(2).repeat(1, 1, P-1, 1).flatten(0, 2)

    intersect_mask = segments_intersect(ego_traj_starts, ego_traj_ends,
                                        bd_inst_starts, bd_inst_ends)
    intersect_mask = intersect_mask.reshape(B, T, P-1)
    intersect_mask = intersect_mask.any(dim=-1)
    intersect_idx = (intersect_mask == True).nonzero()

    target = target.view(target.shape[0], -1, target.shape[-1])
    # [B, fut_ts, num_vec*num_pts]
    dist = torch.norm(pred[:, :, None, :] - target[:, None, :, :], dim=-1)
    min_idxs = torch.argmin(dist, dim=-1).tolist()
    batch_idxs = [[i] for i in range(dist.shape[0])]
    ts_idxs = [[i for i in range(dist.shape[1])] for j in range(dist.shape[0])]
    min_dist = dist[batch_idxs, ts_idxs, min_idxs]
    loss = min_dist
    safe_idx = loss > dis_thresh
    unsafe_idx = loss <= dis_thresh
    loss[safe_idx] = 0
    loss[unsafe_idx] = dis_thresh - loss[unsafe_idx]

    for i in range(len(intersect_idx)):
        loss[intersect_idx[i, 0], intersect_idx[i, 1]:] = 0

    return loss
//  [python]
def plan_map_bound_loss(pred, target, dis_thresh=1.0):
    """Planning map bound constraint (L1 distance).

    Args:
        pred (torch.Tensor): ego_fut_preds, [B, fut_ts, 2].
        target (torch.Tensor): lane_bound_preds, [B, num_vec, num_pts, 2].
        weight (torch.Tensor): [B, fut_ts]

    Returns:
        torch.Tensor: Calculated loss [B, fut_ts]
    """
    pred = pred.cumsum(dim=-2)
    ego_traj_starts = pred[:, :-1, :]
    ego_traj_ends = pred
    B, T, _ = ego_traj_ends.size()
    padding_zeros = torch.zeros((B, 1, 2), dtype=pred.dtype, device=pred.device)  # initial position
    ego_traj_starts = torch.cat((padding_zeros, ego_traj_starts), dim=1)
    _, V, P, _ = target.size()
    ego_traj_expanded = ego_traj_ends.unsqueeze(2).unsqueeze(3)  # [B, T, 1, 1, 2]
    maps_expanded = target.unsqueeze(1)  # [1, 1, M, P, 2]
    dist = torch.linalg.norm(ego_traj_expanded - maps_expanded, dim=-1)  # [B, T, M, P]
    dist = dist.min(dim=-1, keepdim=False)[0]
    min_inst_idxs = torch.argmin(dist, dim=-1).tolist()
    batch_idxs = [[i] for i in range(dist.shape[0])]
    ts_idxs = [[i for i in range(dist.shape[1])] for j in range(dist.shape[0])]
    bd_target = target.unsqueeze(1).repeat(1, pred.shape[1], 1, 1, 1)
    min_bd_insts = bd_target[batch_idxs, ts_idxs, min_inst_idxs]  # [B, T, P, 2]
    bd_inst_starts = min_bd_insts[:, :, :-1, :].flatten(0, 2)
    bd_inst_ends = min_bd_insts[:, :, 1:, :].flatten(0, 2)
    ego_traj_starts = ego_traj_starts.unsqueeze(2).repeat(1, 1, P-1, 1).flatten(0, 2)
    ego_traj_ends = ego_traj_ends.unsqueeze(2).repeat(1, 1, P-1, 1).flatten(0, 2)

    intersect_mask = segments_intersect(ego_traj_starts, ego_traj_ends,
                                        bd_inst_starts, bd_inst_ends)
    intersect_mask = intersect_mask.reshape(B, T, P-1)
    intersect_mask = intersect_mask.any(dim=-1)
    intersect_idx = (intersect_mask == True).nonzero()

    target = target.view(target.shape[0], -1, target.shape[-1])
    # [B, fut_ts, num_vec*num_pts]
    dist = torch.linalg.norm(pred[:, :, None, :] - target[:, None, :, :], dim=-1)
    min_idxs = torch.argmin(dist, dim=-1).tolist()
    batch_idxs = [[i] for i in range(dist.shape[0])]
    ts_idxs = [[i for i in range(dist.shape[1])] for j in range(dist.shape[0])]
    min_dist = dist[batch_idxs, ts_idxs, min_idxs]
    loss = min_dist
    safe_idx = loss > dis_thresh
    unsafe_idx = loss <= dis_thresh
    loss[safe_idx] = 0
    loss[unsafe_idx] = dis_thresh - loss[unsafe_idx]

    for i in range(len(intersect_idx)):
        loss[intersect_idx[i, 0], intersect_idx[i, 1]:] = 0

    return loss
//  [python]
def test_adversarial_robustness(args, simulator, filename, model, steps=10, label=1):
    def calculate_grad(x_,model):
        x_ = torch.autograd.Variable( x_)
        x_.requires_grad = True

        x_reco, log_prob_, _ = model(x_, context=None, mode="dnf")
        loss = log_prob_ + torch.norm( x_reco-x_)
        grad = torch.autograd.grad(loss, x_)
        return grad, loss
    
    alpha = 10 #stepsize
    fig = plt.figure(figsize=(10, 10))
    
    test_data = simulator.load_dataset(train=False, dataset_dir=filename, limit_samplesize=None, joint_score=None, label=label)
    
    x_test = test_data.__getitem__(1)[0].to('cuda')
    x_test = x_test.reshape([1,1,32,32]) 
#     x_test.requires_grad = True
# # output = model(input)
# # v = torch.autograd.Variable(torch.from_numpy(np.ones([1,10])))
#     # loss = log_prob_ - torch.linalg.norm(x_reco-x_test)
    
#     x_reco, log_prob_, _ = model(x_test, context=None, mode="dnf")
#     loss = log_prob_ - torch.linalg.norm( x_reco-x_test)

#     grad = torch.autograd.grad(loss, x_test)
    
    x_old = x_test
    for step in range(steps):
        grad, loss = calculate_grad(x_old,model)
        # import pdb
        # pdb.set_trace() 
        x_step = x_old + alpha * grad[0]
        # x_step = x_step
        x = np.clip(x_step[0,:].detach().cpu().numpy() / 256.0 , 0.0, 1.0)
        img = np.transpose(x, [1, 2, 0])
        ax = fig.add_subplot(1,steps,step+1)  
        ax.imshow(img)
        ax.set_xticks([])
        ax.set_yticks([])
        x_old = x_step
        value = np.round(loss.item(),3)
        plt.title(str(value))
        
    plt.tight_layout() 
    fig.savefig(create_filename("training_plot", "adversarial_gradient_attack", args), bbox_inches = 'tight')
//  [python]
def test_adversarial_robustness(args, simulator, filename, model, steps=10, label=1):
    def calculate_grad(x_,model):
        x_ = torch.autograd.Variable( x_)
        x_.requires_grad = True

        x_reco, log_prob_, _ = model(x_, context=None, mode="dnf")
        loss = log_prob_ + torch.linalg.norm( x_reco-x_)
        grad = torch.autograd.grad(loss, x_)
        return grad, loss
    
    alpha = 10 #stepsize
    fig = plt.figure(figsize=(10, 10))
    
    test_data = simulator.load_dataset(train=False, dataset_dir=filename, limit_samplesize=None, joint_score=None, label=label)
    
    x_test = test_data.__getitem__(1)[0].to('cuda')
    x_test = x_test.reshape([1,1,32,32]) 
#     x_test.requires_grad = True
# # output = model(input)
# # v = torch.autograd.Variable(torch.from_numpy(np.ones([1,10])))
#     # loss = log_prob_ - torch.linalg.norm(x_reco-x_test)
    
#     x_reco, log_prob_, _ = model(x_test, context=None, mode="dnf")
#     loss = log_prob_ - torch.linalg.norm( x_reco-x_test)

#     grad = torch.autograd.grad(loss, x_test)
    
    x_old = x_test
    for step in range(steps):
        grad, loss = calculate_grad(x_old,model)
        # import pdb
        # pdb.set_trace() 
        x_step = x_old + alpha * grad[0]
        # x_step = x_step
        x = np.clip(x_step[0,:].detach().cpu().numpy() / 256.0 , 0.0, 1.0)
        img = np.transpose(x, [1, 2, 0])
        ax = fig.add_subplot(1,steps,step+1)  
        ax.imshow(img)
        ax.set_xticks([])
        ax.set_yticks([])
        x_old = x_step
        value = np.round(loss.item(),3)
        plt.title(str(value))
        
    plt.tight_layout() 
    fig.savefig(create_filename("training_plot", "adversarial_gradient_attack", args), bbox_inches = 'tight')
//  [python]
def plan_map_dir_loss(pred, target, dis_thresh=2.0):
    """Planning ego-map directional loss.

    Args:
        pred (torch.Tensor): ego_fut_preds, [B, fut_ts, 2].
        target (torch.Tensor): lane_div_preds, [B, num_vec, num_pts, 2].
        weight (torch.Tensor): [B, fut_ts]

    Returns:
        torch.Tensor: Calculated loss [B, fut_ts]
    """
    num_map_pts = target.shape[2]
    pred = pred.cumsum(dim=-2)
    traj_dis = torch.norm(pred[:, -1, :] - pred[:, 0, :], dim=-1)
    static_mask = traj_dis < 1.0
    target = target.unsqueeze(1).repeat(1, pred.shape[1], 1, 1, 1)

    # find the closest map instance for ego at each timestamp
    dist = torch.norm(pred[:, :, None, None, :] - target, dim=-1)
    dist = dist.min(dim=-1, keepdim=False)[0]
    min_inst_idxs = torch.argmin(dist, dim=-1).tolist()
    batch_idxs = [[i] for i in range(dist.shape[0])]
    ts_idxs = [[i for i in range(dist.shape[1])] for j in range(dist.shape[0])]
    target_map_inst = target[batch_idxs, ts_idxs, min_inst_idxs]  # [B, fut_ts, num_pts, 2]

    # calculate distance
    dist = torch.norm(pred[:, :, None, :] - target_map_inst, dim=-1)
    min_pts_idxs = torch.argmin(dist, dim=-1)
    min_pts_next_idxs = min_pts_idxs.clone()
    is_end_point = (min_pts_next_idxs == num_map_pts-1)
    not_end_point = (min_pts_next_idxs != num_map_pts-1)
    min_pts_next_idxs[is_end_point] = num_map_pts - 2
    min_pts_next_idxs[not_end_point] = min_pts_next_idxs[not_end_point] + 1
    min_pts_idxs = min_pts_idxs.tolist()
    min_pts_next_idxs = min_pts_next_idxs.tolist()
    traj_yaw = torch.atan2(torch.diff(pred[..., 1]), torch.diff(pred[..., 0]))  # [B, fut_ts-1]
    # last ts yaw assume same as previous
    traj_yaw = torch.cat([traj_yaw, traj_yaw[:, [-1]]], dim=-1)  # [B, fut_ts]
    min_pts = target_map_inst[batch_idxs, ts_idxs, min_pts_idxs]
    dist = torch.norm(min_pts - pred, dim=-1)
    dist_mask = dist > dis_thresh
    min_pts = min_pts.unsqueeze(2)
    min_pts_next = target_map_inst[batch_idxs, ts_idxs, min_pts_next_idxs].unsqueeze(2)
    map_pts = torch.cat([min_pts, min_pts_next], dim=2)
    lane_yaw = torch.atan2(torch.diff(map_pts[..., 1]).squeeze(-1), torch.diff(map_pts[..., 0]).squeeze(-1))  # [B, fut_ts]
    yaw_diff = traj_yaw - lane_yaw
    yaw_diff[yaw_diff > math.pi] =  yaw_diff[yaw_diff > math.pi] - math.pi
    yaw_diff[yaw_diff > math.pi/2] = yaw_diff[yaw_diff > math.pi/2] - math.pi
    yaw_diff[yaw_diff < -math.pi] = yaw_diff[yaw_diff < -math.pi] + math.pi
    yaw_diff[yaw_diff < -math.pi/2] = yaw_diff[yaw_diff < -math.pi/2] + math.pi
    yaw_diff[dist_mask] = 0  # loss = 0 if no lane around ego
    yaw_diff[static_mask] = 0  # loss = 0 if ego is static

    loss = torch.abs(yaw_diff)

    return loss
//  [python]
def plan_map_dir_loss(pred, target, dis_thresh=2.0):
    """Planning ego-map directional loss.

    Args:
        pred (torch.Tensor): ego_fut_preds, [B, fut_ts, 2].
        target (torch.Tensor): lane_div_preds, [B, num_vec, num_pts, 2].
        weight (torch.Tensor): [B, fut_ts]

    Returns:
        torch.Tensor: Calculated loss [B, fut_ts]
    """
    num_map_pts = target.shape[2]
    pred = pred.cumsum(dim=-2)
    traj_dis = torch.linalg.norm(pred[:, -1, :] - pred[:, 0, :], dim=-1)
    static_mask = traj_dis < 1.0
    target = target.unsqueeze(1).repeat(1, pred.shape[1], 1, 1, 1)

    # find the closest map instance for ego at each timestamp
    dist = torch.linalg.norm(pred[:, :, None, None, :] - target, dim=-1)
    dist = dist.min(dim=-1, keepdim=False)[0]
    min_inst_idxs = torch.argmin(dist, dim=-1).tolist()
    batch_idxs = [[i] for i in range(dist.shape[0])]
    ts_idxs = [[i for i in range(dist.shape[1])] for j in range(dist.shape[0])]
    target_map_inst = target[batch_idxs, ts_idxs, min_inst_idxs]  # [B, fut_ts, num_pts, 2]

    # calculate distance
    dist = torch.linalg.norm(pred[:, :, None, :] - target_map_inst, dim=-1)
    min_pts_idxs = torch.argmin(dist, dim=-1)
    min_pts_next_idxs = min_pts_idxs.clone()
    is_end_point = (min_pts_next_idxs == num_map_pts-1)
    not_end_point = (min_pts_next_idxs != num_map_pts-1)
    min_pts_next_idxs[is_end_point] = num_map_pts - 2
    min_pts_next_idxs[not_end_point] = min_pts_next_idxs[not_end_point] + 1
    min_pts_idxs = min_pts_idxs.tolist()
    min_pts_next_idxs = min_pts_next_idxs.tolist()
    traj_yaw = torch.atan2(torch.diff(pred[..., 1]), torch.diff(pred[..., 0]))  # [B, fut_ts-1]
    # last ts yaw assume same as previous
    traj_yaw = torch.cat([traj_yaw, traj_yaw[:, [-1]]], dim=-1)  # [B, fut_ts]
    min_pts = target_map_inst[batch_idxs, ts_idxs, min_pts_idxs]
    dist = torch.linalg.norm(min_pts - pred, dim=-1)
    dist_mask = dist > dis_thresh
    min_pts = min_pts.unsqueeze(2)
    min_pts_next = target_map_inst[batch_idxs, ts_idxs, min_pts_next_idxs].unsqueeze(2)
    map_pts = torch.cat([min_pts, min_pts_next], dim=2)
    lane_yaw = torch.atan2(torch.diff(map_pts[..., 1]).squeeze(-1), torch.diff(map_pts[..., 0]).squeeze(-1))  # [B, fut_ts]
    yaw_diff = traj_yaw - lane_yaw
    yaw_diff[yaw_diff > math.pi] =  yaw_diff[yaw_diff > math.pi] - math.pi
    yaw_diff[yaw_diff > math.pi/2] = yaw_diff[yaw_diff > math.pi/2] - math.pi
    yaw_diff[yaw_diff < -math.pi] = yaw_diff[yaw_diff < -math.pi] + math.pi
    yaw_diff[yaw_diff < -math.pi/2] = yaw_diff[yaw_diff < -math.pi/2] + math.pi
    yaw_diff[dist_mask] = 0  # loss = 0 if no lane around ego
    yaw_diff[static_mask] = 0  # loss = 0 if ego is static

    loss = torch.abs(yaw_diff)

    return loss
//  [python]
def query_embedding(self, HDWD, cam_xyz, photometric_confidence, img_feats, c2ws, w2cs, intrinsics, cam_vid, pointdir_w=False):

        HD, WD = HDWD
        points_embedding = []
        points_dirs = None
        points_conf = None
        points_colors = None
        for feat_str in getattr(self.args, feature_str_lst[cam_vid]):
            if feat_str.startswith("imgfeat"):
                _, view_ids, layer_ids = feat_str.split("_")
                view_ids = [int(a) for a in list(view_ids)]
                layer_ids = [int(a) for a in list(layer_ids)]
                twoD_feats, points_colors = self.extract_2d(img_feats, view_ids, layer_ids, intrinsics, c2ws, w2cs, cam_xyz, HD, WD, cam_vid=cam_vid)
                points_embedding.append(twoD_feats)
            elif feat_str.startswith("dir"):
                _, view_ids = feat_str.split("_")
                view_ids = torch.as_tensor([int(a) for a in list(view_ids)], dtype=torch.int64, device=cam_xyz.device)
                cam_pos_world = c2ws[:, view_ids, :, 3] # B, V, 4
                cam_trans = w2cs[:, cam_vid, ...] # B, 4, 4
                cam_pos_cam = (cam_pos_world @ cam_trans.transpose(1, 2))[...,:3] # B, V, 4
                points_dirs = cam_xyz[:,:, None, :] - cam_pos_cam[:, None, :, :] # B, N, V, 3 in current cam coord
                points_dirs = points_dirs / (torch.norm(points_dirs, dim=-1, keepdim=True) + 1e-6)  # B, N, V, 3
                points_dirs = points_dirs.view(cam_xyz.shape[0], -1, 3) @ c2ws[:, cam_vid, :3, :3].transpose(1, 2)
                if not pointdir_w:
                    points_dirs = points_dirs @ c2ws[:, self.args.ref_vid, :3, :3].transpose(1, 2) # in ref cam coord
                # print("points_dirs", points_dirs.shape)
                points_dirs = points_dirs.view(cam_xyz.shape[0], cam_xyz.shape[1], -1)
            elif feat_str.startswith("point_conf"):
                if photometric_confidence is None:
                    photometric_confidence = torch.ones_like(points_embedding[0][...,0:1])
                points_conf = photometric_confidence
        points_embedding = torch.cat(points_embedding, dim=-1)
        if self.args.shading_feature_mlp_layer0 > 0:
            points_embedding = self.premlp(torch.cat([points_embedding, points_colors, points_dirs, points_conf], dim=-1))
        return points_embedding, points_colors, points_dirs, points_conf
//  [python]
def query_embedding(self, HDWD, cam_xyz, photometric_confidence, img_feats, c2ws, w2cs, intrinsics, cam_vid, pointdir_w=False):

        HD, WD = HDWD
        points_embedding = []
        points_dirs = None
        points_conf = None
        points_colors = None
        for feat_str in getattr(self.args, feature_str_lst[cam_vid]):
            if feat_str.startswith("imgfeat"):
                _, view_ids, layer_ids = feat_str.split("_")
                view_ids = [int(a) for a in list(view_ids)]
                layer_ids = [int(a) for a in list(layer_ids)]
                twoD_feats, points_colors = self.extract_2d(img_feats, view_ids, layer_ids, intrinsics, c2ws, w2cs, cam_xyz, HD, WD, cam_vid=cam_vid)
                points_embedding.append(twoD_feats)
            elif feat_str.startswith("dir"):
                _, view_ids = feat_str.split("_")
                view_ids = torch.as_tensor([int(a) for a in list(view_ids)], dtype=torch.int64, device=cam_xyz.device)
                cam_pos_world = c2ws[:, view_ids, :, 3] # B, V, 4
                cam_trans = w2cs[:, cam_vid, ...] # B, 4, 4
                cam_pos_cam = (cam_pos_world @ cam_trans.transpose(1, 2))[...,:3] # B, V, 4
                points_dirs = cam_xyz[:,:, None, :] - cam_pos_cam[:, None, :, :] # B, N, V, 3 in current cam coord
                points_dirs = points_dirs / (torch.linalg.norm(points_dirs, dim=-1, keepdims=True) + 1e-6)  # B, N, V, 3
                points_dirs = points_dirs.view(cam_xyz.shape[0], -1, 3) @ c2ws[:, cam_vid, :3, :3].transpose(1, 2)
                if not pointdir_w:
                    points_dirs = points_dirs @ c2ws[:, self.args.ref_vid, :3, :3].transpose(1, 2) # in ref cam coord
                # print("points_dirs", points_dirs.shape)
                points_dirs = points_dirs.view(cam_xyz.shape[0], cam_xyz.shape[1], -1)
            elif feat_str.startswith("point_conf"):
                if photometric_confidence is None:
                    photometric_confidence = torch.ones_like(points_embedding[0][...,0:1])
                points_conf = photometric_confidence
        points_embedding = torch.cat(points_embedding, dim=-1)
        if self.args.shading_feature_mlp_layer0 > 0:
            points_embedding = self.premlp(torch.cat([points_embedding, points_colors, points_dirs, points_conf], dim=-1))
        return points_embedding, points_colors, points_dirs, points_conf
//  [python]
def __init__(self, init_pcl, K=4, trajectory=None, node_radius=None):
        super().__init__()
        self.K = K
        self.N = init_pcl.shape[0]
        nn_dist, nn_idxs, _ = pytorch3d.ops.knn_points(init_pcl[None], init_pcl[None], None, None, K=K+1)  # N, K
        nn_dist, nn_idxs = nn_dist[0,:,1:], nn_idxs[0,:,1:]
        nn_dist = 1 / (nn_dist + 1e-7)
        self.nn_idxs = nn_idxs
        self._weight = nn.Parameter(torch.log(nn_dist / (nn_dist.sum(dim=1, keepdim=True) + 1e-5) + 1e-5))
        self.init_pcl = init_pcl
        self.init_pcl_copy = init_pcl.clone()
        self.tensors = {}
        # self.optimizer = torch.optim.Adam([self._weight], lr=1e-5)
        self.mask_control_points = False
        if self.mask_control_points:
            self.generate_mask_init_pcl()
            radius = torch.norm(self.init_pcl_reduced.max(dim=0).values - self.init_pcl_reduced.min(dim=0).values) / 10 * 3
            print("Set ball query radius to %f" % radius.item())
            self.arap_deformer = ARAPDeformer(self.init_pcl_reduced, radius=radius, K=30, point_mask=self.init_pcl_mask, trajectory=trajectory, node_radius=node_radius)
        else:
            radius = torch.norm(self.init_pcl.max(dim=0).values - self.init_pcl.min(dim=0).values) / 8
            print("Set ball query radius to %f" % radius.item())
            self.arap_deformer = ARAPDeformer(init_pcl, radius=radius, K=16, trajectory=trajectory, node_radius=node_radius)

        self.optimizer = torch.optim.Adam([self.arap_deformer.weight], lr=1e-3)
        self.scheduler = torch.optim.lr_scheduler.StepLR(self.optimizer, step_size=100, gamma=0.99)
        self.optim_step = 0
//  [python]
def __init__(self, init_pcl, K=4, trajectory=None, node_radius=None):
        super().__init__()
        self.K = K
        self.N = init_pcl.shape[0]
        nn_dist, nn_idxs, _ = pytorch3d.ops.knn_points(init_pcl[None], init_pcl[None], None, None, K=K+1)  # N, K
        nn_dist, nn_idxs = nn_dist[0,:,1:], nn_idxs[0,:,1:]
        nn_dist = 1 / (nn_dist + 1e-7)
        self.nn_idxs = nn_idxs
        self._weight = nn.Parameter(torch.log(nn_dist / (nn_dist.sum(dim=1, keepdim=True) + 1e-5) + 1e-5))
        self.init_pcl = init_pcl
        self.init_pcl_copy = init_pcl.clone()
        self.tensors = {}
        # self.optimizer = torch.optim.Adam([self._weight], lr=1e-5)
        self.mask_control_points = False
        if self.mask_control_points:
            self.generate_mask_init_pcl()
            radius = torch.linalg.norm(self.init_pcl_reduced.max(dim=0).values - self.init_pcl_reduced.min(dim=0).values) / 10 * 3
            print("Set ball query radius to %f" % radius.item())
            self.arap_deformer = ARAPDeformer(self.init_pcl_reduced, radius=radius, K=30, point_mask=self.init_pcl_mask, trajectory=trajectory, node_radius=node_radius)
        else:
            radius = torch.linalg.norm(self.init_pcl.max(dim=0).values - self.init_pcl.min(dim=0).values) / 8
            print("Set ball query radius to %f" % radius.item())
            self.arap_deformer = ARAPDeformer(init_pcl, radius=radius, K=16, trajectory=trajectory, node_radius=node_radius)

        self.optimizer = torch.optim.Adam([self.arap_deformer.weight], lr=1e-3)
        self.scheduler = torch.optim.lr_scheduler.StepLR(self.optimizer, step_size=100, gamma=0.99)
        self.optim_step = 0
//  [python]
def split_aug(self, profile: torch.Tensor, binary_labels: torch.Tensor, mask: torch.Tensor):
        # B, N
        bsz, dim = profile.shape[0], profile.shape[-1]
        profile_norm = torch.norm(profile, dim=-1, keepdim=False)
        spk_count = binary_labels.sum(dim=1)
        prob = np.random.rand(bsz)
        batch_indices = np.nonzero(prob < self.split_aug_prob)[0]
        for idx in batch_indices:
            valid_spk_idx = torch.nonzero(spk_count[idx] * mask[idx])
            pad_spk_idx = torch.nonzero((spk_count[idx] == 0) * mask[idx])
            if len(valid_spk_idx) == 0 or len(pad_spk_idx) == 0:
                continue
            split_spk_idx = valid_spk_idx[torch.randint(len(valid_spk_idx), ())]
            to_cover_idx = pad_spk_idx[torch.randint(len(pad_spk_idx), ())]
            disturb_vec = torch.randn((dim,)).to(profile)
            disturb_vec = F.normalize(disturb_vec, dim=-1)
            profile[idx, to_cover_idx] = F.normalize(profile[idx, split_spk_idx] +
                                                     self.disturb_alpha * disturb_vec)
            mask[idx, split_spk_idx] = 0
            mask[idx, to_cover_idx] = 0
        return profile, binary_labels, mask
//  [python]
def split_aug(self, profile: torch.Tensor, binary_labels: torch.Tensor, mask: torch.Tensor):
        # B, N
        bsz, dim = profile.shape[0], profile.shape[-1]
        profile_norm = torch.linalg.norm(profile, dim=-1, keepdim=False)
        spk_count = binary_labels.sum(dim=1)
        prob = np.random.rand(bsz)
        batch_indices = np.nonzero(prob < self.split_aug_prob)[0]
        for idx in batch_indices:
            valid_spk_idx = torch.nonzero(spk_count[idx] * mask[idx])
            pad_spk_idx = torch.nonzero((spk_count[idx] == 0) * mask[idx])
            if len(valid_spk_idx) == 0 or len(pad_spk_idx) == 0:
                continue
            split_spk_idx = valid_spk_idx[torch.randint(len(valid_spk_idx), ())]
            to_cover_idx = pad_spk_idx[torch.randint(len(pad_spk_idx), ())]
            disturb_vec = torch.randn((dim,)).to(profile)
            disturb_vec = F.normalize(disturb_vec, dim=-1)
            profile[idx, to_cover_idx] = F.normalize(profile[idx, split_spk_idx] +
                                                     self.disturb_alpha * disturb_vec)
            mask[idx, split_spk_idx] = 0
            mask[idx, to_cover_idx] = 0
        return profile, binary_labels, mask
//  [python]
def merge_aug(self, profile: torch.Tensor, binary_labels: torch.Tensor, mask: torch.Tensor):
        bsz, dim = profile.shape[0], profile.shape[-1]
        profile_norm = torch.norm(profile, dim=-1, keepdim=False)
        spk_count = binary_labels.sum(dim=1)
        prob = np.random.rand(bsz)
        batch_indices = np.nonzero(prob < self.merge_aug_prob)[0]
        for idx in batch_indices:
            valid_spk_idx = torch.nonzero(profile_norm[idx] * mask[idx])
            if len(valid_spk_idx) == 0:
                continue
            to_merge = torch.randint(len(valid_spk_idx), (2, ))
            spk_idx_1, spk_idx_2 = valid_spk_idx[to_merge[0]], valid_spk_idx[to_merge[1]]
            # merge profile
            profile[idx, spk_idx_1] = profile[idx, spk_idx_1] + profile[idx, spk_idx_2]
            profile[idx, spk_idx_1] = F.normalize(profile[idx, spk_idx_1], dim=-1)
            profile[idx, spk_idx_2] = 0
            # merge binary labels
            binary_labels[idx, :, spk_idx_1] = binary_labels[idx, :, spk_idx_1] + binary_labels[idx, :, spk_idx_2]
            binary_labels[idx, :, spk_idx_1] = (binary_labels[idx, :, spk_idx_1] > 0).to(binary_labels)
            binary_labels[idx, :, spk_idx_2] = 0

            mask[idx, spk_idx_1] = 0
            mask[idx, spk_idx_2] = 0

        return profile, binary_labels, mask
//  [python]
def merge_aug(self, profile: torch.Tensor, binary_labels: torch.Tensor, mask: torch.Tensor):
        bsz, dim = profile.shape[0], profile.shape[-1]
        profile_norm = torch.linalg.norm(profile, dim=-1, keepdim=False)
        spk_count = binary_labels.sum(dim=1)
        prob = np.random.rand(bsz)
        batch_indices = np.nonzero(prob < self.merge_aug_prob)[0]
        for idx in batch_indices:
            valid_spk_idx = torch.nonzero(profile_norm[idx] * mask[idx])
            if len(valid_spk_idx) == 0:
                continue
            to_merge = torch.randint(len(valid_spk_idx), (2, ))
            spk_idx_1, spk_idx_2 = valid_spk_idx[to_merge[0]], valid_spk_idx[to_merge[1]]
            # merge profile
            profile[idx, spk_idx_1] = profile[idx, spk_idx_1] + profile[idx, spk_idx_2]
            profile[idx, spk_idx_1] = F.normalize(profile[idx, spk_idx_1], dim=-1)
            profile[idx, spk_idx_2] = 0
            # merge binary labels
            binary_labels[idx, :, spk_idx_1] = binary_labels[idx, :, spk_idx_1] + binary_labels[idx, :, spk_idx_2]
            binary_labels[idx, :, spk_idx_1] = (binary_labels[idx, :, spk_idx_1] > 0).to(binary_labels)
            binary_labels[idx, :, spk_idx_2] = 0

            mask[idx, spk_idx_1] = 0
            mask[idx, spk_idx_2] = 0

        return profile, binary_labels, mask