code: def replace_layernorm(layernorm, masks):
      in_masks, _, _ = masks
      assert isinstance(layernorm, nn.LayerNorm)
      if len(in_masks) != 1:
          raise InputsNumberError()
      in_mask = in_masks[0]
  
      old_normalized_shape = layernorm.normalized_shape
      new_normalized_shape = []
      remained_list = []
      for i in range(-len(old_normalized_shape), 0):
          pruned, remained = convert_to_coarse_mask(in_mask, i)
          new_normalized_shape.append(
              old_normalized_shape[i] - pruned.size()[i]
          )
          remained_list.append(remained)
  
      new_layernorm = nn.LayerNorm(
          tuple(new_normalized_shape),
          layernorm.eps,
          layernorm.elementwise_affine,
      )
      _logger.info(
          f"replace LayerNorm with new normalized_shape: {tuple(new_normalized_shape)}"
      )
  
      if new_layernorm.elementwise_affine:
          new_layernorm.to(layernorm.weight.device)
          # NOTE: should we keep the weight & bias?
          with torch.no_grad():
              tmp_weight_data = layernorm.weight.data
              tmp_bias_data = layernorm.bias.data
              for i, remained in enumerate(remained_list):
                  tmp_weight_data = torch.index_select(
                      tmp_weight_data, i, remained
                  )
                  tmp_bias_data = torch.index_select(
                      tmp_bias_data, i, remained
                  )
              new_layernorm.weight.data = tmp_weight_data
              new_layernorm.bias.data = tmp_bias_data
      return new_layernorm
  
file_path: nni/nni/compression/speedup/replacement.py
imports: import logging
  import math
  import torch
  import torch.nn
  from .error_code import EmptyLayerError, ShapeMisMatchError, InputsNumberError, OutputTypeError, UnBalancedGroupError
label_desc: Use image node_id 13 for calling a global function with example usage: replace_layernorm(layernorm, masks) and returns: new_layernorm