jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def up_sample(self, rays_o, rays_d, z_vals, sdf, n_importance, inv_s):
        """
        Up sampling give a fixed inv_s
        """
        batch_size, n_samples = z_vals.shape

        pts = rays_o[:, None, :] + rays_d[:, None, :] * z_vals[..., :, None]  # n_rays, n_samples, 3
        radius = torch.linalg.norm(pts, ord=2, dim=-1, keepdim=False)
        inside_sphere = (radius[:, :-1] < 1.0) | (radius[:, 1:] < 1.0)
        sdf = sdf.reshape(batch_size, n_samples)

        prev_z_vals, next_z_vals = z_vals[:, :-1], z_vals[:, 1:]
        prev_sdf, next_sdf = sdf[:, :-1], sdf[:, 1:]
        
        mid_sdf = (prev_sdf + next_sdf) * 0.5
        cos_val = (next_sdf - prev_sdf) / (next_z_vals - prev_z_vals + 1e-5)

        # ----------------------------------------------------------------------------------------------------------
        # Use min value of [ cos, prev_cos ]
        # Though it makes the sampling (not rendering) a little bit biased, this strategy can make the sampling more
        # robust when meeting situations like below:
        #
        # SDF
        # ^
        # |\          -----x----...
        # | \        /
        # |  x      x
        # |---\----/-------------> 0 level
        # |    \  /
        # |     \/
        # |
        # ----------------------------------------------------------------------------------------------------------
        prev_cos_val = torch.cat([torch.zeros([batch_size, 1], device=sdf.device), cos_val[:, :-1]], dim=-1)
        cos_val = torch.stack([prev_cos_val, cos_val], dim=-1)
        cos_val, _ = torch.min(cos_val, dim=-1, keepdim=False)
        cos_val = cos_val.clip(-1e3, 0.0) * inside_sphere

        dist = (next_z_vals - prev_z_vals)
        prev_esti_sdf = mid_sdf - cos_val * dist * 0.5
        next_esti_sdf = mid_sdf + cos_val * dist * 0.5
        prev_cdf = torch.sigmoid(prev_esti_sdf * inv_s)
        next_cdf = torch.sigmoid(next_esti_sdf * inv_s)
        alpha = (prev_cdf - next_cdf + 1e-5) / (prev_cdf + 1e-5)

        weights = alpha * torch.cumprod(
            torch.cat([torch.ones([batch_size, 1], device=alpha.device), 1. - alpha + 1e-7], -1), -1)[:, :-1]

        z_samples = sample_pdf(z_vals, weights, n_importance, det=True).detach()

        return z_samples
//  [python]
def prepare_rays_a_view(self, img_idx):
        """
        Generate random rays at world space from one camera.
        """
        tx = torch.linspace(0, self.W - 1, self.W)
        ty = torch.linspace(0, self.H - 1, self.H)
        pixels_x, pixels_y = torch.meshgrid(tx, ty)

        pixels_x = pixels_x.reshape(-1).long()
        pixels_y = pixels_y.reshape(-1).long()
        color = self.images[img_idx][(pixels_y, pixels_x)]    # batch_size, 3
        mask = self.masks[img_idx][(pixels_y, pixels_x)]      # batch_size, 3
        normal = self.normals_world[img_idx][(pixels_y, pixels_x)]      # batch_size, 3
        
        q = torch.stack([(pixels_x / self.W-0.5)*2, (pixels_y / self.H-0.5)*2, torch.zeros_like(pixels_y)], dim=-1).float()  # batch_size, 3
        v = torch.stack([torch.zeros_like(pixels_y), torch.zeros_like(pixels_y), torch.ones_like(pixels_y)], dim=-1).float()
        
        rays_v = v / torch.norm(v, p=2, dim=-1, keepdim=True)    # batch_size, 3
        rays_v = torch.matmul(self.pose_all[img_idx, None, :3, :3], rays_v[:, :, None]).squeeze()  # batch_size, 3
        
        rays_o = torch.matmul(self.pose_all[img_idx, None, :3, :3], q[:, :, None]).squeeze()  # batch_size, 3
        rays_o = self.pose_all[img_idx, None, :3, 3].expand(rays_v.shape) + rays_o # batch_size, 3

        cosines = self.cos(rays_v, normal)

        # pdb.set_trace()
        
        return torch.cat([rays_o.cpu(), rays_v.cpu(), color, mask[:, None], normal, cosines[:, None]], dim=-1)
//  [python]
def prepare_rays_a_view(self, img_idx):
        """
        Generate random rays at world space from one camera.
        """
        tx = torch.linspace(0, self.W - 1, self.W)
        ty = torch.linspace(0, self.H - 1, self.H)
        pixels_x, pixels_y = torch.meshgrid(tx, ty)

        pixels_x = pixels_x.reshape(-1).long()
        pixels_y = pixels_y.reshape(-1).long()
        color = self.images[img_idx][(pixels_y, pixels_x)]    # batch_size, 3
        mask = self.masks[img_idx][(pixels_y, pixels_x)]      # batch_size, 3
        normal = self.normals_world[img_idx][(pixels_y, pixels_x)]      # batch_size, 3
        
        q = torch.stack([(pixels_x / self.W-0.5)*2, (pixels_y / self.H-0.5)*2, torch.zeros_like(pixels_y)], dim=-1).float()  # batch_size, 3
        v = torch.stack([torch.zeros_like(pixels_y), torch.zeros_like(pixels_y), torch.ones_like(pixels_y)], dim=-1).float()
        
        rays_v = v / torch.linalg.norm(v, ord=2, dim=-1, keepdim=True)    # batch_size, 3
        rays_v = torch.matmul(self.pose_all[img_idx, None, :3, :3], rays_v[:, :, None]).squeeze()  # batch_size, 3
        
        rays_o = torch.matmul(self.pose_all[img_idx, None, :3, :3], q[:, :, None]).squeeze()  # batch_size, 3
        rays_o = self.pose_all[img_idx, None, :3, 3].expand(rays_v.shape) + rays_o # batch_size, 3

        cosines = self.cos(rays_v, normal)

        # pdb.set_trace()
        
        return torch.cat([rays_o.cpu(), rays_v.cpu(), color, mask[:, None], normal, cosines[:, None]], dim=-1)
