jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def add_batch(self, inp, out):
        if len(inp.shape) == 2:
            inp = inp.unsqueeze(0)
        tmp = inp.shape[0]
        if isinstance(self.layer, nn.Linear) or isinstance(self.layer, transformers.Conv1D):
            if len(inp.shape) == 3:
                inp = inp.reshape((-1, inp.shape[-1]))
            inp = inp.t()
        self.H *= self.nsamples / (self.nsamples + tmp)

        self.scaler_row *= self.nsamples / (self.nsamples+tmp)

        self.nsamples += tmp
        inp = math.sqrt(2 / self.nsamples) * inp.float()
        self.H += inp.matmul(inp.t())
        self.scaler_row += torch.norm(inp, p=2, dim=1) ** 2  / self.nsamples
//  [python]
def add_batch(self, inp, out):
        if len(inp.shape) == 2:
            inp = inp.unsqueeze(0)
        tmp = inp.shape[0]
        if isinstance(self.layer, nn.Linear) or isinstance(self.layer, transformers.Conv1D):
            if len(inp.shape) == 3:
                inp = inp.reshape((-1, inp.shape[-1]))
            inp = inp.t()
        self.H *= self.nsamples / (self.nsamples + tmp)

        self.scaler_row *= self.nsamples / (self.nsamples+tmp)

        self.nsamples += tmp
        inp = math.sqrt(2 / self.nsamples) * inp.float()
        self.H += inp.matmul(inp.t())
        self.scaler_row += torch.linalg.norm(inp, p=2, dim=1) ** 2  / self.nsamples
//  [python]
def build_loss(self, recons, weights):
        size = self.X.shape[0]
        loss = torch.norm(recons - weights, p='fro')**2 / size
        return loss
//  [python]
def build_loss(self, recons, weights):
        size = self.X.shape[0]
        loss = torch.linalg.norm(recons - weights, ord='fro')**2 / size
        return loss
//  [python]
def magnitude_error(
    prediction: Dict[str, torch.Tensor],
    target: Dict[str, torch.Tensor],
    key=slice(None),
    p: int = 2,
) -> Dict[str, Union[float, int]]:
    assert prediction[key].shape[1] > 1
    error = torch.abs(
        torch.norm(prediction[key], p=p, dim=-1)
        - torch.norm(target[key], p=p, dim=-1)
    )
    return {
        "metric": torch.mean(error).item(),
        "total": torch.sum(error).item(),
        "numel": error.numel(),
    }
//  [python]
def magnitude_error(
    prediction: Dict[str, torch.Tensor],
    target: Dict[str, torch.Tensor],
    key=slice(None),
    p: int = 2,
) -> Dict[str, Union[float, int]]:
    assert prediction[key].shape[1] > 1
    error = torch.abs(
        torch.linalg.norm(prediction[key], ord=p, dim=-1)
        - torch.linalg.norm(target[key], ord=p, dim=-1)
    )
    return {
        "metric": torch.mean(error).item(),
        "total": torch.sum(error).item(),
        "numel": error.numel(),
    }
//  [python]
def _construct_quadratic() -> (
        Tuple[nn.Module, nn.Module, torch.Tensor, torch.Tensor]
    ):
        data = torch.arange(10, dtype=torch.float, device=torch.device("cuda"))
        data /= torch.norm(data)
        model = nn.Sequential(
            nn.Linear(10, 1, bias=False),
            nn.Linear(1, 1, bias=False),
        ).to(device=torch.device("cuda"))
        model[0].weight.data.fill_(1.0)
        model[1].weight.data.fill_(1.0)
        model[1].requires_grad = False
        loss = nn.MSELoss()
        target = torch.tensor([0.0]).to(device=torch.device("cuda"))
        return model, loss, data, target
//  [python]
def _construct_quadratic() -> (
        Tuple[nn.Module, nn.Module, torch.Tensor, torch.Tensor]
    ):
        data = torch.arange(10, dtype=torch.float, device=torch.device("cuda"))
        data /= torch.linalg.norm(data)
        model = nn.Sequential(
            nn.Linear(10, 1, bias=False),
            nn.Linear(1, 1, bias=False),
        ).to(device=torch.device("cuda"))
        model[0].weight.data.fill_(1.0)
        model[1].weight.data.fill_(1.0)
        model[1].requires_grad = False
        loss = nn.MSELoss()
        target = torch.tensor([0.0]).to(device=torch.device("cuda"))
        return model, loss, data, target
//  [python]
def forward_backward_consistency_check(fwd_flow, bwd_flow, alpha=0.01, beta=0.5):
    # fwd_flow, bwd_flow: [B, 2, H, W]
    # alpha and beta values are following UnFlow
    # (https://arxiv.org/abs/1711.07837)
    assert fwd_flow.dim() == 4 and bwd_flow.dim() == 4
    assert fwd_flow.size(1) == 2 and bwd_flow.size(1) == 2
    flow_mag = torch.norm(fwd_flow, dim=1) + torch.norm(bwd_flow, dim=1)  # [B, H, W]

    warped_bwd_flow = flow_warp(bwd_flow, fwd_flow)  # [B, 2, H, W]
    warped_fwd_flow = flow_warp(fwd_flow, bwd_flow)  # [B, 2, H, W]

    diff_fwd = torch.norm(fwd_flow + warped_bwd_flow, dim=1)  # [B, H, W]
    diff_bwd = torch.norm(bwd_flow + warped_fwd_flow, dim=1)

    threshold = alpha * flow_mag + beta

    fwd_occ = (diff_fwd > threshold).float()  # [B, H, W]
    bwd_occ = (diff_bwd > threshold).float()

    return fwd_occ, bwd_occ
//  [python]
def forward_backward_consistency_check(fwd_flow, bwd_flow, alpha=0.01, beta=0.5):
    # fwd_flow, bwd_flow: [B, 2, H, W]
    # alpha and beta values are following UnFlow
    # (https://arxiv.org/abs/1711.07837)
    assert fwd_flow.dim() == 4 and bwd_flow.dim() == 4
    assert fwd_flow.size(1) == 2 and bwd_flow.size(1) == 2
    flow_mag = torch.linalg.norm(fwd_flow, dim=1) + torch.linalg.norm(bwd_flow, dim=1)  # [B, H, W]

    warped_bwd_flow = flow_warp(bwd_flow, fwd_flow)  # [B, 2, H, W]
    warped_fwd_flow = flow_warp(fwd_flow, bwd_flow)  # [B, 2, H, W]

    diff_fwd = torch.linalg.norm(fwd_flow + warped_bwd_flow, dim=1)  # [B, H, W]
    diff_bwd = torch.linalg.norm(bwd_flow + warped_fwd_flow, dim=1)

    threshold = alpha * flow_mag + beta

    fwd_occ = (diff_fwd > threshold).float()  # [B, H, W]
    bwd_occ = (diff_bwd > threshold).float()

    return fwd_occ, bwd_occ
