class_docstr: This class is to speedup the model with provided weight mask
docstr: Infer input shape / output shape based on the module's weight mask / input shape / output shape.
  
  For a module:
      Infer its input and output shape from its weight mask
      Infer its output shape from its input shape
      Infer its input shape from its output shape
  
  If its input shape is changed, continue infering its predecessors
  If its output shape is changed, continue infering its successors
  
  Parameters
  ----------
  module_name : str
      The name of the node
  last_module : str
      The name of last visited node
  mask : tensor of mask or ModuleMasks
      Mask of the weights in this node (i.e., module)
  in_shape : ModuleMasks
      Input shape of this node
  out_shape : ModuleMasks
      Output shape of this node
code: def infer_module_mask(
      self,
      module_name,
      last_module,
      mask=None,
      in_shape=None,
      out_shape=None,
  ):
      """
      Infer input shape / output shape based on the module's weight mask / input shape / output shape.
  
      For a module:
          Infer its input and output shape from its weight mask
          Infer its output shape from its input shape
          Infer its input shape from its output shape
  
      If its input shape is changed, continue infering its predecessors
      If its output shape is changed, continue infering its successors
  
      Parameters
      ----------
      module_name : str
          The name of the node
      last_module : str
          The name of last visited node
      mask : tensor of mask or ModuleMasks
          Mask of the weights in this node (i.e., module)
      in_shape : ModuleMasks
          Input shape of this node
      out_shape : ModuleMasks
          Output shape of this node
      """
      input_cmask = output_cmask = None
      if module_name in self.inferred_masks:
          module_masks = self.inferred_masks[module_name]
      else:
          _, m = get_module_by_name(self.bound_model, module_name)
          module_masks = ModuleMasks(module_name, m)
          self.inferred_masks[module_name] = module_masks
  
      m_type = self.torch_graph.name_to_node[module_name].op_type
      _logger.debug(
          "infer mask of module %s with op_type %s", module_name, m_type
      )
      if mask is not None:
          _logger.debug("mask is not None")
          if not m_type in infer_from_mask:
              raise RuntimeError(
                  "Has not supported infering input/output shape from mask for module/function: `{}`, {}".format(
                      m_type, module_name
                  )
              )
          if m_type in ["Linear"]:
              input_cmask, output_cmask = infer_from_mask[m_type](
                  module_masks,
                  mask,
                  self.torch_graph.name_to_node[module_name].auxiliary,
              )
          else:
              input_cmask, output_cmask = infer_from_mask[m_type](
                  module_masks, mask
              )
      if in_shape is not None:
          _logger.debug("in_shape is not None")
          if not m_type in infer_from_inshape:
              raise RuntimeError(
                  "Has not supported infering output shape from input shape for module/function: `{}`, {}".format(
                      m_type, module_name
                  )
              )
          if m_type in [
              "aten::view",
              "aten::flatten",
              "aten::mean",
              "aten::reshape",
          ]:
              output_cmask = infer_from_inshape[m_type](
                  module_masks,
                  in_shape,
                  self.torch_graph.name_to_node[module_name].auxiliary,
              )
          elif m_type in ["aten::cat"]:
              # To calculate the mask for concat operation, the output shape
              # , cat dimension, and the order of the input parameters.
              output_cmask = infer_from_inshape[m_type](
                  module_masks,
                  in_shape,
                  self.torch_graph.name_to_node[module_name].auxiliary,
                  last_module,
              )
          else:
              output_cmask = infer_from_inshape[m_type](
                  module_masks, in_shape
              )
      if out_shape is not None:
          _logger.debug("out_shape is not None")
          if not m_type in infer_from_outshape:
              raise RuntimeError(
                  "Has not supported infering input shape from output shape for module/function: `{}`, {}".format(
                      m_type, module_name
                  )
              )
          if m_type in [
              "aten::view",
              "aten::flatten",
              "aten::mean",
              "aten::reshape",
          ]:
              input_cmask = infer_from_outshape[m_type](
                  module_masks,
                  out_shape,
                  self.torch_graph.name_to_node[module_name].auxiliary,
              )
          else:
              input_cmask = infer_from_outshape[m_type](
                  module_masks, out_shape
              )
  
      if input_cmask:
          predecessors = self.torch_graph.find_predecessors(module_name)
          for _module_name in predecessors:
              self.infer_module_mask(
                  _module_name, module_name, out_shape=input_cmask
              )
      if output_cmask:
          successors = self.torch_graph.find_successors(module_name)
          for _module_name in successors:
              self.infer_module_mask(
                  _module_name, module_name, in_shape=output_cmask
              )
  
file_path: auptimizer/src/aup/compression/torch/speedup/compressor.py
imports: import logging
  import torch
  from ..utils.mask_conflict import fix_mask_conflict
  from ..utils.utils import get_module_by_name
  from .compress_modules import replace_module
  from .infer_shape import ModuleMasks, infer_from_mask, infer_from_inshape, infer_from_outshape, set_conv_prune_dim
label_desc: Use image node_id 2 for calling the ModelSpeedup obj's underlying member method code with example usage: obj.infer_module_mask(module_name, last_module, mask, in_shape, out_shape) without return types