Repository: microsoft/nni
File: nni/retiarii/operation_def/torch_op_def.py (Python)

Code before:
        elif self.type == 'aten::Float':
            return f'{output} = float({inputs[0]})'


class AtenIndex(PyTorchOperation):
    _ori_type_name = ['aten::index']

    def to_forward_code(self, field: str, output: str, inputs: List[str], inputs_value: List[Any] = None) -> str:
        return f'{output} = {inputs[0]}[{inputs[1]}]'


ManuallyChooseDef = {
    'aten::flatten': [('start_dim', 'int', '0'), ('end_dim', 'int', '-1')],
    'aten::split': [('split_size', 'int', 'None'), ('dim', 'int', '0')],
    # in v1.9 dtype is supported as input argument for view, but torch script does not support it
    'aten::view': [('size', 'List[int]', 'None')],
    # NOTE: dim supports different types: List[int], List[str], Optional[List[int]], now we only support the first two, refactor needed
    # torch.std(input, dim, unbiased, keepdim=False, *, out=None)  Tensor
    # torch.std(input, unbiased)  Tensor
    'aten::std': [('dim', 'List[int]', 'None'), ('unbiased', 'bool', 'True'), ('keepdim', 'bool', 'False')]
}

TensorOpExceptions = {
    'aten::sub': lambda output, inputs: f'{output} = {inputs[0]} - {inputs[1]}',  # example: x.size(1) - 3
    'aten::add': lambda output, inputs: f'{output} = {inputs[0]} + {inputs[1]}',  # example: input.shape[0] + 5
    'aten::mul': lambda output, inputs: f'{output} = {inputs[0]} * {inputs[1]}'  # exmaple: 3 * 5
}

TorchOpExclude = ['aten::Size', 'aten::as_tensor', 'aten::device',
                  'aten::manual_seed', 'aten::quantized_gru', 'aten::quantized_lstm',
                  'aten::save', 'aten::tensor', 'aten::wait'
                  ]


def _hidden(name):
    return name.startswith('_') and not name.startswith('__')


def _emit_args(args):
    # filter out the `out` argument here
    return [(arg.name, str(arg.type), str(arg.default_value)) for arg in args]  # if arg.name != 'out'


def _get_tensor_ops():
    def is_tensor_method(schema):
        if len(schema.arguments) == 0:
            return False
        self = schema.arguments[0]
        if self.name != 'self':
            return False
        if not self.type.isSubtypeOf(torch._C.TensorType.get()):

Review comment (suggestion):
`aten::mul` is different from `aten::sub` and `aten::add`. you can find it blow:

>>> args["aten::add"]
[[('other', 'Tensor', 'None'), ('alpha', 'number', '1')], [('other', 'number', 'None'), ('alpha', 'number', '1')], [('other', 'Tensor', 'None'), ('alpha', 'number', '1'), ('out', 'Tensor', 'None')]]
>>> args["aten::sub"]
[[('other', 'Tensor', 'None'), ('alpha', 'number', '1')], [('other', 'number', 'None'), ('alpha', 'number', '1')], [('other', 'Tensor', 'None'), ('alpha', 'number', '1'), ('out', 'Tensor', 'None')]]
>>> args["aten::mul"]
[[('other', 'Tensor', 'None')], [('other', 'number', 'None')], [('other', 'Tensor', 'None'), ('out', 'Tensor', 'None')]]

there are more than one additional operands for `aten::add` and `aten::sub`, but there is only one additional operand for `aten::mul`

Code after:
    _ori_type_name = ['aten::IntImplicit', 'aten::Float', 'aten::Int', 'aten::ScalarImplicit']

    def to_forward_code(self, field: str, output: str, inputs: List[str], inputs_value: List[Any] = None) -> str:
        if self.type.endswith('Implicit'):
            return f'{output} = {inputs[0]}'
        elif self.type == 'aten::Int':
            return f'{output} = int({inputs[0]})'
        elif self.type == 'aten::Float':
            return f'{output} = float({inputs[0]})'


class AtenIndex(PyTorchOperation):
    _ori_type_name = ['aten::index']

    def to_forward_code(self, field: str, output: str, inputs: List[str], inputs_value: List[Any] = None) -> str:
        return f'{output} = {inputs[0]}[{inputs[1]}]'


ManuallyChooseDef = {
    'aten::flatten': [('start_dim', 'int', '0'), ('end_dim', 'int', '-1')],
    'aten::split': [('split_size', 'int', 'None'), ('dim', 'int', '0')],
    # in v1.9 dtype is supported as input argument for view, but torch script does not support it
    'aten::view': [('size', 'List[int]', 'None')],
    # NOTE: dim supports different types: List[int], List[str], Optional[List[int]], now we only support the first two, refactor needed
    # torch.std(input, dim, unbiased, keepdim=False, *, out=None)  Tensor
    # torch.std(input, unbiased)  Tensor
    'aten::std': [('dim', 'List[int]', 'None'), ('unbiased', 'bool', 'True'), ('keepdim', 'bool', 'False')]
}

TensorOpExceptions = {
    'aten::sub': lambda output, inputs: f'{output} = {inputs[0]} - {inputs[1]}',  # example: x.size(1) - 3
    'aten::add': lambda output, inputs: f'{output} = {inputs[0]} + {inputs[1]}'  # example: input.shape[0] + 5
}

TorchOpExclude = ['aten::Size', 'aten::as_tensor', 'aten::device',
                  'aten::manual_seed', 'aten::quantized_gru', 'aten::quantized_lstm',
                  'aten::save', 'aten::tensor', 'aten::wait'
                  ]


def _hidden(name):
    return name.startswith('_') and not name.startswith('__')


def _emit_args(args):
    # filter out the `out` argument here
    return [(arg.name, str(arg.type), str(arg.default_value)) for arg in args]  # if arg.name != 'out'


def _get_tensor_ops():
    def is_tensor_method(schema):