//  [python]
def inn_pred(W, X, A, B, phi, mitr=300, tol=3e-6, trasposed_A=False, compute_dphi=False):
        # TODO: randomized speed up
        At = A if trasposed_A else torch.transpose(A, 0, 1)
        #X = B if X is None else X

        err = 0
        status = 'max itrs reached'
        for i in range(mitr):
            # WXA
            X_ = W @ X
            support = torch.spmm(At, X_.T).T
            X_new = phi(support + B)
            err = torch.norm(X_new - X, np.inf)
            if err < tol:
                status = 'converged'
                break
            X = X_new

        dphi = None
        if compute_dphi:
            with torch.enable_grad():
                support = torch.spmm(At, (W @ X).T).T
                Z = support + B
                Z.requires_grad_(True)
                X_new = phi(Z)
                dphi = torch.autograd.grad(torch.sum(X_new), Z, only_inputs=True)[0]

        return X_new, err, status, dphi
//  [python]
def inn_pred(W, X, A, B, phi, mitr=300, tol=3e-6, trasposed_A=False, compute_dphi=False):
        At = A if trasposed_A else torch.transpose(A, 0, 1)
        #X = B if X is None else X

        err = 0
        status = 'max itrs reached'
        for i in range(mitr):
            # WXA
            X_ = W @ X
            support = torch.spmm(At, X_.T).T
            X_new = phi(support + B)
            err = torch.linalg.norm(X_new - X, np.inf)
            if err < tol:
                status = 'converged'
                break
            X = X_new

        dphi = None
        if compute_dphi:
            with torch.enable_grad():
                support = torch.spmm(At, (W @ X).T).T
                Z = support + B
                Z.requires_grad_(True)
                X_new = phi(Z)
                dphi = torch.autograd.grad(torch.sum(X_new), Z, only_inputs=True)[0]

        return X_new, err, status, dphi
//  [python]
def get_kmeans_mu(self, x, n_centers, init_times=50, min_delta=1e-3):
        """
        Find an initial value for the mean. Requires a threshold min_delta for the k-means algorithm to stop iterating.
        The algorithm is repeated init_times often, after which the best centerpoint is returned.
        args:
            x:            torch.FloatTensor (n, d) or (n, 1, d)
            init_times:   init
            min_delta:    int
        """
        if len(x.size()) == 3:
            x = x.squeeze(1)
        x_min, x_max = x.min(), x.max()
        x = (x - x_min) / (x_max - x_min)

        min_cost = np.inf

        for i in range(init_times):
            tmp_center = x[np.random.choice(np.arange(x.shape[0]), size=n_centers, replace=True), ...]
            l2_dis = torch.norm((x.unsqueeze(1).repeat(1, n_centers, 1) - tmp_center), p=2, dim=2)
            l2_cls = torch.argmin(l2_dis, dim=1)

            cost = 0
            for c in range(n_centers):
                if not (l2_cls == c).any():
                    continue
                cost += torch.norm(x[l2_cls == c] - tmp_center[c], p=2, dim=1).mean()

                assert not torch.isnan(cost)

            if cost < min_cost:
                min_cost = cost
                center = tmp_center

        delta = np.inf

        while delta > min_delta:
            l2_dis = torch.norm((x.unsqueeze(1).repeat(1, n_centers, 1) - center), p=2, dim=2)
            l2_cls = torch.argmin(l2_dis, dim=1)
            center_old = center.clone()

            for c in range(n_centers):
                if not (l2_cls == c).any():
                    center[c] = center_old[c]
                    continue
                center[c] = x[l2_cls == c].mean(dim=0)

            delta = torch.norm((center_old - center), dim=1).max()

        return (center.unsqueeze(0) * (x_max - x_min) + x_min)
//  [python]
def get_kmeans_mu(self, x, n_centers, init_times=50, min_delta=1e-3):
        """
        Find an initial value for the mean. Requires a threshold min_delta for the k-means algorithm to stop iterating.
        The algorithm is repeated init_times often, after which the best centerpoint is returned.
        args:
            x:            torch.FloatTensor (n, d) or (n, 1, d)
            init_times:   init
            min_delta:    int
        """
        if len(x.size()) == 3:
            x = x.squeeze(1)
        x_min, x_max = x.min(), x.max()
        x = (x - x_min) / (x_max - x_min)

        min_cost = np.inf

        for i in range(init_times):
            tmp_center = x[np.random.choice(np.arange(x.shape[0]), size=n_centers, replace=True), ...]
            l2_dis = torch.linalg.norm((x.unsqueeze(1).repeat(1, n_centers, 1) - tmp_center), ord=2, dim=2)
            l2_cls = torch.argmin(l2_dis, dim=1)

            cost = 0
            for c in range(n_centers):
                if not (l2_cls == c).any():
                    continue
                cost += torch.linalg.norm(x[l2_cls == c] - tmp_center[c], ord=2, dim=1).mean()

                assert not torch.isnan(cost)

            if cost < min_cost:
                min_cost = cost
                center = tmp_center

        delta = np.inf

        while delta > min_delta:
            l2_dis = torch.linalg.norm((x.unsqueeze(1).repeat(1, n_centers, 1) - center), ord=2, dim=2)
            l2_cls = torch.argmin(l2_dis, dim=1)
            center_old = center.clone()

            for c in range(n_centers):
                if not (l2_cls == c).any():
                    center[c] = center_old[c]
                    continue
                center[c] = x[l2_cls == c].mean(dim=0)

            delta = torch.linalg.norm((center_old - center), ord=2, dim=1).max()

        return (center.unsqueeze(0) * (x_max - x_min) + x_min)
