code: def merge_scales_split(self, split_count):
      all_scales = [[] for _ in range(split_count)]
      for dense_scale, qkv_scale, m4hh_scale, mh4h_scale in zip(
          self.dense_scales,
          self.qkv_scales,
          self.mlp4hh_scales,
          self.mlph4h_scales,
      ):
          dense_scale = torch.split(
              dense_scale, dense_scale.numel() // split_count
          )
          qkv_scale = torch.split(
              qkv_scale, qkv_scale.numel() // split_count
          )
          m4hh_scale = torch.split(
              m4hh_scale, m4hh_scale.numel() // split_count
          )
          mh4h_scale = torch.split(
              mh4h_scale, mh4h_scale.numel() // split_count
          )
          for s in range(split_count):
              all_scales[s].append(
                  torch.cat(
                      [
                          torch.cat(
                              (
                                  qkv_scale[s],
                                  torch.zeros_like(qkv_scale[s]),
                              ),
                              dim=1,
                          ),
                          torch.cat(
                              (
                                  dense_scale[s],
                                  torch.zeros_like(dense_scale[s]),
                              ),
                              dim=1,
                          ),
                          mh4h_scale[s],
                          m4hh_scale[s],
                      ]
                  ).unsqueeze(0)
              )
          for scales_a in all_scales:
              torch.cat(scales_a)
      return all_scales
  
file_path: H2O/h2o_flexgen/benchmark/third_party/DeepSpeed/deepspeed/runtime/weight_quantizer.py
imports: import torch
  from ..module_inject.replace_policy import HFBertLayerPolicy, replace_policies
label_desc: Use image node_id 8 for calling the WeightQuantization obj's underlying member method code with example usage: obj.merge_scales_split(split_count) and returns: all_scales