jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def calculate_norm(self, layer_name):
		for group in self.param_groups:
			for p in group['params']:
				param_state = self.state[p]
				if param_state['name'] == layer_name:
					layer_norm = torch.norm(torch.reshape(p.data, (p.data.shape[0], -1)), 2, 1)
		return layer_norm.cpu().detach().numpy()
//  [python]
def calculate_norm(self, layer_name):
		for group in self.param_groups:
			for p in group['params']:
				param_state = self.state[p]
				if param_state['name'] == layer_name:
					layer_norm = torch.linalg.norm(torch.reshape(p.data, (p.data.shape[0], -1)), 2, 1)
		return layer_norm.cpu().detach().numpy()
//  [python]
def _lstsq_exact(fun_with_jac, x, d, b, max_iter=5, mu=1., lambd=1., lr=1.,
                 xtol=1e-5):
    for _ in range(max_iter):
        f, J = fun_with_jac(x)
        grad = mu * J.T.mv(f)
        grad.add_(d - x - b, alpha=-lambd)
        JtJ = mu * J.T.mm(J)
        JtJ.diagonal().add_(lambd)
        try:
            p = torch.cholesky_solve(grad.unsqueeze(1),
                                     torch.linalg.cholesky(JtJ)).squeeze(1)
        except RuntimeError as exc:
            if 'singular' not in exc.args[0]:
                raise
            p = torch.linalg.solve(JtJ, grad)
        x.add_(p, alpha=-lr)
        if torch.norm(p.mul(lr), p=1) <= xtol:
            break

    return x
//  [python]
def _lstsq_exact(fun_with_jac, x, d, b, max_iter=5, mu=1., lambd=1., lr=1.,
                 xtol=1e-5):
    for _ in range(max_iter):
        f, J = fun_with_jac(x)
        grad = mu * J.T.mv(f)
        grad.add_(d - x - b, alpha=-lambd)
        JtJ = mu * J.T.mm(J)
        JtJ.diagonal().add_(lambd)
        try:
            p = torch.cholesky_solve(grad.unsqueeze(1),
                                     torch.linalg.cholesky(JtJ)).squeeze(1)
        except RuntimeError as exc:
            if 'singular' not in exc.args[0]:
                raise
            p = torch.linalg.solve(JtJ, grad)
        x.add_(p, alpha=-lr)
        if torch.linalg.norm(p.mul(lr), ord=1) <= xtol:
            break

    return x
//  [python]
def ball_cost(part,body_support_0,body_support_1):
    '''
    A function which takes particles and returns an L2 loss on the amount of overlap of the balls
    '''
    c_hip_0,c_ass_0,c_mid_0,c_nose_0,c_tip_0,c_impl_0,R_body,R_head,R_nose = body_support_0
    c_hip_1,c_ass_1,c_mid_1,c_nose_1,c_tip_1,c_impl_1,R_body,R_head,R_nose = body_support_1

    s = part[:,2]
    a_hip_00 = a_hip_0 + a_hip_delta * s
    b_hip_00 = b_hip_0 + b_hip_delta * (1.-s)

    s = part[:,11]
    a_hip_01 = a_hip_0 + a_hip_delta * s
    b_hip_01 = b_hip_0 + b_hip_delta * (1.-s)

    # first, we calculate the distances betjupween the centers of the ellipsoids
    # nose2nose
    d_n0n1 = torch.norm(c_nose_0-c_nose_1,dim=1)
    # nose2hip
    d_n0h1 = torch.norm(c_nose_0-c_hip_1,dim=1)
    d_n1h0 = torch.norm(c_nose_1-c_hip_0,dim=1)
    # hip2hip
    d_h0h1 = torch.norm(c_hip_0-c_hip_1,dim=1)
    # implant to other's nose
    d_imp0n1 = torch.norm(c_impl_0-c_nose_1,dim=1)
    # implant to other's hip
    d_imp0h1 = torch.norm(c_impl_0-c_hip_1,dim=1)

    # make a list of the actual distance between the centers
    d_actual = torch.stack([d_n0n1,d_n0h1,d_n1h0,d_h0h1,d_imp0n1,d_imp0h1]).squeeze(2)

    # make a list of the minimum allowed distance between
    cutoff_barrier = 0.8*torch.stack([(b_nose + b_nose)*torch.ones_like(b_hip_01), b_nose+b_hip_01, b_nose+b_hip_00, b_hip_00+b_hip_01, (r_impl + b_nose)*torch.ones_like(b_hip_01), r_impl+b_hip_01 ])

    # clip the overlap
    overlap = torch.clamp(cutoff_barrier-d_actual,0.,None)
    # do a kind of L2 loss, which we add everywhere
    barrier_loss = torch.mean(overlap,dim=0)

    return barrier_loss
//  [python]
def ball_cost(part,body_support_0,body_support_1):
    '''
    A function which takes particles and returns an L2 loss on the amount of overlap of the balls
    '''
    c_hip_0,c_ass_0,c_mid_0,c_nose_0,c_tip_0,c_impl_0,R_body,R_head,R_nose = body_support_0
    c_hip_1,c_ass_1,c_mid_1,c_nose_1,c_tip_1,c_impl_1,R_body,R_head,R_nose = body_support_1

    s = part[:,2]
    a_hip_00 = a_hip_0 + a_hip_delta * s
    b_hip_00 = b_hip_0 + b_hip_delta * (1.-s)

    s = part[:,11]
    a_hip_01 = a_hip_0 + a_hip_delta * s
    b_hip_01 = b_hip_0 + b_hip_delta * (1.-s)

    # first, we calculate the distances betjupween the centers of the ellipsoids
    # nose2nose
    d_n0n1 = torch.linalg.norm(c_nose_0-c_nose_1,dim=1)
    # nose2hip
    d_n0h1 = torch.linalg.norm(c_nose_0-c_hip_1,dim=1)
    d_n1h0 = torch.linalg.norm(c_nose_1-c_hip_0,dim=1)
    # hip2hip
    d_h0h1 = torch.linalg.norm(c_hip_0-c_hip_1,dim=1)
    # implant to other's nose
    d_imp0n1 = torch.linalg.norm(c_impl_0-c_nose_1,dim=1)
    # implant to other's hip
    d_imp0h1 = torch.linalg.norm(c_impl_0-c_hip_1,dim=1)

    # make a list of the actual distance between the centers
    d_actual = torch.stack([d_n0n1,d_n0h1,d_n1h0,d_h0h1,d_imp0n1,d_imp0h1]).squeeze(2)

    # make a list of the minimum allowed distance between
    cutoff_barrier = 0.8*torch.stack([(b_nose + b_nose)*torch.ones_like(b_hip_01), b_nose+b_hip_01, b_nose+b_hip_00, b_hip_00+b_hip_01, (r_impl + b_nose)*torch.ones_like(b_hip_01), r_impl+b_hip_01 ])

    # clip the overlap
    overlap = torch.clamp(cutoff_barrier-d_actual,0.,None)
    # do a kind of L2 loss, which we add everywhere
    barrier_loss = torch.mean(overlap,dim=0)

    return barrier_loss