//  [python]
def forward(self, rotation: torch.Tensor, position: torch.Tensor, offset: torch.Tensor, order='xyz', quater=False,
                world=True):
        '''
        if not quater and rotation.shape[-2] != 3: raise Exception('Unexpected shape of rotation')
        if quater and rotation.shape[-2] != 4: raise Exception('Unexpected shape of rotation')
        rotation = rotation.permute(0, 3, 1, 2)
        position = position.permute(0, 2, 1)
        '''
        result = torch.empty(rotation.shape[:-1] + (3,), device=position.device)

        norm = torch.norm(rotation, dim=-1, keepdim=True)
        rotation = rotation / norm

        # if quater:
        transform = self.transform_from_quaternion(rotation)
        # else:
        #     transform = self.transform_from_euler(rotation, order)

        offset = offset.reshape((-1, 1, offset.shape[-2], offset.shape[-1], 1))

        result[..., 0, :] = position
        for i, pi in enumerate(self.parents):
            if pi == -1:
                assert i == 0
                continue

            result[..., i, :] = torch.matmul(transform[..., pi, :, :], offset[..., i, :, :]).squeeze()
            transform[..., i, :, :] = torch.matmul(transform[..., pi, :, :].clone(), transform[..., i, :, :].clone())
            if world: result[..., i, :] += result[..., pi, :]
        return result
//  [python]
def forward(self, rotation: torch.Tensor, position: torch.Tensor, offset: torch.Tensor, order='xyz', quater=False,
                world=True):
        '''
        if not quater and rotation.shape[-2] != 3: raise Exception('Unexpected shape of rotation')
        if quater and rotation.shape[-2] != 4: raise Exception('Unexpected shape of rotation')
        rotation = rotation.permute(0, 3, 1, 2)
        position = position.permute(0, 2, 1)
        '''
        result = torch.empty(rotation.shape[:-1] + (3,), device=position.device)

        norm = torch.linalg.norm(rotation, dim=-1, keepdim=True)
        rotation = rotation / norm

        # if quater:
        transform = self.transform_from_quaternion(rotation)
        # else:
        #     transform = self.transform_from_euler(rotation, order)

        offset = offset.reshape((-1, 1, offset.shape[-2], offset.shape[-1], 1))

        result[..., 0, :] = position
        for i, pi in enumerate(self.parents):
            if pi == -1:
                assert i == 0
                continue

            result[..., i, :] = torch.matmul(transform[..., pi, :, :], offset[..., i, :, :]).squeeze()
            transform[..., i, :, :] = torch.matmul(transform[..., pi, :, :].clone(), transform[..., i, :, :].clone())
            if world: result[..., i, :] += result[..., pi, :]
        return result
//  [python]
def test_norm() -> None:
    x: Tensor[torch.float32, L[2], L[3], L[4]]
    x_out: Tensor[torch.float32, L[2], L[3], L[4]]

    y1: Tensor[torch.float32] = torch.norm(x)
    y2: Tensor[torch.float32, L[3], L[4]] = torch.norm(x, dim=0, out=x_out, p=1)
    # pyre-fixme[9]: Expected error.
    y2_error: Tensor[torch.float32, L[3], L[99]] = torch.norm(x, dim=0)
    y3: Tensor[torch.float32, L[1], L[3], L[4]] = torch.norm(x, dim=0, keepdim=True)
//  [python]
def test_norm() -> None:
    x: Tensor[torch.float32, L[2], L[3], L[4]]
    x_out: Tensor[torch.float32, L[2], L[3], L[4]]

    y1: Tensor[torch.float32] = torch.linalg.norm(x)
    y2: Tensor[torch.float32, L[3], L[4]] = torch.linalg.norm(x, dim=0, out=x_out, ord=1)
    # pyre-fixme[9]: Expected error.
    y2_error: Tensor[torch.float32, L[3], L[99]] = torch.linalg.norm(x, dim=0)
    y3: Tensor[torch.float32, L[1], L[3], L[4]] = torch.linalg.norm(x, dim=0, keepdim=True)
//  [python]
def corner_loss(self, pred_bbox3d, gt_bbox3d, delta=1):
        """Calculate corner loss of given boxes.

        Args:
            pred_bbox3d (torch.FloatTensor): Predicted boxes in shape (N, 7).
            gt_bbox3d (torch.FloatTensor): Ground truth boxes in shape (N, 7).

        Returns:
            torch.FloatTensor: Calculated corner loss in shape (N).
        """
        assert pred_bbox3d.shape[0] == gt_bbox3d.shape[0]

        gt_boxes_structure = LiDARInstance3DBoxes(gt_bbox3d)
        pred_box_corners = LiDARInstance3DBoxes(pred_bbox3d).corners
        gt_box_corners = gt_boxes_structure.corners

        # This flip only changes the heading direction of GT boxes
        gt_bbox3d_flip = gt_boxes_structure.clone()
        gt_bbox3d_flip.tensor[:, 6] += np.pi
        gt_box_corners_flip = gt_bbox3d_flip.corners

        corner_dist = torch.min(
            torch.norm(pred_box_corners - gt_box_corners, dim=2),
            torch.norm(pred_box_corners - gt_box_corners_flip,
                       dim=2))  # (N, 8)
        # huber loss
        abs_error = torch.abs(corner_dist)
        quadratic = torch.clamp(abs_error, max=delta)
        linear = (abs_error - quadratic)
        corner_loss = 0.5 * quadratic**2 + delta * linear

        return corner_loss.mean(1)
//  [python]
def corner_loss(self, pred_bbox3d, gt_bbox3d, delta=1):
        """Calculate corner loss of given boxes.

        Args:
            pred_bbox3d (torch.FloatTensor): Predicted boxes in shape (N, 7).
            gt_bbox3d (torch.FloatTensor): Ground truth boxes in shape (N, 7).

        Returns:
            torch.FloatTensor: Calculated corner loss in shape (N).
        """
        assert pred_bbox3d.shape[0] == gt_bbox3d.shape[0]

        gt_boxes_structure = LiDARInstance3DBoxes(gt_bbox3d)
        pred_box_corners = LiDARInstance3DBoxes(pred_bbox3d).corners
        gt_box_corners = gt_boxes_structure.corners

        # This flip only changes the heading direction of GT boxes
        gt_bbox3d_flip = gt_boxes_structure.clone()
        gt_bbox3d_flip.tensor[:, 6] += np.pi
        gt_box_corners_flip = gt_bbox3d_flip.corners

        corner_dist = torch.min(
            torch.linalg.norm(pred_box_corners - gt_box_corners, dim=2),
            torch.linalg.norm(pred_box_corners - gt_box_corners_flip,
                       dim=2))  # (N, 8)
        # huber loss
        abs_error = torch.abs(corner_dist)
        quadratic = torch.clamp(abs_error, max=delta)
        linear = (abs_error - quadratic)
        corner_loss = 0.5 * quadratic**2 + delta * linear

        return corner_loss.mean(1)