//  [python]
def _encode_image(self, retrieved_images, batch_size):
        if len(retrieved_images[0]) == 0:
            return None
        for i in range(len(retrieved_images)):
            retrieved_images[i] = normalize_images(retrieved_images[i])
            retrieved_images[i] = preprocess_images(retrieved_images[i], self.feature_extractor).to(
                self.clip.device, dtype=self.clip.dtype
            )
        _, c, h, w = retrieved_images[0].shape

        retrieved_images = torch.reshape(torch.cat(retrieved_images, dim=0), (-1, c, h, w))
        image_embeddings = self.clip.get_image_features(retrieved_images)
        image_embeddings = image_embeddings / torch.norm(image_embeddings, dim=-1, keepdim=True)
        _, d = image_embeddings.shape
        image_embeddings = torch.reshape(image_embeddings, (batch_size, -1, d))
        return image_embeddings
//  [python]
def _encode_image(self, retrieved_images, batch_size):
        if len(retrieved_images[0]) == 0:
            return None
        for i in range(len(retrieved_images)):
            retrieved_images[i] = normalize_images(retrieved_images[i])
            retrieved_images[i] = preprocess_images(retrieved_images[i], self.feature_extractor).to(
                self.clip.device, dtype=self.clip.dtype
            )
        _, c, h, w = retrieved_images[0].shape

        retrieved_images = torch.reshape(torch.cat(retrieved_images, dim=0), (-1, c, h, w))
        image_embeddings = self.clip.get_image_features(retrieved_images)
        image_embeddings = image_embeddings / torch.linalg.norm(image_embeddings, dim=-1, keepdim=True)
        _, d = image_embeddings.shape
        image_embeddings = torch.reshape(image_embeddings, (batch_size, -1, d))
        return image_embeddings
//  [python]
def log_param_norms(model, optimizer, param_id_to_buffer):
    weight_norms = {}
    other_norms = {}
    for name, param in model.named_parameters():
        buf = param_id_to_buffer[id(param)]
        param_norm = torch.norm(buf) ** 2
        other_norm = torch.norm(param.ds_tensor.data) ** 2
        torch.distributed.all_reduce(param_norm, group=optimizer.ds_param_shard_group)
        torch.distributed.all_reduce(other_norm, group=optimizer.ds_param_shard_group)
        weight_norms[name] = torch.sqrt(param_norm).item()
        other_norms[name] = torch.sqrt(other_norm).item()
        if smp.rank() == 0:
            print(f"{name}: {weight_norms[name]} {other_norms[name]}")
//  [python]
def log_param_norms(model, optimizer, param_id_to_buffer):
    weight_norms = {}
    other_norms = {}
    for name, param in model.named_parameters():
        buf = param_id_to_buffer[id(param)]
        param_norm = torch.linalg.norm(buf) ** 2
        other_norm = torch.linalg.norm(param.ds_tensor.data) ** 2
        torch.distributed.all_reduce(param_norm, group=optimizer.ds_param_shard_group)
        torch.distributed.all_reduce(other_norm, group=optimizer.ds_param_shard_group)
        weight_norms[name] = torch.sqrt(param_norm).item()
        other_norms[name] = torch.sqrt(other_norm).item()
        if smp.rank() == 0:
            print(f"{name}: {weight_norms[name]} {other_norms[name]}")
//  [python]
def forward(self, text):
        tokens = self.tokenize(self.tokenizer, text, context_length=self.max_length).to(self.device)
        z = self.model(tokens)
        if self.normalize:
            z = z / torch.norm(z, dim=1, keepdim=True)
        return z
//  [python]
def forward(self, text):
        tokens = self.tokenize(self.tokenizer, text, context_length=self.max_length).to(self.device)
        z = self.model(tokens)
        if self.normalize:
            z = z / torch.linalg.norm(z, dim=1, keepdim=True)
        return z
//  [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 __init__(
        self,
        rot_mats: Optional[Tensor] = None,
        quats: Optional[Tensor] = None,
        normalize_quats: bool = True,
    ):
        """
        Args:
            rot_mats:
                A [*, 3, 3] rotation matrix tensor. Mutually exclusive with
                quats
            quats:
                A [*, 4] quaternion. Mutually exclusive with rot_mats. If
                normalize_quats is not True, must be a unit quaternion
            normalize_quats:
                If quats is specified, whether to normalize quats
        """
        if (rot_mats is None and quats is None) or (rot_mats is not None and quats is not None):
            raise ValueError("Exactly one input argument must be specified")

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

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

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

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

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

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

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

        self._rot_mats = rot_mats
        self._quats = quats
//  [python]
def forward(self, inputs):
        inputs = inputs * self.scale
        if self.embed_fn_fine is not None:
            inputs = self.embed_fn_fine(inputs)

        x = inputs
        for l in range(0, self.num_layers - 1):
            lin = getattr(self, "lin" + str(l))

            if l in self.skip_in:
                x = torch.cat([x, inputs], 1) / np.sqrt(2)

            x = lin(x)

            if l < self.num_layers - 2:
                x = self.activation(x)

        if self.use_normal_pred:
            normal = x[:, 1:4] / torch.norm(x[:, 1:4], dim=-1, keepdim=True)
            return torch.cat([x[:, :1] / self.scale, normal, x[:, 4:]], dim=-1)
        else:
            return torch.cat([x[:, :1] / self.scale, x[:, 1:]], dim=-1)