//  [python]
def calculate_proportion(self, layer_name):
		#self.calculate_all_w_star()
		for group in self.param_groups:
			for p in group['params']:
				param_state = self.state[p]
				if 'w_star' in param_state and param_state['name'] == layer_name:
					#print(layer_name)
					self.calculate_w_star_by_layer(layer_name)
					if len(p.data.size())==4:
						ts_reshape = torch.reshape(param_state['w_star'], (param_state['w_star'].shape[0], -1))
						ts_norm = torch.norm(ts_reshape, 2, 1)
						num_selected_filters = torch.sum(ts_norm != 0).item()
						return num_selected_filters/p.data.size()[0]
					elif len(p.data.size())==2:
						num_selected_units = (param_state['w_star'] > 0.0).sum().item()
						return num_selected_units/p.data.size()[0] * p.data.size()[1]
					else:
						pass
//  [python]
def calculate_proportion(self, layer_name):
		#self.calculate_all_w_star()
		for group in self.param_groups:
			for p in group['params']:
				param_state = self.state[p]
				if 'w_star' in param_state and param_state['name'] == layer_name:
					#print(layer_name)
					self.calculate_w_star_by_layer(layer_name)
					if len(p.data.size())==4:
						ts_reshape = torch.reshape(param_state['w_star'], (param_state['w_star'].shape[0], -1))
						ts_norm = torch.linalg.norm(ts_reshape, 2, 1)
						num_selected_filters = torch.sum(ts_norm != 0).item()
						return num_selected_filters/p.data.size()[0]
					elif len(p.data.size())==2:
						num_selected_units = (param_state['w_star'] > 0.0).sum().item()
						return num_selected_units/p.data.size()[0] * p.data.size()[1]
					else:
						pass
//  [python]
def _lstsq_cg(fun, x, d, b, max_iter=5, mu=1., lambd=1., lr=1.,
              xtol=1e-5, cg_kwargs=None):
    if cg_kwargs is None:
        cg_kwargs = {}

    def obj(u):
        """used for strong-wolfe line search"""
        return 0.5 * (mu * fun(u).square().sum() + lambd * (d-u-b).square().sum())

    def dir_evaluate(u, t, p):
        """used for strong-wolfe line search"""
        u = u.add(p, alpha=t).view_as(x).requires_grad_(True)
        with torch.enable_grad():
            loss = obj(u)
        grad = autograd.grad(loss, u)[0]
        return float(loss), grad.view(-1)

    loss = obj(x)
    for _ in range(max_iter):
        J = JacobianLinearOperator(fun, x)
        f = J.f.detach()
        grad = mu * J.rmv(f)
        grad.add_(d - x - b, alpha=-lambd)
        JtJ = LinearOperator(shape=(x.numel(), x.numel()),
                             mv=lambda v: mu * J.rmv(J.mv(v)) + lambd * v)
        p = - cg(JtJ, grad, **cg_kwargs)[0]
        loss, _, t, _ = \
            _strong_wolfe(dir_evaluate, x.view(-1), lr, p.view(-1), loss,
                          grad.view(-1), grad.mul(p).sum())
        x.add_(p, alpha=t)
        if torch.norm(p.mul(t), p=1) <= xtol:
            break

    return x
//  [python]
def _lstsq_cg(fun, x, d, b, max_iter=5, mu=1., lambd=1., lr=1.,
              xtol=1e-5, cg_kwargs=None):
    if cg_kwargs is None:
        cg_kwargs = {}

    def obj(u):
        """used for strong-wolfe line search"""
        return 0.5 * (mu * fun(u).square().sum() + lambd * (d-u-b).square().sum())

    def dir_evaluate(u, t, p):
        """used for strong-wolfe line search"""
        u = u.add(p, alpha=t).view_as(x).requires_grad_(True)
        with torch.enable_grad():
            loss = obj(u)
        grad = autograd.grad(loss, u)[0]
        return float(loss), grad.view(-1)

    loss = obj(x)
    for _ in range(max_iter):
        J = JacobianLinearOperator(fun, x)
        f = J.f.detach()
        grad = mu * J.rmv(f)
        grad.add_(d - x - b, alpha=-lambd)
        JtJ = LinearOperator(shape=(x.numel(), x.numel()),
                             mv=lambda v: mu * J.rmv(J.mv(v)) + lambd * v)
        p = - cg(JtJ, grad, **cg_kwargs)[0]
        loss, _, t, _ = \
            _strong_wolfe(dir_evaluate, x.view(-1), lr, p.view(-1), loss,
                          grad.view(-1), grad.mul(p).sum())
        x.add_(p, alpha=t)
        if torch.linalg.norm(p.mul(t), ord=1) <= xtol:
            break

    return x
//  [python]
def use_w_star(self):
		#use sparse params to replace original params
		self.calculate_all_w_star()
		for group in self.param_groups:
			for p in group['params']:
				param_state = self.state[p]
				if 'w_star' in param_state:
					if len(p.data.size())==4:
						ts_reshape = torch.reshape(param_state['w_star'], (param_state['w_star'].shape[0], -1))
						ts_norm = torch.norm(ts_reshape, 2, 1)
						num_selected_filters = torch.sum(ts_norm != 0).item()
						param_state['original_params'] = copy.deepcopy(p.data)
						p.data = param_state['w_star']
						print('max:', torch.max(param_state['w_star']))
						print('min:', torch.min(param_state['w_star']))
						print('number of filters: ',  p.data.size()[0])
						print('number of selected filter:', num_selected_filters )
					elif len(p.data.size())==2:
						num_selected_units = (param_state['w_star'] > 0.0).sum().item()
						param_state['original_params'] = copy.deepcopy(p.data)
						p.data = param_state['w_star']
						print('max:', torch.max(param_state['w_star']))
						print('min:', torch.min(param_state['w_star']))
						print('number of filters: ',  p.data.size()[0] * p.data.size()[1])
						print('number of selected units:', num_selected_units )
					else:
						pass
//  [python]
def use_w_star(self):
		#use sparse params to replace original params
		self.calculate_all_w_star()
		for group in self.param_groups:
			for p in group['params']:
				param_state = self.state[p]
				if 'w_star' in param_state:
					if len(p.data.size())==4:
						ts_reshape = torch.reshape(param_state['w_star'], (param_state['w_star'].shape[0], -1))
						ts_norm = torch.linalg.norm(ts_reshape, 2, 1)
						num_selected_filters = torch.sum(ts_norm != 0).item()
						param_state['original_params'] = copy.deepcopy(p.data)
						p.data = param_state['w_star']
						print('max:', torch.max(param_state['w_star']))
						print('min:', torch.min(param_state['w_star']))
						print('number of filters: ',  p.data.size()[0])
						print('number of selected filter:', num_selected_filters )
					elif len(p.data.size())==2:
						num_selected_units = (param_state['w_star'] > 0.0).sum().item()
						param_state['original_params'] = copy.deepcopy(p.data)
						p.data = param_state['w_star']
						print('max:', torch.max(param_state['w_star']))
						print('min:', torch.min(param_state['w_star']))
						print('number of filters: ',  p.data.size()[0] * p.data.size()[1])
						print('number of selected units:', num_selected_units )
					else:
						pass
