Project: s41467-025-56078-4 | Function: compute_snip_per_weight (python)

system: You are an expert Python programmer implementing algorithms from scientific papers.
Complete the masked function based on the algorithm description.
Follow Python best practices and maintain consistent code style.
Return **only** the complete Python function code.

user: Complete this Python function:

**DESCRIPTION:** Implements the core SNIP (Single-shot Network Pruning) zero-shot neural architecture search algorithm that computes connection sensitivity for weakly physics-informed network adaptation without requiring target domain data, enabling robust domain-shift-resistant parameter selection

**KEY TERMS:** {
  "weight_mask": "binary mask m for connection sensitivity analysis",
  "weight_mask.grad": "gradient ∇_θ L w.r.t. mask parameters",
  "snip_forward": "masked forward pass f(x; θ ⊙ m)",
  "torch.abs": "absolute gradient magnitude |∇_θ L|",
  "zero_shot": "no target domain data required"
}

**CONTEXT CODE:**             def hook_factory(layer):
                def hook(module, grad_input, grad_output):
                    act = layer.act.detach()
                    grad = grad_output[0].detach()
                    if len(act.shape) > 2:
                        g_nk = torch.sum((act * grad), list(range(2,len(act.shape))))
                    else:
                        g_nk = act * grad
                    del_k = g_nk.pow(2).mean(0).mul(0.5)
                    if layer.fisher is None:
                        layer.fisher = del_k
                    else:
                        layer.fisher += del_k
                    del layer.act #without deleting this, a nasty memory leak occurs! related: https://discuss.pytorch.org/t/memory-leak-when-using-forward-hook-and-backward-hook-simultaneously/27555
                return hook

            #register backward hook on identity fcn to compute fisher info
            layer.dummy.register_backward_hook(hook_factory(layer))

    N = inputs.shape[0]
    for sp in range(split_data):
        st=sp*N//split_data
        en=(sp+1)*N//split_data

        net.zero_grad()
        outputs = net(inputs[st:en])
        loss = loss_fn(outputs, targets[st:en])
        loss.backward()

    # retrieve fisher info
    def fisher(layer):
        if layer.fisher is not None:
            return torch.abs(layer.fisher.detach())
        else:
            return torch.zeros(layer.weight.shape[0]) #size=ch

    grads_abs_ch = get_layer_metric_array(net, fisher, mode)
    shapes = get_layer_metric_array(net, lambda l : l.weight.shape[1:], mode)

    grads_abs = reshape_elements(grads_abs_ch, shapes, device)

    return grads_abs

def snip_forward_conv2d(self, x):
        return F.conv2d(x, self.weight * self.weight_mask, self.bias,
                        self.stride, self.padding, self.dilation, self.groups)

def snip_forward_linear(self, x):
        return F.linear(x, self.weight * self.weight_mask, self.bias)

def compute_snip_per_weight(net, inputs, targets, loss_fn,  mode = 'param', split_data=1):
    # MASKED: Implements the core SNIP (Single-shot Network Pruning) zero-shot neural architecture search algorithm that computes connection sensitivity for weakly physics-informed network adaptation without requiring target domain data, enabling robust domain-shift-resistant parameter selection
    raise NotImplementedError("This function is masked for testing purposes.")

def compute_grasp_per_weight(net, inputs, targets, loss_fn, mode = 'param',  T=1, num_iters=1, split_data=1):
    weights = []
    for layer in net.modules():
        if isinstance(layer, nn.Conv2d) or isinstance(layer, nn.Linear):
            weights.append(layer.weight)
            layer.weight.requires_grad_(True) # TODO isn't this already true?

    # NOTE original code had some input/target splitting into 2
    # I am guessing this was because of GPU mem limit
    net.zero_grad()
    N = inputs.shape[0]
    # ['input_ids'].
    for sp in range(split_data):
        st=sp*N//split_data
        en=(sp+1)*N//split_data

        #forward/grad pass #1
        grad_w = None
        for _ in range(num_iters):
            #TODO get new data, otherwise num_iters is useless!
            # inputs = inputs[st:en]
            # {k: v[st:en].cuda() for k, v in inputs.items()}
            # inputs[st:en]
            outputs = net.forward(inputs[st:en])/T
            # .logits
            loss = loss_fn(outputs, targets[st:en])
            # loss.backward()
            
            # outputs = net.forward(inputs[st:en])/T
            # # .logits
            # loss = loss_fn(outputs, targets[st:en])
            grad_w_p = autograd.grad(loss, weights, allow_unused=True)
            if grad_w is None:
                grad_w = list(grad_w_p)
            else:
                for idx in range(len(grad_w)):
                    grad_w[idx] += grad_w_p[idx]

    
    for sp in range(split_data):
        st=sp*N//split_data
        en=(sp+1)*N//split_data
        # inputs = inputs[st:en]
        outputs = net.forward(inputs[st:en])/T
        loss = loss_fn(outputs, targets[st:en])

        grad_f = autograd.grad(loss, weights, create_graph=True, allow_unused=True)
        
        # accumulate gradients computed in previous step and call backwards
        z, count = 0,0
        for layer in net.modules():
            if isinstance(layer, nn.Conv2d) or isinstance(layer, nn.Linear):
                if grad_w[count] is not None:
                    z += (grad_w[count].data * grad_f[count]).sum()
                count += 1
        z.backward()

    # compute final sensitivity metric and put in grads
    def grasp(layer):
        if layer.weight.grad is not None:
            return -layer.weight.data * layer.weight.grad   # -theta_q Hg
            #NOTE in the grasp code they take the *bottom* (1-p)% of values
            #but we take the *top* (1-p)%, therefore we remove the -ve sign
            #EDIT accuracy seems to be negatively correlated with this metric, so we add -ve sign here!
        else:
            return torch.zeros_like(layer.weight)
    
    grads = get_layer_metric_array(net, grasp, mode)
    return grads

def compute_plain_per_weight(net, inputs, targets, loss_fn, mode = 'param', split_data=1):
    net.zero_grad()
    N = inputs.shape[0]
    # ['input_ids']
    for sp in range(split_data):
        st=sp*N//split_data
        en=(sp+1)*N//split_data
        # inputs = inputs[st:en]
        # {k: v[st:en].cuda() for k, v in inputs.items()}
        outputs = net.forward( inputs[st:en])
        # .logits
        loss = loss_fn(outputs, targets[st:en])
        loss.backward()

    # select the gradients that we want to use for search/prune
    def plain(layer):
        if layer.weight.grad is not None:
            return layer.weight.grad * layer.weight
        else:
            return torch.zeros_like(layer.weight)

    grads_abs = get_layer_metric_array(net, plain, mode)
    return grads_abs

def compute_synflow_per_weight(net, inputs, targets, mode, split_data=1, loss_fn=None):



Complete the function implementation replacing "raise NotImplementedError".