//  [python]
def forward(self, inputs):
        inputs = inputs * self.scale
        if self.embed_fn_fine is not None:
            inputs = self.embed_fn_fine(inputs)

        x = inputs
        for l in range(0, self.num_layers - 1):
            lin = getattr(self, "lin" + str(l))

            if l in self.skip_in:
                x = torch.cat([x, inputs], 1) / np.sqrt(2)

            x = lin(x)

            if l < self.num_layers - 2:
                x = self.activation(x)

        if self.use_normal_pred:
            normal = x[:, 1:4] / torch.linalg.norm(x[:, 1:4], dim=-1, keepdim=True)
            return torch.cat([x[:, :1] / self.scale, normal, x[:, 4:]], dim=-1)
        else:
            return torch.cat([x[:, :1] / self.scale, x[:, 1:]], dim=-1)
//  [python]
def get_traj_cls_target(self,
             traj_preds,
             traj_targets,
             gt_fut_masks,
             neg_inds):
        """"Get Trajectory mode classification target.
        Args:
            traj_preds (Tensor): MultiModal traj preds with shape (num_box_preds, fut_mode, fut_ts, 2).
            traj_targets (Tensor): Ground truth traj for each pred box with shape (num_box_preds, fut_ts, 2).
            gt_fut_masks (Tensor): Ground truth traj mask with shape (num_box_preds, fut_ts).
            neg_inds (Tensor): Negtive indices with shape (num_box_preds,)

        Returns:
            traj_labels (Tensor): traj cls labels (num_box_preds,).
        """

        cum_traj_preds = traj_preds.cumsum(dim=-2)
        cum_traj_targets = traj_targets.cumsum(dim=-2)

        # Get min pred mode indices.
        # (num_box_preds, fut_mode, fut_ts)
        dist = torch.norm(cum_traj_targets[:, None, :, :] - cum_traj_preds, dim=-1)
        dist = dist * gt_fut_masks[:, None, :]
        dist = dist[..., -1]
        dist[torch.isnan(dist)] = dist[torch.isnan(dist)] * 0
        traj_labels = torch.argmin(dist, dim=-1)
        traj_labels[neg_inds] = self.fut_mode

        return traj_labels
//  [python]
def get_traj_cls_target(self,
             traj_preds,
             traj_targets,
             gt_fut_masks,
             neg_inds):
        """"Get Trajectory mode classification target.
        Args:
            traj_preds (Tensor): MultiModal traj preds with shape (num_box_preds, fut_mode, fut_ts, 2).
            traj_targets (Tensor): Ground truth traj for each pred box with shape (num_box_preds, fut_ts, 2).
            gt_fut_masks (Tensor): Ground truth traj mask with shape (num_box_preds, fut_ts).
            neg_inds (Tensor): Negtive indices with shape (num_box_preds,)

        Returns:
            traj_labels (Tensor): traj cls labels (num_box_preds,).
        """

        cum_traj_preds = traj_preds.cumsum(dim=-2)
        cum_traj_targets = traj_targets.cumsum(dim=-2)

        # Get min pred mode indices.
        # (num_box_preds, fut_mode, fut_ts)
        dist = torch.linalg.norm(cum_traj_targets[:, None, :, :] - cum_traj_preds, dim=-1)
        dist = dist * gt_fut_masks[:, None, :]
        dist = dist[..., -1]
        dist[torch.isnan(dist)] = dist[torch.isnan(dist)] * 0
        traj_labels = torch.argmin(dist, dim=-1)
        traj_labels[neg_inds] = self.fut_mode

        return traj_labels
//  [python]
def get_best_fut_preds(self,
             traj_preds,
             traj_targets,
             gt_fut_masks):
        """"Choose best preds among all modes.
        Args:
            traj_preds (Tensor): MultiModal traj preds with shape (num_box_preds, fut_mode, fut_ts, 2).
            traj_targets (Tensor): Ground truth traj for each pred box with shape (num_box_preds, fut_ts, 2).
            gt_fut_masks (Tensor): Ground truth traj mask with shape (num_box_preds, fut_ts).
            pred_box_centers (Tensor): Pred box centers with shape (num_box_preds, 2).
            gt_box_centers (Tensor): Ground truth box centers with shape (num_box_preds, 2).

        Returns:
            best_traj_preds (Tensor): best traj preds (min displacement error with gt)
                with shape (num_box_preds, fut_ts*2).
        """

        cum_traj_preds = traj_preds.cumsum(dim=-2)
        cum_traj_targets = traj_targets.cumsum(dim=-2)

        # Get min pred mode indices.
        # (num_box_preds, fut_mode, fut_ts)
        dist = torch.norm(cum_traj_targets[:, None, :, :] - cum_traj_preds, dim=-1)
        dist = dist * gt_fut_masks[:, None, :]
        dist = dist[..., -1]
        dist[torch.isnan(dist)] = dist[torch.isnan(dist)] * 0
        min_mode_idxs = torch.argmin(dist, dim=-1).tolist()
        box_idxs = torch.arange(traj_preds.shape[0]).tolist()
        best_traj_preds = traj_preds[box_idxs, min_mode_idxs, :, :].reshape(-1, self.fut_ts*2)

        return best_traj_preds
//  [python]
def get_best_fut_preds(self,
             traj_preds,
             traj_targets,
             gt_fut_masks):
        """"Choose best preds among all modes.
        Args:
            traj_preds (Tensor): MultiModal traj preds with shape (num_box_preds, fut_mode, fut_ts, 2).
            traj_targets (Tensor): Ground truth traj for each pred box with shape (num_box_preds, fut_ts, 2).
            gt_fut_masks (Tensor): Ground truth traj mask with shape (num_box_preds, fut_ts).
            pred_box_centers (Tensor): Pred box centers with shape (num_box_preds, 2).
            gt_box_centers (Tensor): Ground truth box centers with shape (num_box_preds, 2).

        Returns:
            best_traj_preds (Tensor): best traj preds (min displacement error with gt)
                with shape (num_box_preds, fut_ts*2).
        """

        cum_traj_preds = traj_preds.cumsum(dim=-2)
        cum_traj_targets = traj_targets.cumsum(dim=-2)

        # Get min pred mode indices.
        # (num_box_preds, fut_mode, fut_ts)
        dist = torch.linalg.norm(cum_traj_targets[:, None, :, :] - cum_traj_preds, dim=-1)
        dist = dist * gt_fut_masks[:, None, :]
        dist = dist[..., -1]
        dist[torch.isnan(dist)] = dist[torch.isnan(dist)] * 0
        min_mode_idxs = torch.argmin(dist, dim=-1).tolist()
        box_idxs = torch.arange(traj_preds.shape[0]).tolist()
        best_traj_preds = traj_preds[box_idxs, min_mode_idxs, :, :].reshape(-1, self.fut_ts*2)

        return best_traj_preds