//  [python]
def test_frobenius_norm(self):
        x = torch.randn(2, 3, 4).float()
        self.assertONNX(lambda x: torch.norm(x, p="fro", dim=(0, 1), keepdim=True), x)
//  [python]
def test_frobenius_norm(self):
        x = torch.randn(2, 3, 4).float()
        self.assertONNX(lambda x: torch.linalg.norm(x, ord="fro", dim=(0, 1), keepdim=True), x)
//  [python]
def calculate_mask(self, layer_name):
		self.calculate_all_w_star()
		for group in self.param_groups:
			for p in group['params']:
				param_state = self.state[p]
				if 'w_star' in param_state and param_state['name'] == layer_name:
					if len(p.data.size())==4:
						ts_reshape = torch.reshape(param_state['w_star'], (param_state['w_star'].shape[0], -1))
						ts_norm = torch.norm(ts_reshape, 2, 1)
						num_selected_filters = torch.sum(ts_norm != 0).item()
						mask = torch.ones_like(param_state['w_star'])
						mask[ts_norm != 0, :, :, :] = 0
						return mask
					elif len(p.data.size())==2:
						return torch.ones_like(param_state['w_star'])
						#return torch.gt(torch.abs(param_state['w_star']), 0.0).float()
					else:
						pass
//  [python]
def calculate_mask(self, layer_name):
		self.calculate_all_w_star()
		for group in self.param_groups:
			for p in group['params']:
				param_state = self.state[p]
				if 'w_star' in param_state and param_state['name'] == layer_name:
					if len(p.data.size())==4:
						ts_reshape = torch.reshape(param_state['w_star'], (param_state['w_star'].shape[0], -1))
						ts_norm = torch.linalg.norm(ts_reshape, 2, 1)
						num_selected_filters = torch.sum(ts_norm != 0).item()
						mask = torch.ones_like(param_state['w_star'])
						mask[ts_norm != 0, :, :, :] = 0
						return mask
					elif len(p.data.size())==2:
						return torch.ones_like(param_state['w_star'])
						#return torch.gt(torch.abs(param_state['w_star']), 0.0).float()
					else:
						pass
//  [python]
def _reset_parameters(self, init_weight):
        scale = torch.norm(init_weight, dim=-1).unsqueeze(-1)
        self.extra_output.weight = nn.Parameter(init_weight / (scale ** 2))
        init.constant_(self.extra_output.bias, 0.)
//  [python]
def _reset_parameters(self, init_weight):
        scale = torch.linalg.norm(init_weight, dim=-1).unsqueeze(-1)
        self.extra_output.weight = nn.Parameter(init_weight / (scale ** 2))
        init.constant_(self.extra_output.bias, 0.)
//  [python]
def quaternion_to_axis_angle(quaternions: torch.Tensor) -> torch.Tensor:
    """
    Convert rotations given as quaternions to axis/angle.
    Args:
        quaternions: quaternions with real part first,
            as tensor of shape (..., 4).
    Returns:
        Rotations given as a vector in axis angle form, as a tensor
            of shape (..., 3), where the magnitude is the angle
            turned anticlockwise in radians around the vector's
            direction.
    """
    norms = torch.norm(quaternions[..., 1:], p=2, dim=-1, keepdim=True)
    half_angles = torch.atan2(norms, quaternions[..., :1])
    angles = 2 * half_angles
    eps = 1e-6
    small_angles = angles.abs() < eps
    sin_half_angles_over_angles = torch.empty_like(angles)
    sin_half_angles_over_angles[~small_angles] = (
        torch.sin(half_angles[~small_angles]) / angles[~small_angles]
    )
    # for x small, sin(x/2) is about x/2 - (x/2)^3/6
    # so sin(x/2)/x is about 1/2 - (x*x)/48
    sin_half_angles_over_angles[small_angles] = (
        0.5 - (angles[small_angles] * angles[small_angles]) / 48
    )
    return quaternions[..., 1:] / sin_half_angles_over_angles
//  [python]
def quaternion_to_axis_angle(quaternions: torch.Tensor) -> torch.Tensor:
    """
    Convert rotations given as quaternions to axis/angle.
    Args:
        quaternions: quaternions with real part first,
            as tensor of shape (..., 4).
    Returns:
        Rotations given as a vector in axis angle form, as a tensor
            of shape (..., 3), where the magnitude is the angle
            turned anticlockwise in radians around the vector's
            direction.
    """
    norms = torch.linalg.norm(quaternions[..., 1:], ord=2, dim=-1, keepdim=True)
    half_angles = torch.atan2(norms, quaternions[..., :1])
    angles = 2 * half_angles
    eps = 1e-6
    small_angles = angles.abs() < eps
    sin_half_angles_over_angles = torch.empty_like(angles)
    sin_half_angles_over_angles[~small_angles] = (
        torch.sin(half_angles[~small_angles]) / angles[~small_angles]
    )
    # for x small, sin(x/2) is about x/2 - (x/2)^3/6
    # so sin(x/2)/x is about 1/2 - (x*x)/48
    sin_half_angles_over_angles[small_angles] = (
        0.5 - (angles[small_angles] * angles[small_angles]) / 48
    )
    return quaternions[..., 1:] / sin_half_angles_over_angles
//  [python]
def forward(self, x):
        logits = []
        bbox_reg = []
        centerness = []

        dyhead_tower = self.dyhead_tower(x)

        for l, feature in enumerate(x):
            if self.cls_logits_bias is None:
                logit = self.cls_logits(dyhead_tower[l])
            else:
                # CosineSimOutputLayers: https://github.com/ucbdrive/few-shot-object-detection/blob/master/fsdet/modeling/roi_heads/fast_rcnn.py#L448-L464
                # normalize the input x along the `channel` dimension
                x_norm = torch.norm(dyhead_tower[l], p=2, dim=1, keepdim=True).expand_as(dyhead_tower[l])
                x_normalized = dyhead_tower[l].div(x_norm + 1e-5)
                # normalize weight
                temp_norm = (
                    torch.norm(self.cls_logits.weight.data, p=2, dim=1, keepdim=True)
                        .expand_as(self.cls_logits.weight.data)
                )
                self.cls_logits.weight.data = self.cls_logits.weight.data.div(
                    temp_norm + 1e-5
                )
                cos_dist = self.cls_logits(x_normalized)
                logit = self.cosine_scale * cos_dist + self.cls_logits_bias.reshape(1, len(self.cls_logits_bias), 1, 1)
            logits.append(logit)

            bbox_pred = self.scales[l](self.bbox_pred(dyhead_tower[l]))
            bbox_reg.append(bbox_pred)

            centerness.append(self.centerness(dyhead_tower[l]))
        return logits, bbox_reg, centerness