//  [python]
def normalize_Trans(new_cameras):
    
    t_gt = new_cameras.T.clone()
    t_gt = t_gt[1:,:]
    t_gt_scale = torch.norm(t_gt, dim=(0, 1))
    t_gt_scale = t_gt_scale / math.sqrt(len(t_gt))
    t_gt_scale = t_gt_scale / 2
    t_gt_scale = t_gt_scale.clamp(min=0.01, max = 100)
    new_cameras.T = new_cameras.T / t_gt_scale
    
    return new_cameras
//  [python]
def normalize_Trans(new_cameras):
    
    t_gt = new_cameras.T.clone()
    t_gt = t_gt[1:,:]
    t_gt_scale = torch.linalg.norm(t_gt, dim=(0, 1))
    t_gt_scale = t_gt_scale / math.sqrt(len(t_gt))
    t_gt_scale = t_gt_scale / 2
    t_gt_scale = t_gt_scale.clamp(min=0.01, max = 100)
    new_cameras.T = new_cameras.T / t_gt_scale
    
    return new_cameras
//  [python]
def forward(self, embeddings, x, y):
        B, _, N, _ = x.size()
        src_keypts = x.squeeze(dim=1)[:, :, 0:2]
        tgt_keypts = x.squeeze(dim=1)[:, :, 2:4]
        with torch.no_grad():
            src_dist = torch.norm((src_keypts[:, :, None, :] - src_keypts[:, None, :, :]), dim=-1)
            len_sim = src_dist - torch.norm((tgt_keypts[:, :, None, :] - tgt_keypts[:, None, :, :]), dim=-1)

        out = get_graph_feature(embeddings, k=self.k_num)
        out = self.gab(out)
        out = self.aggregator(out)
        out = self.encoder(out)
        w0 = self.linear_0(out).view(B, -1)

        out_g = self.transformer(out.transpose(1, 2).squeeze(dim=-1), out.transpose(1, 2).squeeze(dim=-1), len_sim)[0]
        out_g = out_g.unsqueeze(dim=-1).transpose(1, 2)
        out = torch.cat([out, out_g], dim=1)
        out = self.resnet_block(out)
        w1 = self.linear_1(out).view(B, -1)

        if self.predict == False:
            w1_ds, indices = torch.sort(w1, dim=-1, descending=True)
            w1_ds = w1_ds[:, :int(N*self.sr)]
            x_ds, y_ds, w0_ds = self.down_sampling(x, y, w0, indices, None, self.predict)
            return x_ds, y_ds, [w0, w1], [w0_ds, w1_ds]
        else:
            w1_ds, indices = torch.sort(w1, dim=-1, descending=True)
            w1_ds = w1_ds[:, :int(N*self.sr)]
            x_ds, y_ds, w0_ds, out = self.down_sampling(x, y, w0, indices, out, self.predict)
            w2 = self.inlier_predictor(out)
            e_hat = weighted_8points(x_ds, w2)
            return x_ds, y_ds, [w0, w1, w2[:, 0, :, 0]], [w0_ds, w1_ds], e_hat
//  [python]
def forward(self, embeddings, x, y):
        B, _, N, _ = x.size()
        src_keypts = x.squeeze(dim=1)[:, :, 0:2]
        tgt_keypts = x.squeeze(dim=1)[:, :, 2:4]
        with torch.no_grad():
            src_dist = torch.linalg.norm((src_keypts[:, :, None, :] - src_keypts[:, None, :, :]), dim=-1)
            len_sim = src_dist - torch.linalg.norm((tgt_keypts[:, :, None, :] - tgt_keypts[:, None, :, :]), dim=-1)

        out = get_graph_feature(embeddings, k=self.k_num)
        out = self.gab(out)
        out = self.aggregator(out)
        out = self.encoder(out)
        w0 = self.linear_0(out).view(B, -1)

        out_g = self.transformer(out.transpose(1, 2).squeeze(dim=-1), out.transpose(1, 2).squeeze(dim=-1), len_sim)[0]
        out_g = out_g.unsqueeze(dim=-1).transpose(1, 2)
        out = torch.cat([out, out_g], dim=1)
        out = self.resnet_block(out)
        w1 = self.linear_1(out).view(B, -1)

        if self.predict == False:
            w1_ds, indices = torch.sort(w1, dim=-1, descending=True)
            w1_ds = w1_ds[:, :int(N*self.sr)]
            x_ds, y_ds, w0_ds = self.down_sampling(x, y, w0, indices, None, self.predict)
            return x_ds, y_ds, [w0, w1], [w0_ds, w1_ds]
        else:
            w1_ds, indices = torch.sort(w1, dim=-1, descending=True)
            w1_ds = w1_ds[:, :int(N*self.sr)]
            x_ds, y_ds, w0_ds, out = self.down_sampling(x, y, w0, indices, out, self.predict)
            w2 = self.inlier_predictor(out)
            e_hat = weighted_8points(x_ds, w2)
            return x_ds, y_ds, [w0, w1, w2[:, 0, :, 0]], [w0_ds, w1_ds], e_hat
//  [python]
def fltrust(gradients, net, lr, f, byz, device):
    """
    Based on the description in https://arxiv.org/abs/2012.13995
    gradients: list of gradients. The last one is the server update.
    net: model parameters.
    lr: learning rate.
    f: number of malicious clients. The first f clients are malicious.
    byz: attack type.
    device: computation device.
    """
    param_list = [torch.cat([xx.reshape((-1, 1)) for xx in x], dim=0) for x in gradients]
    # let the malicious clients (first f clients) perform the byzantine attack
    param_list = byz(param_list, net, lr, f, device)
    n = len(param_list) - 1

    # use the last gradient (server update) as the trusted source
    baseline = param_list[-1].squeeze()
    sim = []
    new_param_list = []

    # compute similarity
    for each_param_list in param_list:
        each_param_array = each_param_list.squeeze()

        sim.append(torch.dot(baseline, each_param_array) / (torch.norm(baseline) + 1e-9) / (
                    torch.norm(each_param_array) + 1e-9))

    sim = torch.stack(sim)[:-1]

    # clip similarities and get trust scores
    sim = F.relu(sim)
    normalized_weights = sim / (torch.sum(sim).item() + 1e-9)

    # normalize the magnitudes and weight by the trust score
    for i in range(n):
        new_param_list.append(
            param_list[i] * normalized_weights[i] / (torch.norm(param_list[i]) + 1e-9) * torch.norm(baseline))

    # compute global update
    global_update = torch.sum(torch.cat(new_param_list, dim=1), dim=-1)

    # update global model
    idx = 0
    for j, (param) in enumerate(net.parameters()):
        param.add_(global_update[idx:(idx + torch.numel(param))].reshape(tuple(param.size())), alpha=-lr)
        idx += torch.numel(param)