//  [python]
def so3exp(self, w):
        """
            Maps so(3) --> SO(3) group with closed form expression.
        """
        w.squeeze()
        theta = torch.norm(w)
        if theta < 1e-8:
            return torch.eye(3, device=w.device, dtype=w.dtype)
        else:
            w_hat = self.skew_symmetric(w)
            R = torch.eye(3, device=w.device, dtype=w.dtype) + (torch.sin(theta) / theta) * w_hat + ((1 - torch.cos(theta)) / theta**2) * (w_hat @ w_hat)
            return R
//  [python]
def so3exp(self, w):
        """
            Maps so(3) --> SO(3) group with closed form expression.
        """
        w.squeeze()
        theta = torch.linalg.norm(w)
        if theta < 1e-8:
            return torch.eye(3, device=w.device, dtype=w.dtype)
        else:
            w_hat = self.skew_symmetric(w)
            R = torch.eye(3, device=w.device, dtype=w.dtype) + (torch.sin(theta) / theta) * w_hat + ((1 - torch.cos(theta)) / theta**2) * (w_hat @ w_hat)
            return R
//  [python]
def test_norm_vector_degenerate_shapes(self, device, dtype):
        def run_test_case(input, ord, dim, keepdim, should_error):
            msg = f'input.size()={input.size()}, ord={ord}, dim={dim}, keepdim={keepdim}, dtype={dtype}'
            input_numpy = input.cpu().numpy()
            if should_error:
                with self.assertRaises(ValueError):
                    np.linalg.norm(input_numpy, ord, dim, keepdim)
                with self.assertRaises(RuntimeError):
                    t.norm(input, ord, dim, keepdim)
            else:
                result_numpy = np.linalg.norm(input_numpy, ord, dim, keepdim)
                result = t.norm(input, ord, dim, keepdim)
                self.assertEqual(result, result_numpy, msg=msg)

        ord_vector = [0, 0.5, 1, 2, 3, inf, -0.5, -1, -2, -3, -inf, None]
        S = 10
        test_cases = [
            # input size, p settings that cause error, dim
            ((0, ), [inf, -inf], None),
            ((0, S), [inf, -inf], 0),
            ((0, S), [], 1),
            ((S, 0), [], 0),
            ((S, 0), [inf, -inf], 1),
        ]
        for keepdim in [True, False]:
            for input_size, error_ords, dim in test_cases:
                input = torch.randn(*input_size, dtype=dtype, device=device)
                for ord in ord_vector:
                    run_test_case(input, ord, dim, keepdim, ord in error_ords)
//  [python]
def test_norm_vector_degenerate_shapes(self, device, dtype):
        def run_test_case(input, ord, dim, keepdim, should_error):
            msg = f'input.size()={input.size()}, ord={ord}, dim={dim}, keepdim={keepdim}, dtype={dtype}'
            input_numpy = input.cpu().numpy()
            if should_error:
                with self.assertRaises(ValueError):
                    np.linalg.norm(input_numpy, ord, dim, keepdim)
                with self.assertRaises(RuntimeError):
                    torch.linalg.norm(input, ord, dim, keepdim)
            else:
                result_numpy = np.linalg.norm(input_numpy, ord, dim, keepdim)
                result = torch.linalg.norm(input, ord, dim, keepdim)
                self.assertEqual(result, result_numpy, msg=msg)

        ord_vector = [0, 0.5, 1, 2, 3, inf, -0.5, -1, -2, -3, -inf, None]
        S = 10
        test_cases = [
            # input size, p settings that cause error, dim
            ((0, ), [inf, -inf], None),
            ((0, S), [inf, -inf], 0),
            ((0, S), [], 1),
            ((S, 0), [], 0),
            ((S, 0), [inf, -inf], 1),
        ]
        for keepdim in [True, False]:
            for input_size, error_ords, dim in test_cases:
                input = torch.randn(*input_size, dtype=dtype, device=device)
                for ord in ord_vector:
                    run_test_case(input, ord, dim, keepdim, ord in error_ords)
//  [python]
def test_norm_complex(self, device, dtype):
        def gen_error_message(input_size, ord, keepdim, dim=None):
            return "complex norm failed for input size %s, ord=%s, keepdim=%s, dim=%s" % (
                input_size, ord, keepdim, dim)

        vector_ords = [None, 0, 1, 2, 3, inf, -1, -2, -3, -inf]
        matrix_ords = [None, 'fro', 'nuc', 1, 2, inf, -1, -2, -inf]

        # Test supported ords
        for keepdim in [False, True]:
            # vector norm
            x = torch.randn(25, device=device, dtype=dtype)
            xn = x.cpu().numpy()
            for ord in vector_ords:
                res = torch.norm(x, ord, keepdim=keepdim).cpu()
                expected = np.linalg.norm(xn, ord, keepdims=keepdim)
                msg = gen_error_message(x.size(), ord, keepdim)
                self.assertEqual(res.shape, expected.shape, msg=msg)
                self.assertEqual(res, expected, msg=msg)

                res_out = torch.Tensor().to(device)
                torch.norm(x, ord, keepdim=keepdim, out=res_out)
                self.assertEqual(res_out.shape, expected.shape, msg=msg)
                self.assertEqual(res_out.cpu(), expected, msg=msg)

            # matrix norm
            x = torch.randn(25, 25, device=device, dtype=dtype)
            xn = x.cpu().numpy()
            for ord in matrix_ords:
                res = torch.norm(x, ord, keepdim=keepdim).cpu()
                expected = np.linalg.norm(xn, ord, keepdims=keepdim)
                msg = gen_error_message(x.size(), ord, keepdim)
                self.assertEqual(res.shape, expected.shape, msg=msg)
                self.assertEqual(res, expected, msg=msg)

                res_out = torch.Tensor().to(device)
                torch.norm(x, ord, keepdim=keepdim, out=res_out)
                self.assertEqual(res_out.shape, expected.shape, msg=msg)
                self.assertEqual(res_out.cpu(), expected, msg=msg)