//  [python]
def forward(self, x):
        logits = []
        bbox_reg = []
        centerness = []

        dyhead_tower = self.dyhead_tower(x)

        for l, feature in enumerate(x):
            if self.cls_logits_bias is None:
                logit = self.cls_logits(dyhead_tower[l])
            else:
                # CosineSimOutputLayers: https://github.com/ucbdrive/few-shot-object-detection/blob/master/fsdet/modeling/roi_heads/fast_rcnn.py#L448-L464
                # normalize the input x along the `channel` dimension
                x_norm = torch.linalg.norm(dyhead_tower[l], ord=2, dim=1, keepdims=True).expand_as(dyhead_tower[l])
                x_normalized = dyhead_tower[l].div(x_norm + 1e-5)
                # normalize weight
                temp_norm = (
                    torch.linalg.norm(self.cls_logits.weight.data, ord=2, dim=1, keepdims=True)
                        .expand_as(self.cls_logits.weight.data)
                )
                self.cls_logits.weight.data = self.cls_logits.weight.data.div(
                    temp_norm + 1e-5
                )
                cos_dist = self.cls_logits(x_normalized)
                logit = self.cosine_scale * cos_dist + self.cls_logits_bias.reshape(1, len(self.cls_logits_bias), 1, 1)
            logits.append(logit)

            bbox_pred = self.scales[l](self.bbox_pred(dyhead_tower[l]))
            bbox_reg.append(bbox_pred)

            centerness.append(self.centerness(dyhead_tower[l]))
        return logits, bbox_reg, centerness
//  [python]
def graph_fusion(graphlist, wlist):
    torch.cuda.empty_cache()
    with torch.no_grad():
        fusiongraph =torch.sum(torch.mul(graphlist, wlist.T.unsqueeze(dim=1)), dim=0)/ (torch.sum(wlist))
    for v in range(graphlist.shape[0]):
        w = 1 / (2 * (torch.norm((graphlist[v] - fusiongraph), p='fro')))
        wlist.T[v] = w

    return fusiongraph, wlist
//  [python]
def graph_fusion(graphlist, wlist):
    torch.cuda.empty_cache()
    with torch.no_grad():
        fusiongraph =torch.sum(torch.mul(graphlist, wlist.T.unsqueeze(dim=1)), dim=0)/ (torch.sum(wlist))
    for v in range(graphlist.shape[0]):
        w = 1 / (2 * (torch.linalg.norm((graphlist[v] - fusiongraph), ord='fro')))
        wlist.T[v] = w

    return fusiongraph, wlist
//  [python]
def _geometric_median(inputs, weights, eps=1e-6, maxiter=100, ftol=1e-20):
        weighted_average = (
            lambda inputs, weights: torch.sum(weights.view(-1, 1) * inputs, dim=0)
            / weights.sum()
        )

        def obj_func(median, inputs, weights):
            # This function is not used so far,
            # as the numpy version appears to be more accurate (and I don't know why).

            # norms = torch.norm(inputs - median, dim=1)
            # return (torch.sum(norms * weights) / torch.sum(weights)).item()

            return np.average(
                [torch.norm(p - median).item() for p in inputs],
                weights=weights.cpu(),
            )

        with torch.no_grad():
            median = weighted_average(inputs, weights)
            new_weights = weights
            objective_value = obj_func(median, inputs, weights)

            # Weiszfeld iterations
            for _ in range(maxiter):
                prev_obj_value = objective_value
                denom = torch.stack([torch.norm(p - median) for p in inputs])
                new_weights = weights / torch.clamp(denom, min=eps)
                median = weighted_average(inputs, new_weights)

                objective_value = obj_func(median, inputs, weights)
                if abs(prev_obj_value - objective_value) <= ftol * objective_value:
                    break

        median = weighted_average(inputs, new_weights)
        return median
//  [python]
def _geometric_median(inputs, weights, eps=1e-6, maxiter=100, ftol=1e-20):
        weighted_average = (
            lambda inputs, weights: torch.sum(weights.view(-1, 1) * inputs, dim=0)
            / weights.sum()
        )

        def obj_func(median, inputs, weights):
            # This function is not used so far,
            # as the numpy version appears to be more accurate (and I don't know why).

            # norms = torch.linalg.norm(inputs - median, dim=1)
            # return (torch.sum(norms * weights) / torch.sum(weights)).item()

            return np.average(
                [torch.linalg.norm(p - median).item() for p in inputs],
                weights=weights.cpu(),
            )

        with torch.no_grad():
            median = weighted_average(inputs, weights)
            new_weights = weights
            objective_value = obj_func(median, inputs, weights)

            # Weiszfeld iterations
            for _ in range(maxiter):
                prev_obj_value = objective_value
                denom = torch.stack([torch.linalg.norm(p - median) for p in inputs])
                new_weights = weights / torch.clamp(denom, min=eps)
                median = weighted_average(inputs, new_weights)

                objective_value = obj_func(median, inputs, weights)
                if abs(prev_obj_value - objective_value) <= ftol * objective_value:
                    break

        median = weighted_average(inputs, new_weights)
        return median
//  [python]
def update_prune_order(self, epoch):
		self.calculate_all_w_star()
		for group in self.param_groups:
			for p in group['params']:
				param_state = self.state[p]
				if 'z_buffer' in param_state:
					if len(p.data.size())==4:
						if 'epoch_record' not in param_state:
							param_state['epoch_record'] = torch.zeros_like(p.data).add_(2000.0)
							mask = (torch.gt(torch.abs(param_state['gamma_buffer']), 0.0)).float()
							param_state['epoch_record'] = torch.min((1.0 - mask) * 2000.0 + mask * float(epoch),  param_state['epoch_record'])
							epoch_per_filer, _ = torch.min(torch.reshape(param_state['epoch_record'], (param_state['epoch_record'].shape[0], -1)), dim=1) 
							param_state['prune_order'] = torch.norm(torch.reshape(p.data, (p.data.shape[0], -1)), 2, 1) - epoch_per_filer
						else:
							mask = (torch.gt(torch.abs(param_state['gamma_buffer']), 0.0)).float()
							#print(mask.size())
							#print(param_state['epoch_record'].size())
							param_state['epoch_record'] = torch.min((1.0 - mask) * 2000.0 + mask * float(epoch),  param_state['epoch_record'])
							epoch_per_filer, _ = torch.min(torch.reshape(param_state['epoch_record'], (param_state['epoch_record'].shape[0], -1)), dim=1) 
							param_state['prune_order'] = torch.norm(torch.reshape(p.data, (p.data.shape[0], -1)), 2, 1) - epoch_per_filer		
					elif len(p.data.size()) == 2:
						if 'epoch_record' not in param_state:
							param_state['epoch_record'] = torch.zeros_like(p.data).add_(2000.0)
							mask = (torch.gt(torch.abs(param_state['gamma_buffer']), 0.0)).float()
							param_state['epoch_record'] = torch.min((1.0 - mask) * 2000.0 + mask * float(epoch),  param_state['epoch_record'])
							param_state['prune_order'] = torch.abs(p.data) - param_state['epoch_record']
						else:
							mask = (torch.gt(torch.abs(param_state['gamma_buffer']), 0.0)).float()
							param_state['epoch_record'] = torch.min((1.0 - mask) * 2000.0 + mask * float(epoch),  param_state['epoch_record'])
							#param_state['prune_order'] = torch.abs(param_state['w_star']) - param_state['epoch_record']
							param_state['prune_order'] = torch.abs(p.data) - param_state['epoch_record']
					else:
						pass