//  [python]
def fltrust(gradients, net, lr, f, byz, device):
    """
    Based on the description in https://arxiv.org/abs/2012.13995
    gradients: list of gradients. The last one is the server update.
    net: model parameters.
    lr: learning rate.
    f: number of malicious clients. The first f clients are malicious.
    byz: attack type.
    device: computation device.
    """
    param_list = [torch.cat([xx.reshape((-1, 1)) for xx in x], dim=0) for x in gradients]
    # let the malicious clients (first f clients) perform the byzantine attack
    param_list = byz(param_list, net, lr, f, device)
    n = len(param_list) - 1

    # use the last gradient (server update) as the trusted source
    baseline = param_list[-1].squeeze()
    sim = []
    new_param_list = []

    # compute similarity
    for each_param_list in param_list:
        each_param_array = each_param_list.squeeze()

        sim.append(torch.dot(baseline, each_param_array) / (torch.linalg.norm(baseline) + 1e-9) / (
                    torch.linalg.norm(each_param_array) + 1e-9))

    sim = torch.stack(sim)[:-1]

    # clip similarities and get trust scores
    sim = F.relu(sim)
    normalized_weights = sim / (torch.sum(sim).item() + 1e-9)

    # normalize the magnitudes and weight by the trust score
    for i in range(n):
        new_param_list.append(
            param_list[i] * normalized_weights[i] / (torch.linalg.norm(param_list[i]) + 1e-9) * torch.linalg.norm(baseline))

    # compute global update
    global_update = torch.sum(torch.cat(new_param_list, dim=1), dim=-1)

    # update global model
    idx = 0
    for j, (param) in enumerate(net.parameters()):
        param.add_(global_update[idx:(idx + torch.numel(param))].reshape(tuple(param.size())), alpha=-lr)
        idx += torch.numel(param)
//  [python]
def krum(gradients, net, lr, f, byz, device):
    """
    Based on the description in https://papers.nips.cc/paper/2017/hash/f4b9ec30ad9f68f89b29639786cb62ef-Abstract.html
    gradients: list of gradients.
    net: model parameters.
    lr: learning rate.
    f: number of malicious clients. The first f clients are malicious.
    byz: attack type.
    device: computation device.
    """
    param_list = [torch.cat([xx.reshape((-1, 1)) for xx in x], dim=0) for x in gradients]
    # let the malicious clients (first f clients) perform the byzantine attack
    if byz == attacks.fltrust_attack:
        param_list = byz(param_list, net, lr, f, device)[:-1]
    else:
        param_list = byz(param_list, net, lr, f, device)
    n = len(param_list)

    # compute pairwise Euclidean distance
    dist = torch.zeros((n, n)).to(device)
    for i in range(n):
        for j in range(i + 1, n):
            d = torch.norm(param_list[i] - param_list[j])
            dist[i, j], dist[j, i] = d, d

    # sort distances and get model with smallest sum of distances to closest n-f-2 models
    sorted_dist, _ = torch.sort(dist, dim=-1)
    global_update = param_list[torch.argmin(torch.sum(sorted_dist[:, 0:(n - f - 1)], dim=-1))]

    # update global model
    idx = 0
    for j, param in enumerate(net.parameters()):
        param.add_(global_update[idx:(idx + torch.numel(param))].reshape(tuple(param.size())), alpha=-lr)
        idx += torch.numel(param)
//  [python]
def krum(gradients, net, lr, f, byz, device):
    """
    Based on the description in https://papers.nips.cc/paper/2017/hash/f4b9ec30ad9f68f89b29639786cb62ef-Abstract.html
    gradients: list of gradients.
    net: model parameters.
    lr: learning rate.
    f: number of malicious clients. The first f clients are malicious.
    byz: attack type.
    device: computation device.
    """
    param_list = [torch.cat([xx.reshape((-1, 1)) for xx in x], dim=0) for x in gradients]
    # let the malicious clients (first f clients) perform the byzantine attack
    if byz == attacks.fltrust_attack:
        param_list = byz(param_list, net, lr, f, device)[:-1]
    else:
        param_list = byz(param_list, net, lr, f, device)
    n = len(param_list)

    # compute pairwise Euclidean distance
    dist = torch.zeros((n, n)).to(device)
    for i in range(n):
        for j in range(i + 1, n):
            d = torch.linalg.norm(param_list[i] - param_list[j])
            dist[i, j], dist[j, i] = d, d

    # sort distances and get model with smallest sum of distances to closest n-f-2 models
    sorted_dist, _ = torch.sort(dist, dim=-1)
    global_update = param_list[torch.argmin(torch.sum(sorted_dist[:, 0:(n - f - 1)], dim=-1))]

    # update global model
    idx = 0
    for j, param in enumerate(net.parameters()):
        param.add_(global_update[idx:(idx + torch.numel(param))].reshape(tuple(param.size())), alpha=-lr)
        idx += torch.numel(param)
//  [python]
def normalize_cameras(cameras, compute_optical=True, first_camera=True, scale=1.0, normalize_T = False):
    """
    Normalizes cameras such that the optical axes point to the origin and the average
    distance to the origin is 1.

    Args:
        cameras (List[camera]).
    """
    # Let distance from first camera to origin be unit
    new_cameras = cameras.clone()

    if compute_optical:
        new_transform = new_cameras.get_world_to_view_transform()

        (p_intersect, dist, p_line_intersect, pp, r) = compute_optical_axis_intersection(cameras)
        t = Translate(p_intersect)

        scale = dist.squeeze()[0]

        # Degenerate case
        if scale == 0:
            scale = torch.norm(new_cameras.T, dim=(0, 1))
            scale = torch.sqrt(scale)
            new_cameras.T = new_cameras.T / scale
        else:
            new_matrix = t.compose(new_transform).get_matrix()
            new_cameras.R = new_matrix[:, :3, :3]
            new_cameras.T = new_matrix[:, 3, :3] / scale
    else:
        scale = torch.norm(new_cameras.T, dim=(0, 1))
        scale = torch.sqrt(scale)
        new_cameras.T = new_cameras.T / scale

    if first_camera:
        new_cameras = first_camera_transform(new_cameras)
    
    if normalize_T:
        new_cameras = normalize_Trans(new_cameras)

    return new_cameras
