code: def _fill_one_on_dims(
      mask: torch.Tensor, dims: int | List[int]
  ) -> torch.Tensor:
      dims = dims if isinstance(dims, list) else [dims]
      dims = [d if d >= 0 else d + len(mask.shape) for d in dims]
      new_mask = torch.ones_like(mask)
      for i in range(len(mask.shape)):
          if i in dims:
              continue
          dim_mask = (
              mask.sum([_ for _ in range(len(mask.shape)) if _ != i])
              == 0.0
          )
          new_mask = new_mask.transpose(0, i)
          new_mask[
              torch.arange(len(dim_mask), device=new_mask.device)[
                  dim_mask
              ]
              .long()
              .tolist()
          ] = 0.0
          new_mask = new_mask.transpose(0, i)
      return new_mask
  
file_path: nni/nni/compression/utils/external/external_replacer.py
imports: from __future__ import annotations
  from collections import defaultdict
  import logging
  import re
  from typing import TYPE_CHECKING, List
  import torch
  from torch.utils._pytree import tree_map
  from .huggingface import parser_factory, HuggingfaceModelParser
  from ..attr import get_nested_attr
  from ...speedup.replacer import Replacer
label_desc: Use image node_id 4 for calling a global function with example usage: _fill_one_on_dims(mask, dims) and returns: new_mask