//  [python]
def update_prune_order(self, epoch):
		self.calculate_all_w_star()
		for group in self.param_groups:
			for p in group['params']:
				param_state = self.state[p]
				if 'z_buffer' in param_state:
					if len(p.data.size())==4:
						if 'epoch_record' not in param_state:
							param_state['epoch_record'] = torch.zeros_like(p.data).add_(2000.0)
							mask = (torch.gt(torch.abs(param_state['gamma_buffer']), 0.0)).float()
							param_state['epoch_record'] = torch.min((1.0 - mask) * 2000.0 + mask * float(epoch),  param_state['epoch_record'])
							epoch_per_filer, _ = torch.min(torch.reshape(param_state['epoch_record'], (param_state['epoch_record'].shape[0], -1)), dim=1) 
							param_state['prune_order'] = torch.linalg.norm(torch.reshape(p.data, (p.data.shape[0], -1)), 2, 1) - epoch_per_filer
						else:
							mask = (torch.gt(torch.abs(param_state['gamma_buffer']), 0.0)).float()
							#print(mask.size())
							#print(param_state['epoch_record'].size())
							param_state['epoch_record'] = torch.min((1.0 - mask) * 2000.0 + mask * float(epoch),  param_state['epoch_record'])
							epoch_per_filer, _ = torch.min(torch.reshape(param_state['epoch_record'], (param_state['epoch_record'].shape[0], -1)), dim=1) 
							param_state['prune_order'] = torch.linalg.norm(torch.reshape(p.data, (p.data.shape[0], -1)), 2, 1) - epoch_per_filer		
					elif len(p.data.size()) == 2:
						if 'epoch_record' not in param_state:
							param_state['epoch_record'] = torch.zeros_like(p.data).add_(2000.0)
							mask = (torch.gt(torch.abs(param_state['gamma_buffer']), 0.0)).float()
							param_state['epoch_record'] = torch.min((1.0 - mask) * 2000.0 + mask * float(epoch),  param_state['epoch_record'])
							param_state['prune_order'] = torch.abs(p.data) - param_state['epoch_record']
						else:
							mask = (torch.gt(torch.abs(param_state['gamma_buffer']), 0.0)).float()
							param_state['epoch_record'] = torch.min((1.0 - mask) * 2000.0 + mask * float(epoch),  param_state['epoch_record'])
							#param_state['prune_order'] = torch.abs(param_state['w_star']) - param_state['epoch_record']
							param_state['prune_order'] = torch.abs(p.data) - param_state['epoch_record']
					else:
						pass
//  [python]
def forward(self, mels):
        self.lstm.flatten_parameters()
        _, (hidden, _) = self.lstm(mels)
        embeds_raw = self.relu(self.linear(hidden[-1]))
        return embeds_raw / torch.norm(embeds_raw, dim=1, keepdim=True)
//  [python]
def forward(self, mels):
        self.lstm.flatten_parameters()
        _, (hidden, _) = self.lstm(mels)
        embeds_raw = self.relu(self.linear(hidden[-1]))
        return embeds_raw / torch.linalg.norm(embeds_raw, dim=1, keepdim=True)
//  [python]
def _compute_reg_loss(self, vd, ue, edge_group_to_vd, vd_num_edges):
        """
        Regularizer L_dev as in Equation 8
        """
        dist = torch.norm(ue - torch.index_select(input=vd, index=edge_group_to_vd, dim=0), dim=-1)
        mean_l2 = torch.zeros_like(vd[:, 0])
        mean_l2 = (mean_l2).index_add_(0, edge_group_to_vd, dist) / vd_num_edges.squeeze(1).float()
        mad = (dist - torch.index_select(input=mean_l2, index=edge_group_to_vd, dim=0)).abs()
        return mad
//  [python]
def _compute_reg_loss(self, vd, ue, edge_group_to_vd, vd_num_edges):
        """
        Regularizer L_dev as in Equation 8
        """
        dist = torch.linalg.norm(ue - torch.index_select(input=vd, index=edge_group_to_vd, dim=0), dim=-1)
        mean_l2 = torch.zeros_like(vd[:, 0])
        mean_l2 = (mean_l2).index_add_(0, edge_group_to_vd, dist) / vd_num_edges.squeeze(1).float()
        mad = (dist - torch.index_select(input=mean_l2, index=edge_group_to_vd, dim=0)).abs()
        return mad
//  [python]
def create_projection_matrix(m, d, seed=0, scaling=0, struct_mode=False):
    nb_full_blocks = int(m/d)
    block_list = []
    current_seed = seed
    for _ in range(nb_full_blocks):
        torch.manual_seed(current_seed)
        if struct_mode:
            q = create_products_of_givens_rotations(d, current_seed)
        else:
            unstructured_block = torch.randn((d, d))
            q, _ = torch.qr(unstructured_block)
            q = torch.t(q)
        block_list.append(q)
        current_seed += 1
    remaining_rows = m - nb_full_blocks * d
    if remaining_rows > 0:
        torch.manual_seed(current_seed)
        if struct_mode:
            q = create_products_of_givens_rotations(d, current_seed)
        else:
            unstructured_block = torch.randn((d, d))
            q, _ = torch.qr(unstructured_block)
            q = torch.t(q)
        block_list.append(q[0:remaining_rows])
    final_matrix = torch.vstack(block_list)

    current_seed += 1
    torch.manual_seed(current_seed)
    if scaling == 0:
        multiplier = torch.norm(torch.randn((m, d)), dim=1)
    elif scaling == 1:
        multiplier = torch.sqrt(torch.tensor(float(d))) * torch.ones(m)
    else:
        raise ValueError("Scaling must be one of {0, 1}. Was %s" % scaling)

    return torch.matmul(torch.diag(multiplier), final_matrix)