//  [python]
def normalize_cameras(cameras, compute_optical=True, first_camera=True, scale=1.0, normalize_T = False):
    """
    Normalizes cameras such that the optical axes point to the origin and the average
    distance to the origin is 1.

    Args:
        cameras (List[camera]).
    """
    # Let distance from first camera to origin be unit
    new_cameras = cameras.clone()

    if compute_optical:
        new_transform = new_cameras.get_world_to_view_transform()

        (p_intersect, dist, p_line_intersect, pp, r) = compute_optical_axis_intersection(cameras)
        t = Translate(p_intersect)

        scale = dist.squeeze()[0]

        # Degenerate case
        if scale == 0:
            scale = torch.linalg.norm(new_cameras.T, dim=(0, 1))
            scale = torch.sqrt(scale)
            new_cameras.T = new_cameras.T / scale
        else:
            new_matrix = t.compose(new_transform).get_matrix()
            new_cameras.R = new_matrix[:, :3, :3]
            new_cameras.T = new_matrix[:, 3, :3] / scale
    else:
        scale = torch.linalg.norm(new_cameras.T, dim=(0, 1))
        scale = torch.sqrt(scale)
        new_cameras.T = new_cameras.T / scale

    if first_camera:
        new_cameras = first_camera_transform(new_cameras)
    
    if normalize_T:
        new_cameras = normalize_Trans(new_cameras)

    return new_cameras
//  [python]
def get_random_spherical_points(batch_size, point_num):
    pcs = torch.randn((batch_size, 3, point_num), dtype=torch.float32)
    pcs = pcs/(torch.norm(pcs, dim=1, keepdim=True) + 1e-7)
    return pcs
//  [python]
def get_random_spherical_points(batch_size, point_num):
    pcs = torch.randn((batch_size, 3, point_num), dtype=torch.float32)
    pcs = pcs/(torch.linalg.norm(pcs, dim=1, keepdim=True) + 1e-7)
    return pcs
//  [python]
def get_avg_init_norm(layer, param_type=None, ord=2, repetitions=100):
    """Computes the average norm of default layer initialization"""
    output = 0
    for _ in range(repetitions):
        layer.reset_parameters()
        warnings.warn("torch.norm is deprecated. Think about updating this.")
        output += torch.norm(getattr(layer, param_type), p=ord).item()
    return float(output) / repetitions
//  [python]
def get_avg_init_norm(layer, param_type=None, ord=2, repetitions=100):
    """Computes the average norm of default layer initialization"""
    output = 0
    for _ in range(repetitions):
        layer.reset_parameters()
        output += torch.linalg.norm(getattr(layer, param_type), ord=ord).item()
    return float(output) / repetitions
//  [python]
def foolsgold(gradients, net, lr, f, byz, device, gradient_history):
    """
    Based on the description in https://arxiv.org/abs/1808.04866
    gradients: list of gradients.
    net: model parameters.
    lr: learning rate.
    f: number of malicious clients. The first f clients are malicious.
    byz: attack type.
    device: computation device.
    gradient_history: aggregation of previous gradients per worker.
    """

    # reference implementation: https://github.com/DistributedML/FoolsGold
    # FoolsGold has individual learning rates. Global learning rate lr has no effect

    kappa = 1
    eps = 10e-5

    param_list = [torch.cat([xx.reshape((-1, 1)) for xx in x], dim=0) for x in gradients]
    # let the malicious clients (first f clients) perform the byzantine attack
    if byz == attacks.fltrust_attack:
        param_list = byz(param_list, net, lr, f, device)[:-1]
    else:
        param_list = byz(param_list, net, lr, f, device)

    n = len(param_list)

    for i in range(n):
        norm = torch.norm(param_list[i])
        if(norm > 1):
            param_list[i] /= norm

    # updates history
    gradient_history = [gradient_history[i] + param_list[i] for i in range(n)]

    cos_dist = torch.zeros((n, n), dtype=torch.double).to(device)
    for i in range(n):
        for j in range(i + 1, n):
            d = F.cosine_similarity(gradient_history[i], gradient_history[j], dim=0, eps=1e-9)
            cos_dist[i, j], cos_dist[j, i] = d, d

    v, _ = torch.max(cos_dist, dim=1)

    # pardoning
    for i in range(n):
        for j in range(n):
            if v[j] > v[i]:
                cos_dist[i][j] *= v[i]/v[j]

    alpha = torch.clamp(1 - torch.max(cos_dist, dim=1)[0], min=0, max=1)

    # logit function
    alpha /= torch.max(alpha, dim=0, keepdim=True)[0]
    alpha = kappa * (torch.logit(alpha, eps=eps) + 0.5)
    alpha = torch.clamp(alpha, min=0, max=1)

    # calculate global update
    global_update = torch.zeros(param_list[0].size()).to(device)
    for i, grad in enumerate(param_list):
        global_update += grad * alpha[i]

    # update the global model
    idx = 0
    for j, param in enumerate(net.parameters()):
        param.add_(global_update[idx:(idx + torch.numel(param))].reshape(tuple(param.size())), alpha=-1)        # FoolsGold uses individual learning rates
        idx += torch.numel(param)

    return gradient_history