//  [python]
def test_norm_complex(self, device, dtype):
        def gen_error_message(input_size, ord, keepdim, dim=None):
            return "complex norm failed for input size %s, ord=%s, keepdim=%s, dim=%s" % (
                input_size, ord, keepdim, dim)

        vector_ords = [None, 0, 1, 2, 3, inf, -1, -2, -3, -inf]
        matrix_ords = [None, 'fro', 'nuc', 1, 2, inf, -1, -2, -inf]

        # Test supported ords
        for keepdim in [False, True]:
            # vector norm
            x = torch.randn(25, device=device, dtype=dtype)
            xn = x.cpu().numpy()
            for ord in vector_ords:
                res = torch.linalg.norm(x, ord, keepdim=keepdim).cpu()
                expected = np.linalg.norm(xn, ord, keepdims=keepdim)
                msg = gen_error_message(x.size(), ord, keepdim)
                self.assertEqual(res.shape, expected.shape, msg=msg)
                self.assertEqual(res, expected, msg=msg)

                res_out = torch.Tensor().to(device)
                torch.linalg.norm(x, ord, keepdim=keepdim, out=res_out)
                self.assertEqual(res_out.shape, expected.shape, msg=msg)
                self.assertEqual(res_out.cpu(), expected, msg=msg)

            # matrix norm
            x = torch.randn(25, 25, device=device, dtype=dtype)
            xn = x.cpu().numpy()
            for ord in matrix_ords:
                res = torch.linalg.norm(x, ord, keepdim=keepdim).cpu()
                expected = np.linalg.norm(xn, ord, keepdims=keepdim)
                msg = gen_error_message(x.size(), ord, keepdim)
                self.assertEqual(res.shape, expected.shape, msg=msg)
                self.assertEqual(res, expected, msg=msg)

                res_out = torch.Tensor().to(device)
                torch.linalg.norm(x, ord, keepdim=keepdim, out=res_out)
                self.assertEqual(res_out.shape, expected.shape, msg=msg)
                self.assertEqual(res_out.cpu(), expected, msg=msg)
//  [python]
def test_norm_matrix_degenerate_shapes(self, device, dtype):
        def run_test_case(input, ord, dim, keepdim, should_error):
            msg = f'input.size()={input.size()}, ord={ord}, dim={dim}, keepdim={keepdim}, dtype={dtype}'
            input_numpy = input.cpu().numpy()
            if should_error:
                with self.assertRaises(ValueError):
                    np.linalg.norm(input_numpy, ord, dim, keepdim)
                with self.assertRaises(RuntimeError):
                    t.norm(input, ord, dim, keepdim)
            else:
                result_numpy = np.linalg.norm(input_numpy, ord, dim, keepdim)
                result = t.norm(input, ord, dim, keepdim)
                self.assertEqual(result, result_numpy, msg=msg)

        ord_matrix = ['fro', 'nuc', 1, 2, inf, -1, -2, -inf, None]
        S = 10
        test_cases = [
            # input size, p settings that cause error, dim
            ((0, 0), [1, 2, inf, -1, -2, -inf], None),
            ((0, S), [2, inf, -2, -inf], None),
            ((S, 0), [1, 2, -1, -2], None),
            ((S, S, 0), [], (0, 1)),
            ((1, S, 0), [], (0, 1)),
            ((0, 0, S), [1, 2, inf, -1, -2, -inf], (0, 1)),
            ((0, 0, S), [1, 2, inf, -1, -2, -inf], (1, 0)),
        ]
        for keepdim in [True, False]:
            for input_size, error_ords, dim in test_cases:
                input = torch.randn(*input_size, dtype=dtype, device=device)
                for ord in ord_matrix:
                    run_test_case(input, ord, dim, keepdim, ord in error_ords)
//  [python]
def test_norm_matrix_degenerate_shapes(self, device, dtype):
        def run_test_case(input, ord, dim, keepdim, should_error):
            msg = f'input.size()={input.size()}, ord={ord}, dim={dim}, keepdim={keepdim}, dtype={dtype}'
            input_numpy = input.cpu().numpy()
            if should_error:
                with self.assertRaises(ValueError):
                    np.linalg.norm(input_numpy, ord, dim, keepdim)
                with self.assertRaises(RuntimeError):
                    torch.linalg.norm(input, ord, dim, keepdim)
            else:
                result_numpy = np.linalg.norm(input_numpy, ord, dim, keepdim)
                result = torch.linalg.norm(input, ord, dim, keepdim)
                self.assertEqual(result, result_numpy, msg=msg)

        ord_matrix = ['fro', 'nuc', 1, 2, inf, -1, -2, -inf, None]
        S = 10
        test_cases = [
            # input size, p settings that cause error, dim
            ((0, 0), [1, 2, inf, -1, -2, -inf], None),
            ((0, S), [2, inf, -2, -inf], None),
            ((S, 0), [1, 2, -1, -2], None),
            ((S, S, 0), [], (0, 1)),
            ((1, S, 0), [], (0, 1)),
            ((0, 0, S), [1, 2, inf, -1, -2, -inf], (0, 1)),
            ((0, 0, S), [1, 2, inf, -1, -2, -inf], (1, 0)),
        ]
        for keepdim in [True, False]:
            for input_size, error_ords, dim in test_cases:
                input = torch.randn(*input_size, dtype=dtype, device=device)
                for ord in ord_matrix:
                    run_test_case(input, ord, dim, keepdim, ord in error_ords)