//  [python]
def create_projection_matrix(m, d, seed=0, scaling=0, struct_mode=False):
    nb_full_blocks = int(m/d)
    block_list = []
    current_seed = seed
    for _ in range(nb_full_blocks):
        torch.manual_seed(current_seed)
        if struct_mode:
            q = create_products_of_givens_rotations(d, current_seed)
        else:
            unstructured_block = torch.randn((d, d))
            q, _ = torch.qr(unstructured_block)
            q = torch.t(q)
        block_list.append(q)
        current_seed += 1
    remaining_rows = m - nb_full_blocks * d
    if remaining_rows > 0:
        torch.manual_seed(current_seed)
        if struct_mode:
            q = create_products_of_givens_rotations(d, current_seed)
        else:
            unstructured_block = torch.randn((d, d))
            q, _ = torch.qr(unstructured_block)
            q = torch.t(q)
        block_list.append(q[0:remaining_rows])
    final_matrix = torch.vstack(block_list)

    current_seed += 1
    torch.manual_seed(current_seed)
    if scaling == 0:
        multiplier = torch.linalg.norm(torch.randn((m, d)), dim=1)
    elif scaling == 1:
        multiplier = torch.sqrt(torch.tensor(float(d))) * torch.ones(m)
    else:
        raise ValueError("Scaling must be one of {0, 1}. Was %s" % scaling)

    return torch.matmul(torch.diag(multiplier), final_matrix)
//  [python]
def rodrigues(r):
        """
        Rodrigues' rotation formula that turns axis-angle tensor into rotation
        matrix in a batch-ed manner.

        Parameter:
        ----------
        r: Axis-angle rotation tensor of shape [batch_size * angle_num, 1, 3].

        Return:
        -------
        Rotation matrix of shape [batch_size * angle_num, 3, 3].

        """
        eps = r.clone().normal_(std=1e-8)
        theta = torch.norm(r + eps, dim=(1, 2), keepdim=True)  # dim cannot be tuple
        theta_dim = theta.shape[0]
        r_hat = r / theta
        cos = torch.cos(theta)
        z_stick = torch.zeros(theta_dim, dtype=r.dtype).to(r.device)
        m = torch.stack(
        (z_stick, -r_hat[:, 0, 2], r_hat[:, 0, 1], r_hat[:, 0, 2], z_stick,
            -r_hat[:, 0, 0], -r_hat[:, 0, 1], r_hat[:, 0, 0], z_stick), dim=1)
        m = torch.reshape(m, (-1, 3, 3))
        i_cube = (torch.eye(3, dtype=r.dtype).unsqueeze(dim=0) \
                + torch.zeros((theta_dim, 3, 3), dtype=r.dtype)).to(r.device)
        A = r_hat.permute(0, 2, 1)
        dot = torch.matmul(A, r_hat)
        R = cos * i_cube + (1 - cos) * dot + torch.sin(theta) * m
        return R
//  [python]
def rodrigues(r):
        """
        Rodrigues' rotation formula that turns axis-angle tensor into rotation
        matrix in a batch-ed manner.

        Parameter:
        ----------
        r: Axis-angle rotation tensor of shape [batch_size * angle_num, 1, 3].

        Return:
        -------
        Rotation matrix of shape [batch_size * angle_num, 3, 3].

        """
        eps = r.clone().normal_(std=1e-8)
        theta = torch.linalg.norm(r + eps, dim=(1, 2), keepdim=True)  # dim cannot be tuple
        theta_dim = theta.shape[0]
        r_hat = r / theta
        cos = torch.cos(theta)
        z_stick = torch.zeros(theta_dim, dtype=r.dtype).to(r.device)
        m = torch.stack(
        (z_stick, -r_hat[:, 0, 2], r_hat[:, 0, 1], r_hat[:, 0, 2], z_stick,
            -r_hat[:, 0, 0], -r_hat[:, 0, 1], r_hat[:, 0, 0], z_stick), dim=1)
        m = torch.reshape(m, (-1, 3, 3))
        i_cube = (torch.eye(3, dtype=r.dtype).unsqueeze(dim=0) \
                + torch.zeros((theta_dim, 3, 3), dtype=r.dtype)).to(r.device)
        A = r_hat.permute(0, 2, 1)
        dot = torch.matmul(A, r_hat)
        R = cos * i_cube + (1 - cos) * dot + torch.sin(theta) * m
        return R
//  [python]
def get_complete_graph(self, idx, n_atoms, start):
        if idx in self.complete_graphs:
            return self.complete_graphs[idx].to(self.device)
        else:
            src, dst = self.get_pairwise(n_atoms)
            g = dgl.graph((src, dst), device=self.device)
            g.ndata['feat'] = self.features_tensor[start: start + n_atoms].to(self.device)
            g.ndata['x'] = self.coordinates[start: start + n_atoms].to(self.device)
            g.edata['d'] = torch.norm(g.ndata['x'][g.edges()[0]] - g.ndata['x'][g.edges()[1]], p=2, dim=-1).unsqueeze(
                -1).detach()
            self.complete_graphs[idx] = g.to('cpu')
            return g
//  [python]
def get_complete_graph(self, idx, n_atoms, start):
        if idx in self.complete_graphs:
            return self.complete_graphs[idx].to(self.device)
        else:
            src, dst = self.get_pairwise(n_atoms)
            g = dgl.graph((src, dst), device=self.device)
            g.ndata['feat'] = self.features_tensor[start: start + n_atoms].to(self.device)
            g.ndata['x'] = self.coordinates[start: start + n_atoms].to(self.device)
            g.edata['d'] = torch.linalg.norm(g.ndata['x'][g.edges()[0]] - g.ndata['x'][g.edges()[1]], ord=2, dim=-1).unsqueeze(
                -1).detach()
            self.complete_graphs[idx] = g.to('cpu')
            return g
//  [python]
def Stable_loss(net,Nstate):
    x_ref = np.zeros(Nstate)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    x_ref_lift = net.encode_only(torch.DoubleTensor(x_ref).to(device))
    loss = torch.norm(x_ref_lift)
    return loss
//  [python]
def Stable_loss(net,Nstate):
    x_ref = np.zeros(Nstate)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    x_ref_lift = net.encode_only(torch.DoubleTensor(x_ref).to(device))
    loss = torch.linalg.norm(x_ref_lift)
    return loss
//  [python]
def batch_rodrigues(theta):
    """Convert axis-angle representation to rotation matrix.
    Args:
        theta: size = [B, 3]
    Returns:
        Rotation matrix corresponding to the quaternion -- size = [B, 3, 3]
    """
    l1norm = torch.norm(theta + 1e-8, p = 2, dim = 1)
    angle = torch.unsqueeze(l1norm, -1)
    normalized = torch.div(theta, angle)
    angle = angle * 0.5
    v_cos = torch.cos(angle)
    v_sin = torch.sin(angle)
    quat = torch.cat([v_cos, v_sin * normalized], dim = 1)
    return quat_to_rotmat(quat)