//  [python]
def foolsgold(gradients, net, lr, f, byz, device, gradient_history):
    """
    Based on the description in https://arxiv.org/abs/1808.04866
    gradients: list of gradients.
    net: model parameters.
    lr: learning rate.
    f: number of malicious clients. The first f clients are malicious.
    byz: attack type.
    device: computation device.
    gradient_history: aggregation of previous gradients per worker.
    """

    # reference implementation: https://github.com/DistributedML/FoolsGold
    # FoolsGold has individual learning rates. Global learning rate lr has no effect

    kappa = 1
    eps = 10e-5

    param_list = [torch.cat([xx.reshape((-1, 1)) for xx in x], dim=0) for x in gradients]
    # let the malicious clients (first f clients) perform the byzantine attack
    if byz == attacks.fltrust_attack:
        param_list = byz(param_list, net, lr, f, device)[:-1]
    else:
        param_list = byz(param_list, net, lr, f, device)

    n = len(param_list)

    for i in range(n):
        norm = torch.linalg.norm(param_list[i])
        if(norm > 1):
            param_list[i] /= norm

    # updates history
    gradient_history = [gradient_history[i] + param_list[i] for i in range(n)]

    cos_dist = torch.zeros((n, n), dtype=torch.double).to(device)
    for i in range(n):
        for j in range(i + 1, n):
            d = F.cosine_similarity(gradient_history[i], gradient_history[j], dim=0, eps=1e-9)
            cos_dist[i, j], cos_dist[j, i] = d, d

    v, _ = torch.max(cos_dist, dim=1)

    # pardoning
    for i in range(n):
        for j in range(n):
            if v[j] > v[i]:
                cos_dist[i][j] *= v[i]/v[j]

    alpha = torch.clamp(1 - torch.max(cos_dist, dim=1)[0], min=0, max=1)

    # logit function
    alpha /= torch.max(alpha, dim=0, keepdim=True)[0]
    alpha = kappa * (torch.logit(alpha, eps=eps) + 0.5)
    alpha = torch.clamp(alpha, min=0, max=1)

    # calculate global update
    global_update = torch.zeros(param_list[0].size()).to(device)
    for i, grad in enumerate(param_list):
        global_update += grad * alpha[i]

    # update the global model
    idx = 0
    for j, param in enumerate(net.parameters()):
        param.add_(global_update[idx:(idx + torch.numel(param))].reshape(tuple(param.size())), alpha=-1)        # FoolsGold uses individual learning rates
        idx += torch.numel(param)

    return gradient_history
//  [python]
def edge_init(self, edges):
        # Calculate angles k -> j -> i
        R1, R2 = edges.src["o"], edges.dst["o"]
        x = torch.sum(R1 * R2, dim=-1)
        y = torch.cross(R1, R2)
        y = torch.norm(y, dim=-1)
        angle = torch.atan2(y, x)
        # Transform via angles
        cbf = [f(angle) for f in self.sbf_layer.get_sph_funcs()]
        cbf = torch.stack(cbf, dim=1)  # [None, 7]
        cbf = cbf.repeat_interleave(self.num_radial, dim=1)  # [None, 42]
        sbf = edges.src["rbf_env"] * cbf  # [None, 42]
        return {"sbf": sbf}
//  [python]
def edge_init(self, edges):
        # Calculate angles k -> j -> i
        R1, R2 = edges.src["o"], edges.dst["o"]
        x = torch.sum(R1 * R2, dim=-1)
        y = torch.cross(R1, R2)
        y = torch.linalg.norm(y, dim=-1)
        angle = torch.atan2(y, x)
        # Transform via angles
        cbf = [f(angle) for f in self.sbf_layer.get_sph_funcs()]
        cbf = torch.stack(cbf, dim=1)  # [None, 7]
        cbf = cbf.repeat_interleave(self.num_radial, dim=1)  # [None, 42]
        sbf = edges.src["rbf_env"] * cbf  # [None, 42]
        return {"sbf": sbf}
//  [python]
def forward(ctx, input, emb):
        if input.size(1) != emb.size(0):
            raise RuntimeError('invalid argument: input.size(1) ({}) must be equal to emb.size(0) ({})'.
                               format(input.size(1), emb.size(0)))

        # save sizes for backward
        ctx.batch_size = input.size(0)
        ctx.num_latents = int(np.prod(np.array(input.size()[2:])))
        ctx.emb_dim = emb.size(0)
        ctx.num_emb = emb.size(1)
        ctx.input_type = type(input)
        ctx.dims = list(range(len(input.size())))

        # expand so it broadcastable
        x_expanded = input.unsqueeze(-1)
        num_arbitrary_dims = len(ctx.dims) - 2
        if num_arbitrary_dims:
            emb_expanded = emb.view(emb.shape[0], *([1] * num_arbitrary_dims), emb.shape[1])
        else:
            emb_expanded = emb

        # find nearest neighbors
        dist = torch.norm(x_expanded - emb_expanded, 2, 1)
        _, argmin = dist.min(-1)
        shifted_shape = [input.shape[0], *list(input.shape[2:]) ,input.shape[1]]
        result = emb.t().index_select(0, argmin.view(-1)).view(shifted_shape).permute(0, ctx.dims[-1], *ctx.dims[1:-1])

        ctx.argmin = argmin
        return result.contiguous(), argmin
//  [python]
def forward(ctx, input, emb):
        if input.size(1) != emb.size(0):
            raise RuntimeError('invalid argument: input.size(1) ({}) must be equal to emb.size(0) ({})'.
                               format(input.size(1), emb.size(0)))

        # save sizes for backward
        ctx.batch_size = input.size(0)
        ctx.num_latents = int(np.prod(np.array(input.size()[2:])))
        ctx.emb_dim = emb.size(0)
        ctx.num_emb = emb.size(1)
        ctx.input_type = type(input)
        ctx.dims = list(range(len(input.size())))

        # expand so it broadcastable
        x_expanded = input.unsqueeze(-1)
        num_arbitrary_dims = len(ctx.dims) - 2
        if num_arbitrary_dims:
            emb_expanded = emb.view(emb.shape[0], *([1] * num_arbitrary_dims), emb.shape[1])
        else:
            emb_expanded = emb

        # find nearest neighbors
        dist = torch.linalg.norm(x_expanded - emb_expanded, 2, 1)
        _, argmin = dist.min(-1)
        shifted_shape = [input.shape[0], *list(input.shape[2:]) ,input.shape[1]]
        result = emb.t().index_select(0, argmin.view(-1)).view(shifted_shape).permute(0, ctx.dims[-1], *ctx.dims[1:-1])

        ctx.argmin = argmin
        return result.contiguous(), argmin