//  [python]
def build(self, device):
        x, y = torch.meshgrid(
            torch.arange(self.width, device=device),
            torch.arange(self.height, device=device),
            indexing="xy",
        )
        directions = F.pad(
            torch.stack(
                [
                    (x - self.K[0, 2] + 0.5) / self.K[0, 0],
                    (y - self.K[1, 2] + 0.5) / self.K[1, 1] * self.sign_z,
                ],
                dim=-1,
            ),
            (0, 1),
            value=self.sign_z,
        )  # [H,W,3]
        # Distance from each unit-norm direction vector to its x-axis neighbor
        dx = torch.norm(
            (directions[:, :-1, :] - directions[:, 1:, :]),
            dim=-1,
            keepdim=True,
        )  # [H,W-1,1]
        dx = torch.cat([dx, dx[:, -2:-1, :]], 1)  # [H,W,1]
        dy = torch.norm(
            (directions[:-1, :, :] - directions[1:, :, :]),
            dim=-1,
            keepdim=True,
        )  # [H-1,W,1]
        dy = torch.cat([dy, dy[-2:-1, :, :]], 0)  # [H,W,1]
        # Cut the distance in half, and then round it out so that it's
        # halfway between inscribed by / circumscribed about the pixel.
        area = dx * dy
        radii = torch.sqrt(area / torch.pi)
        if self.normalize_ray:
            directions = directions / torch.norm(
                directions, dim=-1, keepdim=True
            )
        self.ray_bundle = RayBundle(
            origins=torch.zeros_like(directions),
            directions=directions,
            radiis=radii,
            ray_cos=torch.matmul(
                directions,
                torch.tensor([[0.0, 0.0, self.sign_z]], device=device).T,
            ),
        )
        return self.ray_bundle
//  [python]
def build(self, device):
        x, y = torch.meshgrid(
            torch.arange(self.width, device=device),
            torch.arange(self.height, device=device),
            indexing="xy",
        )
        directions = F.pad(
            torch.stack(
                [
                    (x - self.K[0, 2] + 0.5) / self.K[0, 0],
                    (y - self.K[1, 2] + 0.5) / self.K[1, 1] * self.sign_z,
                ],
                dim=-1,
            ),
            (0, 1),
            value=self.sign_z,
        )  # [H,W,3]
        # Distance from each unit-norm direction vector to its x-axis neighbor
        dx = torch.linalg.norm(
            (directions[:, :-1, :] - directions[:, 1:, :]),
            dim=-1,
            keepdims=True,
        )  # [H,W-1,1]
        dx = torch.cat([dx, dx[:, -2:-1, :]], 1)  # [H,W,1]
        dy = torch.linalg.norm(
            (directions[:-1, :, :] - directions[1:, :, :]),
            dim=-1,
            keepdims=True,
        )  # [H-1,W,1]
        dy = torch.cat([dy, dy[-2:-1, :, :]], 0)  # [H,W,1]
        # Cut the distance in half, and then round it out so that it's
        # halfway between inscribed by / circumscribed about the pixel.
        area = dx * dy
        radii = torch.sqrt(area / torch.pi)
        if self.normalize_ray:
            directions = directions / torch.linalg.norm(
                directions, dim=-1, keepdims=True
            )
        self.ray_bundle = RayBundle(
            origins=torch.zeros_like(directions),
            directions=directions,
            radiis=radii,
            ray_cos=torch.matmul(
                directions,
                torch.tensor([[0.0, 0.0, self.sign_z]], device=device).T,
            ),
        )
        return self.ray_bundle
//  [python]
def main() -> None:
    core_v2.init(
        defaults=core_v2.DefaultConfig(
            name="pytorch-profiler-sync-test",
        ),
    )

    with torch.profiler.profile(
        activities=[torch.profiler.ProfilerActivity.CPU],
        with_flops=True,
        on_trace_ready=torch.profiler.tensorboard_trace_handler(
            str(core_v2.train.get_tensorboard_path())
        ),
    ):
        state = torch.rand(100, 100)
        for _i in range(1000):
            state = torch.matmul(state, torch.rand(100, 100))
            norm = torch.norm(state)
            state = state / norm

    core_v2.close()
//  [python]
def main() -> None:
    core_v2.init(
        defaults=core_v2.DefaultConfig(
            name="pytorch-profiler-sync-test",
        ),
    )

    with torch.profiler.profile(
        activities=[torch.profiler.ProfilerActivity.CPU],
        with_flops=True,
        on_trace_ready=torch.profiler.tensorboard_trace_handler(
            str(core_v2.train.get_tensorboard_path())
        ),
    ):
        state = torch.rand(100, 100)
        for _i in range(1000):
            state = torch.matmul(state, torch.rand(100, 100))
            norm = torch.linalg.norm(state)
            state = state / norm

    core_v2.close()
//  [python]
def unit(v: Tensor, dim: int = -1) -> Tensor:
    '''
    Unit vectors
    '''
    return v / torch.norm(v, dim=dim, keepdim=True)
//  [python]
def unit(v: Tensor, dim: int = -1) -> Tensor:
    '''
    Unit vectors
    '''
    return v / torch.linalg.norm(v, dim=dim, keepdims=True)
//  [python]
def distance_to_depth(K, dist, uv=None):
    if uv is None and len(dist.shape) >= 2:
        # create mesh grid according to d
        uv = np.stack(np.meshgrid(np.arange(dist.shape[1]), np.arange(dist.shape[0])), -1)
        uv = uv.reshape(-1, 2)
        dist = dist.reshape(-1)
        if not isinstance(dist, np.ndarray):
            uv = torch.from_numpy(uv).to(dist)
    if isinstance(dist, np.ndarray):
        # z * np.sqrt(x_temp**2+y_temp**2+z_temp**2) = dist
        uvh = np.concatenate([uv, np.ones((len(uv), 1))], -1)
        temp_point = dot(np.linalg.inv(K), uvh)
        z = dist / np.linalg.norm(temp_point, axis=1)

    else:
        uvh = torch.cat([uv, torch.ones(len(uv), 1).to(uv)], -1)
        temp_point = dot(torch.inverse(K), uvh)
        z = dist / torch.norm(temp_point, dim=1)
    return z