//  [python]
def batch_rodrigues(theta):
    """Convert axis-angle representation to rotation matrix.
    Args:
        theta: size = [B, 3]
    Returns:
        Rotation matrix corresponding to the quaternion -- size = [B, 3, 3]
    """
    l1norm = torch.linalg.norm(theta + 1e-8, ord = 2, dim = 1)
    angle = torch.unsqueeze(l1norm, -1)
    normalized = torch.div(theta, angle)
    angle = angle * 0.5
    v_cos = torch.cos(angle)
    v_sin = torch.sin(angle)
    quat = torch.cat([v_cos, v_sin * normalized], dim = 1)
    return quat_to_rotmat(quat)
//  [python]
def rigid_from_3_points(N, Ca, C, is_na=None, eps=1e-4):
    dims = N.shape[:-1]

    v1 = C-Ca
    v2 = N-Ca
    e1 = v1/(torch.norm(v1, dim=-1, keepdim=True)+eps)
    u2 = v2-(torch.einsum('...li, ...li -> ...l', e1, v2)[...,None]*e1)
    e2 = u2/(torch.norm(u2, dim=-1, keepdim=True)+eps)
    e3 = torch.cross(e1, e2, dim=-1)
    R = torch.cat([e1[...,None], e2[...,None], e3[...,None]], axis=-1) #[B,L,3,3] - rotation matrix

    v2 = v2/(torch.norm(v2, dim=-1, keepdim=True)+eps)
    cosref = torch.sum(e1*v2, dim=-1)

    costgt = torch.full(dims, -0.3616, device=N.device)
    if is_na is not None:
       costgt[is_na] = ChemData().costgtNA

    cos2del = torch.clamp( cosref*costgt + torch.sqrt((1-cosref*cosref)*(1-costgt*costgt)+eps), min=-1.0, max=1.0 )

    cosdel = torch.sqrt(0.5*(1+cos2del)+eps)

    sindel = torch.sign(costgt-cosref) * torch.sqrt(1-0.5*(1+cos2del)+eps)

    Rp = torch.eye(3, device=N.device).repeat(*dims,1,1)
    Rp[...,0,0] = cosdel
    Rp[...,0,1] = -sindel
    Rp[...,1,0] = sindel
    Rp[...,1,1] = cosdel
    R = torch.einsum('...ij,...jk->...ik', R,Rp)

    return R, Ca
//  [python]
def rigid_from_3_points(N, Ca, C, is_na=None, eps=1e-4):
    dims = N.shape[:-1]

    v1 = C-Ca
    v2 = N-Ca
    e1 = v1/(torch.linalg.norm(v1, dim=-1, keepdim=True)+eps)
    u2 = v2-(torch.einsum('...li, ...li -> ...l', e1, v2)[...,None]*e1)
    e2 = u2/(torch.linalg.norm(u2, dim=-1, keepdim=True)+eps)
    e3 = torch.cross(e1, e2, dim=-1)
    R = torch.cat([e1[...,None], e2[...,None], e3[...,None]], axis=-1) #[B,L,3,3] - rotation matrix

    v2 = v2/(torch.linalg.norm(v2, dim=-1, keepdim=True)+eps)
    cosref = torch.sum(e1*v2, dim=-1)

    costgt = torch.full(dims, -0.3616, device=N.device)
    if is_na is not None:
       costgt[is_na] = ChemData().costgtNA

    cos2del = torch.clamp( cosref*costgt + torch.sqrt((1-cosref*cosref)*(1-costgt*costgt)+eps), min=-1.0, max=1.0 )

    cosdel = torch.sqrt(0.5*(1+cos2del)+eps)

    sindel = torch.sign(costgt-cosref) * torch.sqrt(1-0.5*(1+cos2del)+eps)

    Rp = torch.eye(3, device=N.device).repeat(*dims,1,1)
    Rp[...,0,0] = cosdel
    Rp[...,0,1] = -sindel
    Rp[...,1,0] = sindel
    Rp[...,1,1] = cosdel
    R = torch.einsum('...ij,...jk->...ik', R,Rp)

    return R, Ca
//  [python]
def _h_sig(p_sigma: torch.Tensor, c_sigma: float, iter: int) -> torch.Tensor:
    """Boolean flag for stalling the update to the evolution path for rank-1 updates
    Args:
        p_sigma (torch.Tensor): The evolution path for step-size updates
        c_sigma (float): The learning rate for step-size updates
        iter (int): The current iteration (generation)
    Returns:
        stall (torch.Tensor): Whether to stall the update to p_c, expressed as a single torch float with 0 = continue, 1 = stall
    """
    # infer dimension from p_sigma
    d = p_sigma.shape[-1]
    # Compute the discounted squared sum
    squared_sum = torch.norm(p_sigma).pow(2.0) / (1 - (1 - c_sigma) ** (2 * iter + 1))
    # Check boolean flag and return
    stall = (squared_sum / d) - 1 < 1 + 4.0 / (d + 1)
    return stall.any().to(p_sigma.dtype)
//  [python]
def _h_sig(p_sigma: torch.Tensor, c_sigma: float, iter: int) -> torch.Tensor:
    """Boolean flag for stalling the update to the evolution path for rank-1 updates
    Args:
        p_sigma (torch.Tensor): The evolution path for step-size updates
        c_sigma (float): The learning rate for step-size updates
        iter (int): The current iteration (generation)
    Returns:
        stall (torch.Tensor): Whether to stall the update to p_c, expressed as a single torch float with 0 = continue, 1 = stall
    """
    # infer dimension from p_sigma
    d = p_sigma.shape[-1]
    # Compute the discounted squared sum
    squared_sum = torch.linalg.norm(p_sigma).pow(2.0) / (1 - (1 - c_sigma) ** (2 * iter + 1))
    # Check boolean flag and return
    stall = (squared_sum / d) - 1 < 1 + 4.0 / (d + 1)
    return stall.any().to(p_sigma.dtype)
//  [python]
def rigid_from_3_points(N, Ca, C, non_ideal=False, eps=1e-8):
    # N, Ca, C - [B,L, 3]
    # R - [B,L, 3, 3], det(R)=1, inv(R) = R.T, R is a rotation matrix
    B, L = N.shape[:2]

    v1 = C - Ca
    v2 = N - Ca
    e1 = v1 / (torch.norm(v1, dim=-1, keepdim=True) + eps)
    u2 = v2 - (torch.einsum("bli, bli -> bl", e1, v2)[..., None] * e1)
    e2 = u2 / (torch.norm(u2, dim=-1, keepdim=True) + eps)
    e3 = torch.cross(e1, e2, dim=-1)
    R = torch.cat(
        [e1[..., None], e2[..., None], e3[..., None]], axis=-1
    )  # [B,L,3,3] - rotation matrix

    if non_ideal:
        v2 = v2 / (torch.norm(v2, dim=-1, keepdim=True) + eps)
        cosref = torch.sum(e1 * v2, dim=-1)  # cosine of current N-CA-C bond angle
        costgt = cos_ideal_NCAC.item()
        cos2del = torch.clamp(
            cosref * costgt
            + torch.sqrt((1 - cosref * cosref) * (1 - costgt * costgt) + eps),
            min=-1.0,
            max=1.0,
        )
        cosdel = torch.sqrt(0.5 * (1 + cos2del) + eps)
        sindel = torch.sign(costgt - cosref) * torch.sqrt(1 - 0.5 * (1 + cos2del) + eps)
        Rp = torch.eye(3, device=N.device).repeat(B, L, 1, 1)
        Rp[:, :, 0, 0] = cosdel
        Rp[:, :, 0, 1] = -sindel
        Rp[:, :, 1, 0] = sindel
        Rp[:, :, 1, 1] = cosdel

        R = torch.einsum("blij,bljk->blik", R, Rp)

    return R, Ca