//  [python]
def forward_affinity(id1, id2, mesh_dir=None):
    """
    Calculate affinity matrix according to id1 and id2.
    Save the affinity matrix to mesh_dir if it is not None.
    """
    valid_bs = 1
    affinity_pred = torch.zeros([max_rois, max_rois], dtype=torch.float32)
    if id1.shape[0] == 0 or id2.shape[0] == 0:
        affinity_pred = affinity_pred.numpy()
        if mesh_dir is not None:
            plt.figure()
            sns.heatmap(affinity_pred, vmin=0.0, vmax=1.0)
            plt.savefig(os.path.join(mesh_dir, 'affinity_pred.png'))
            plt.close()
        return affinity_pred
    affinity_pred = affinity_pred.cuda()
    objs_in_view1 = id1
    objs_in_view2 = id2
    
    norms = (torch.norm(objs_in_view1, p=2, dim=1) + 1e-3).detach()
    objs_in_view1 = objs_in_view1.div(norms.view(-1,1).expand_as(objs_in_view1))
    norms = (torch.norm(objs_in_view2, p=2, dim=1) + 1e-3).detach()
    objs_in_view2 = objs_in_view2.div(norms.view(-1,1).expand_as(objs_in_view2))
    prod = torch.mm(
        objs_in_view1, torch.t(objs_in_view2)
    ) * 5
    prod = torch.sigmoid(prod)
    sz_i = prod.shape[0]
    sz_j = prod.shape[1]
    #if sz_i > 20 or sz_j > 20:
    #shutil.rmtree(mesh_dir)
    #return None
    affinity_pred[:sz_i, :sz_j] = prod
        
    #match = affinity_pred.argmax(dim=1).cpu().numpy()
    affinity_pred = affinity_pred.cpu().detach().numpy()

    # visualization
    if mesh_dir is not None:
        max_sz = max(sz_i, sz_j)
        if max_sz < 5:
            max_sz = 5
        elif max_sz < 10:
            max_sz = 10
        affinity_vis = affinity_pred[:max_sz, :max_sz]
        plt.figure()
        sns.heatmap(affinity_vis, annot=True, vmin=0.0, vmax=1.0)
        plt.savefig(os.path.join(mesh_dir, 'affinity_pred.png'))
        plt.close()

    return affinity_pred
//  [python]
def forward_affinity(id1, id2, mesh_dir=None):
    """
    Calculate affinity matrix according to id1 and id2.
    Save the affinity matrix to mesh_dir if it is not None.
    """
    valid_bs = 1
    affinity_pred = torch.zeros([max_rois, max_rois], dtype=torch.float32)
    if id1.shape[0] == 0 or id2.shape[0] == 0:
        affinity_pred = affinity_pred.numpy()
        if mesh_dir is not None:
            plt.figure()
            sns.heatmap(affinity_pred, vmin=0.0, vmax=1.0)
            plt.savefig(os.path.join(mesh_dir, 'affinity_pred.png'))
            plt.close()
        return affinity_pred
    affinity_pred = affinity_pred.cuda()
    objs_in_view1 = id1
    objs_in_view2 = id2
    
    norms = (torch.linalg.norm(objs_in_view1, ord=2, dim=1) + 1e-3).detach()
    objs_in_view1 = objs_in_view1.div(norms.view(-1,1).expand_as(objs_in_view1))
    norms = (torch.linalg.norm(objs_in_view2, ord=2, dim=1) + 1e-3).detach()
    objs_in_view2 = objs_in_view2.div(norms.view(-1,1).expand_as(objs_in_view2))
    prod = torch.mm(
        objs_in_view1, torch.t(objs_in_view2)
    ) * 5
    prod = torch.sigmoid(prod)
    sz_i = prod.shape[0]
    sz_j = prod.shape[1]
    #if sz_i > 20 or sz_j > 20:
    #shutil.rmtree(mesh_dir)
    #return None
    affinity_pred[:sz_i, :sz_j] = prod
        
    #match = affinity_pred.argmax(dim=1).cpu().numpy()
    affinity_pred = affinity_pred.cpu().detach().numpy()

    # visualization
    if mesh_dir is not None:
        max_sz = max(sz_i, sz_j)
        if max_sz < 5:
            max_sz = 5
        elif max_sz < 10:
            max_sz = 10
        affinity_vis = affinity_pred[:max_sz, :max_sz]
        plt.figure()
        sns.heatmap(affinity_vis, annot=True, vmin=0.0, vmax=1.0)
        plt.savefig(os.path.join(mesh_dir, 'affinity_pred.png'))
        plt.close()

    return affinity_pred
//  [python]
def signguard(gradients, net, lr, f, byz, device, seed):
    """
    Based on the description in https://arxiv.org/abs/2109.05872
    gradients: list of gradients.
    net: model parameters.
    lr: learning rate.
    f: number of malicious clients. The first f clients are malicious.
    byz: attack type.
    device: computation device.
    seed: seed for randomness
    """
    param_list = [torch.cat([xx.reshape((-1, 1)) for xx in x], dim=0) for x in gradients]
    # let the malicious clients (first f clients) perform the byzantine attack
    if byz == attacks.fltrust_attack:
        param_list = byz(param_list, net, lr, f, device)[:-1]
    else:
        param_list = byz(param_list, net, lr, f, device)
    n = len(param_list)

    num_params = param_list[0].size(0)
    selection_fraction = 0.1

    # lower and upper bound L,R for gradient norm
    L = 0.1
    R = 3.0
    S1 = []
    S2 = []

    # compute l2-norm
    l2_norm = torch.stack([torch.norm(g.flatten(), p=2.0) for g in param_list])

    # compute element wise sign
    num_selection = int(num_params * selection_fraction)
    perm = torch.randperm(num_params)
    idx = perm[:num_selection]
    sign_grads = [torch.sign(g[idx]) for g in param_list]

    # norm-threshold filtering
    M = torch.median(l2_norm)
    for i in range(n):
        if L <= l2_norm[i] / M and l2_norm[i] / M <= R:
            S1.append(i)

    # compute sign statistics
    sign_pos = torch.stack([grad.eq(1.0).float().mean() for grad in sign_grads])
    sign_zero = torch.stack([grad.eq(0.0).float().mean() for grad in sign_grads])
    sign_neg = torch.stack([grad.eq(-1.0).float().mean() for grad in sign_grads])

    # sign-based clustering
    sign_feat = torch.stack([sign_pos, sign_zero, sign_neg], dim=1).detach().cpu().numpy()
    cluster = KMeans(n_clusters=2, max_iter=20, random_state=seed)
    labels = cluster.fit_predict(sign_feat)

    labels_tensor = torch.from_numpy(labels).to(device)
    count = torch.bincount(labels_tensor)
    largest_cluster = torch.argmax(count)

    for i, value in enumerate(labels_tensor):
        if value == largest_cluster:
            S2.append(i)

    # compute intersection of S1 and S2
    S = [i for i in S1 if i in S2]

    # global update
    global_update = torch.zeros(param_list[0].size()).to(device)
    for i in S:
        global_update += param_list[i]
    global_update *= 1 / len(S)

    # update the global model
    idx = 0
    for j, param in enumerate(net.parameters()):
        param.add_(global_update[idx:(idx + torch.numel(param))].reshape(tuple(param.size())), alpha=-lr)
        idx += torch.numel(param)