//  [python]
def distance_to_depth(K, dist, uv=None):
    if uv is None and len(dist.shape) >= 2:
        # create mesh grid according to d
        uv = np.stack(np.meshgrid(np.arange(dist.shape[1]), np.arange(dist.shape[0])), -1)
        uv = uv.reshape(-1, 2)
        dist = dist.reshape(-1)
        if not isinstance(dist, np.ndarray):
            uv = torch.from_numpy(uv).to(dist)
    if isinstance(dist, np.ndarray):
        # z * np.sqrt(x_temp**2+y_temp**2+z_temp**2) = dist
        uvh = np.concatenate([uv, np.ones((len(uv), 1))], -1)
        temp_point = dot(np.linalg.inv(K), uvh)
        z = dist / np.linalg.norm(temp_point, axis=1)

    else:
        uvh = torch.cat([uv, torch.ones(len(uv), 1).to(uv)], -1)
        temp_point = dot(torch.inverse(K), uvh)
        z = dist / torch.linalg.norm(temp_point, dim=1)
    return z
//  [python]
def aggregation_by_graph(cfg, graph_matrix, nets_this_round, global_w):
    tmp_client_state_dict = {}
    cluster_model_vectors = {}
    for client_id in nets_this_round.keys():
        tmp_client_state_dict[client_id] = copy.deepcopy(global_w)
        cluster_model_vectors[client_id] = torch.zeros_like(weight_flatten_all(global_w))
        for key in tmp_client_state_dict[client_id]:
            tmp_client_state_dict[client_id][key] = torch.zeros_like(tmp_client_state_dict[client_id][key])

    for client_id in nets_this_round.keys():
        tmp_client_state = tmp_client_state_dict[client_id]
        cluster_model_state = cluster_model_vectors[client_id]
        aggregation_weight_vector = graph_matrix[client_id]

        # if client_id==0:
        #     print(f'Aggregation weight: {aggregation_weight_vector}. Summation: {aggregation_weight_vector.sum()}')
        
        for neighbor_id in nets_this_round.keys():
            net_para = nets_this_round[neighbor_id].state_dict()
            for key in tmp_client_state:
                tmp_client_state[key] += net_para[key] * aggregation_weight_vector[neighbor_id]

        for neighbor_id in nets_this_round.keys():
            net_para = weight_flatten_all(nets_this_round[neighbor_id].state_dict())
            cluster_model_state += net_para * (aggregation_weight_vector[neighbor_id] / torch.norm(net_para))
               
    for client_id in nets_this_round.keys():
        nets_this_round[client_id].load_state_dict(tmp_client_state_dict[client_id])
    
    return cluster_model_vectors
//  [python]
def aggregation_by_graph(cfg, graph_matrix, nets_this_round, global_w):
    tmp_client_state_dict = {}
    cluster_model_vectors = {}
    for client_id in nets_this_round.keys():
        tmp_client_state_dict[client_id] = copy.deepcopy(global_w)
        cluster_model_vectors[client_id] = torch.zeros_like(weight_flatten_all(global_w))
        for key in tmp_client_state_dict[client_id]:
            tmp_client_state_dict[client_id][key] = torch.zeros_like(tmp_client_state_dict[client_id][key])

    for client_id in nets_this_round.keys():
        tmp_client_state = tmp_client_state_dict[client_id]
        cluster_model_state = cluster_model_vectors[client_id]
        aggregation_weight_vector = graph_matrix[client_id]

        # if client_id==0:
        #     print(f'Aggregation weight: {aggregation_weight_vector}. Summation: {aggregation_weight_vector.sum()}')
        
        for neighbor_id in nets_this_round.keys():
            net_para = nets_this_round[neighbor_id].state_dict()
            for key in tmp_client_state:
                tmp_client_state[key] += net_para[key] * aggregation_weight_vector[neighbor_id]

        for neighbor_id in nets_this_round.keys():
            net_para = weight_flatten_all(nets_this_round[neighbor_id].state_dict())
            cluster_model_state += net_para * (aggregation_weight_vector[neighbor_id] / torch.linalg.norm(net_para))
               
    for client_id in nets_this_round.keys():
        nets_this_round[client_id].load_state_dict(tmp_client_state_dict[client_id])
    
    return cluster_model_vectors
//  [python]
def rotation_matrix(a, b):
    """Compute the rotation matrix that rotates vector a to vector b.
    Args:
        a: The vector to rotate.
        b: The vector to rotate to.
    Returns:
        The rotation matrix.
    borrowed from from nerfstudio
    """
    a = a / t.norm(a)
    b = b / t.norm(b)
    v = t.cross(a, b)
    c = t.dot(a, b)
    # If vectors are exactly opposite, we add a little noise to one of them
    if c < -1 + 1e-8:
        eps = (t.rand(3) - 0.5) * 0.01
        return rotation_matrix(a + eps, b)
    s = t.norm(v)
    skew_sym_mat = t.Tensor(
        [
            [0, -v[2], v[1]],
            [v[2], 0, -v[0]],
            [-v[1], v[0], 0],
        ]
    )
    return t.eye(3) + skew_sym_mat + skew_sym_mat @ skew_sym_mat * ((1 - c) / (s ** 2 + 1e-8))