//  [python]
def rigid_from_3_points(N, Ca, C, non_ideal=False, eps=1e-8):
    # N, Ca, C - [B,L, 3]
    # R - [B,L, 3, 3], det(R)=1, inv(R) = R.T, R is a rotation matrix
    B, L = N.shape[:2]

    v1 = C - Ca
    v2 = N - Ca
    e1 = v1 / (torch.linalg.norm(v1, dim=-1, keepdim=True) + eps)
    u2 = v2 - (torch.einsum("bli, bli -> bl", e1, v2)[..., None] * e1)
    e2 = u2 / (torch.linalg.norm(u2, dim=-1, keepdim=True) + eps)
    e3 = torch.cross(e1, e2, dim=-1)
    R = torch.cat(
        [e1[..., None], e2[..., None], e3[..., None]], axis=-1
    )  # [B,L,3,3] - rotation matrix

    if non_ideal:
        v2 = v2 / (torch.linalg.norm(v2, dim=-1, keepdim=True) + eps)
        cosref = torch.sum(e1 * v2, dim=-1)  # cosine of current N-CA-C bond angle
        costgt = cos_ideal_NCAC.item()
        cos2del = torch.clamp(
            cosref * costgt
            + torch.sqrt((1 - cosref * cosref) * (1 - costgt * costgt) + eps),
            min=-1.0,
            max=1.0,
        )
        cosdel = torch.sqrt(0.5 * (1 + cos2del) + eps)
        sindel = torch.sign(costgt - cosref) * torch.sqrt(1 - 0.5 * (1 + cos2del) + eps)
        Rp = torch.eye(3, device=N.device).repeat(B, L, 1, 1)
        Rp[:, :, 0, 0] = cosdel
        Rp[:, :, 0, 1] = -sindel
        Rp[:, :, 1, 0] = sindel
        Rp[:, :, 1, 1] = cosdel

        R = torch.einsum("blij,bljk->blik", R, Rp)

    return R, Ca
//  [python]
def update_sigma(self) -> None:
        """Update the step size sigma according to its evolution path p_sigma
        This operation is bounded O(d), with the most expensive component being the norm of the evolution path, a d-dimensional vector.
        """
        d = self._problem.solution_length
        # Compute the exponential update
        if self.csa_squared:
            # Exponential update based on natural gradient maximizing squared norm of p_sigma
            exponential_update = (torch.norm(self.p_sigma).pow(2.0) / d - 1) / 2
        else:
            # Exponential update increasing likelihood p_sigma having expected norm
            exponential_update = torch.norm(self.p_sigma) / self.unbiased_expectation - 1
        # Rescale exponential update based on learning rate + damping factor
        exponential_update = (self.c_sigma / self.damp_sigma) * exponential_update
        # Multiplicative update to sigma
        self.sigma = self.sigma * torch.exp(exponential_update)
//  [python]
def update_sigma(self) -> None:
        """Update the step size sigma according to its evolution path p_sigma
        This operation is bounded O(d), with the most expensive component being the norm of the evolution path, a d-dimensional vector.
        """
        d = self._problem.solution_length
        # Compute the exponential update
        if self.csa_squared:
            # Exponential update based on natural gradient maximizing squared norm of p_sigma
            exponential_update = (torch.linalg.norm(self.p_sigma).pow(2.0) / d - 1) / 2
        else:
            # Exponential update increasing likelihood p_sigma having expected norm
            exponential_update = torch.linalg.norm(self.p_sigma) / self.unbiased_expectation - 1
        # Rescale exponential update based on learning rate + damping factor
        exponential_update = (self.c_sigma / self.damp_sigma) * exponential_update
        # Multiplicative update to sigma
        self.sigma = self.sigma * torch.exp(exponential_update)
//  [python]
def fpcm_MOVE_sl(a, env, ctm_args=cfg.ctm_args, global_args=cfg.global_args,
    past_steps_data=None):
    verbosity= ctm_args.verbosity_fpcm_move
    # 0) extract raw tensors as tuple
    tensors= tuple([a,env.C[env.keyC],env.T[env.keyT]])
    
    # function wrapping up the core of the FPCM MOVE segment of CTM algorithm
    def fpcm_MOVE_sl_c(*tensors):
        a, C, T= tensors

        # 1) compute fixed points for U, T
        i=0
        e0=1.0e+16
        tmp_=[(T,C)]
        while e0 > ctm_args.fpcm_fpt_tol:
            if verbosity>0: log.info(f"fpcm_MOVE_sl_c iteration {i}")
            # 1.1) get C, U by solving pulling-through eq
            #      --Cp--T-- = --U--Cp-- where 0--U--1
            #            |       |                2
            Cp, U= isogauge_MPS(T, C0=tmp_[-1][1], isogauge_tol=ctm_args.fpcm_isogauge_tol,\
                verbosity=verbosity)
            Tp= fp_T(a, U, T0=tmp_[-1][0])
            
            # (optional) check symmetry error on Tp
            Tp= 0.5*(Tp + Tp.permute(1,0,2))
            
            # 1.2) check convergence 
            tmp_.append((Tp,Cp))
            if len(tmp_)>1:
                e0= torch.norm(tmp_[1][0]-tmp_[0][0])/max(Tp.size())
                if verbosity>0: log.info(f"fpcm_MOVE_sl_c iteration {i} error (p=2) {e0}"\
                    +f" (p=inf) {torch.norm(tmp_[1][0]-tmp_[0][0],float('inf'))}")
                tmp_.pop(0)
            i+=1

        # 2) compute fixed point of the corner
        nC= fp_C(a, Tp, U, verbosity=verbosity)
        
        # 4) symmetrize, normalize and assign new C,T
        nC= 0.5*(nC + nC.t())
        nC= nC/torch.max(torch.abs(nC))
        # nC= nC/nC.norm()
        Tp= Tp/torch.max(torch.abs(Tp))
        # nT= nT/nT.norm()

        return nC, Tp

    # Call the core function, allowing for checkpointing
    if ctm_args.fwd_checkpoint_move:
        new_tensors= checkpoint(fpcm_MOVE_sl_c,*tensors)
    else:
        new_tensors= fpcm_MOVE_sl_c(*tensors)

    env.C[env.keyC]= new_tensors[0]
    env.T[env.keyT]= new_tensors[1]