//  [python]
def rotation_matrix(a, b):
    """Compute the rotation matrix that rotates vector a to vector b.
    Args:
        a: The vector to rotate.
        b: The vector to rotate to.
    Returns:
        The rotation matrix.
    borrowed from from nerfstudio
    """
    a = a / torch.linalg.norm(a)
    b = b / torch.linalg.norm(b)
    v = torch.cross(a, b)
    c = torch.dot(a, b)
    # If vectors are exactly opposite, we add a little noise to one of them
    if c < -1 + 1e-8:
        eps = (torch.rand(3) - 0.5) * 0.01
        return rotation_matrix(a + eps, b)
    s = torch.linalg.norm(v)
    skew_sym_mat = torch.Tensor(
        [
            [0, -v[2], v[1]],
            [v[2], 0, -v[0]],
            [-v[1], v[0], 0],
        ]
    )
    return torch.eye(3) + skew_sym_mat + skew_sym_mat @ skew_sym_mat * ((1 - c) / (s ** 2 + 1e-8))
//  [python]
def forward(self, x0, x1, y):
        #方向,点乘小丸子
        #欧式，最后的分类的
        self.idx+=1
        x0_soft = F.softmax(x0,1)
        x1_soft = F.softmax(x1,1)
        # cosdis = torch.sum(x0_soft*x0_soft) + torch.sum(x1_soft*x1_soft) - 2* torch.mm(x0_soft, torch.transpose(x1_soft,0,1))
        cosdis = torch.matmul(x0_soft, x1_soft.T)
        coslabel = torch.diag(cosdis)
        coslabel = torch.div(coslabel, torch.norm(x0_soft, 1) * torch.norm(x1_soft, 1))

        if self.idx==60:
            coslabel = torch.div(coslabel, torch.norm(x0_soft, 1) * torch.norm(x1_soft, 1))
        e_diff = x0 - x1
        e_dist_sq = torch.sum(torch.pow(e_diff, 2), 1)
        e_dist = torch.sqrt(e_dist_sq)
        e_mdist = self.margin - e_dist
        e_dist = torch.clamp(e_mdist, min=0.0)
        # e_loss = y*coslabel*e_dist_sq + (1-y)*coslabel* torch.pow(e_dist, 2)
        e_loss = 2*y*abs(y-0.5)*e_dist_sq + 2*(1 - y)* abs(y-0.5) * torch.pow(e_dist, 2) + 400*y*(1-y)*e_dist_sq* coslabel
        # e_loss = y*e_dist_sq + (1 - y)*torch.pow(e_dist, 2)

        loss_1 = torch.sum(e_loss) / 2.0 / x0.size()[0]

        return loss_1
//  [python]
def forward(self, x0, x1, y):
        #方向,点乘小丸子
        #欧式，最后的分类的
        self.idx+=1
        x0_soft = F.softmax(x0,1)
        x1_soft = F.softmax(x1,1)
        # cosdis = torch.sum(x0_soft*x0_soft) + torch.sum(x1_soft*x1_soft) - 2* torch.mm(x0_soft, torch.transpose(x1_soft,0,1))
        cosdis = torch.matmul(x0_soft, x1_soft.T)
        coslabel = torch.diag(cosdis)
        coslabel = torch.div(coslabel, torch.linalg.norm(x0_soft, 1) * torch.linalg.norm(x1_soft, 1))

        if self.idx==60:
            coslabel = torch.div(coslabel, torch.linalg.norm(x0_soft, 1) * torch.linalg.norm(x1_soft, 1))
        e_diff = x0 - x1
        e_dist_sq = torch.sum(torch.pow(e_diff, 2), 1)
        e_dist = torch.sqrt(e_dist_sq)
        e_mdist = self.margin - e_dist
        e_dist = torch.clamp(e_mdist, min=0.0)
        # e_loss = y*coslabel*e_dist_sq + (1-y)*coslabel* torch.pow(e_dist, 2)
        e_loss = 2*y*abs(y-0.5)*e_dist_sq + 2*(1 - y)* abs(y-0.5) * torch.pow(e_dist, 2) + 400*y*(1-y)*e_dist_sq* coslabel
        # e_loss = y*e_dist_sq + (1 - y)*torch.pow(e_dist, 2)

        loss_1 = torch.sum(e_loss) / 2.0 / x0.size()[0]

        return loss_1
//  [python]
def normalize(x: torch.Tensor) -> TensorType[...]:
    """Returns a normalized vector."""
    return x / torch.norm(x)
//  [python]
def normalize(x: torch.Tensor) -> TensorType[...]:
    """Returns a normalized vector."""
    return x / torch.linalg.norm(x)
//  [python]
def forward(self, flow, _):
        return torch.mean(torch.norm(self.L(flow))**2)
//  [python]
def forward(self, flow, _):
        return torch.mean(torch.linalg.norm(self.L(flow))**2)
//  [python]
def pixel_norm(x, eps=1e-6):
    """Pixel Normalization.

    This normalization is proposed in:
    Progressive Growing of GANs for Improved Quality, Stability, and Variation

    Args:
        x (torch.Tensor): Tensor to be normalized.
        eps (float, optional): Epsilon to avoid dividing zero.
            Defaults to 1e-6.

    Returns:
        torch.Tensor: Normalized tensor.
    """
    norm = torch.norm(x, p=2, dim=1, keepdim=True)
    norm = norm / torch.sqrt(torch.tensor(x.shape[1]).to(x))

    return x / (norm + eps)
//  [python]
def pixel_norm(x, eps=1e-6):
    """Pixel Normalization.

    This normalization is proposed in:
    Progressive Growing of GANs for Improved Quality, Stability, and Variation

    Args:
        x (torch.Tensor): Tensor to be normalized.
        eps (float, optional): Epsilon to avoid dividing zero.
            Defaults to 1e-6.

    Returns:
        torch.Tensor: Normalized tensor.
    """
    if version.parse(torch.__version__) >= version.parse('1.7.0'):
        norm = torch.linalg.norm(x, ord=2, dim=1, keepdim=True)
    # support older pytorch version
    else:
        norm = torch.norm(x, p=2, dim=1, keepdim=True)
    norm = norm / torch.sqrt(torch.tensor(x.shape[1]).to(x))

    return x / (norm + eps)