from __future__ import annotations import functools import inspect import sys import typing import warnings from typing import ( Any, Callable, List, Literal, NoReturn, Optional, Sequence, Set, Tuple, Union, ) import torch import torch._C._onnx as _C_onnx from torch import _C # Monkey-patch graph manipulation methods on Graph, used for the ONNX symbolics from torch.onnx import _constants, _type_utils, errors from torch.onnx._globals import GLOBALS from torch.onnx._internal import _beartype, jit_utils from torch.types import Number __all__ = [ "args_have_same_dtype", "cast_pytorch_to_onnx", "check_training_mode", "dequantize_helper", "is_caffe2_aten_fallback", "is_complex_value", "parse_args", "pytorch_name_to_type", "quantize_helper", "quantized_args", "requantize_bias_helper", "scalar_name_to_pytorch", "scalar_type_to_onnx", "scalar_type_to_pytorch_type", ] # --------------------------------------------------------------------------------- # Helper functions # --------------------------------------------------------------------------------- _ValueDescriptor = Literal[ "v", "i", "is", "f", "fs", "b", "s", "t", "none", ] @_beartype.beartype def _parse_arg( value, desc: _ValueDescriptor, arg_name: Optional[str] = None, node_name: Optional[str] = None, ): if desc == "none": return value if desc == "v" or not _is_value(value): return value node = value.node() if node.mustBeNone(): return None if node.kind() == "onnx::Constant": node_val = _node_get(node, "value") if desc == "i": return int(node_val) elif desc == "f": return float(node_val) elif desc == "b": return bool(node_val) elif desc == "s": return str(node_val) elif desc == "t": return node_val elif desc == "is": return [int(v) for v in node_val] elif desc == "fs": return [float(v) for v in node_val] else: raise errors.SymbolicValueError( f"ONNX symbolic does not understand the Constant node '{node}' " f"specified with descriptor '{desc}'.", value, ) elif node.kind() == "prim::ListConstruct": if desc == "is": for v in node.inputs(): element_node = v.node() if element_node.kind() != "onnx::Constant": raise errors.SymbolicValueError( f"Failed to export a node '{element_node}' " f"(in list node {node}) " f"because it is not constant. " f"Please try to make things (e.g. kernel sizes) static if possible.", value, ) return [int(_node_get(v.node(), "value")) for v in value.node().inputs()] else: raise errors.SymbolicValueError( f"ONNX symbolic does not know how to unpack the ListConstruct node that " f"is not a list of integers: '{node}'", value, ) if arg_name is None or node_name is None: raise errors.SymbolicValueError( f"Expected node type 'onnx::Constant', got '{node.kind()}'.", value, ) raise errors.SymbolicValueError( "Expected node type 'onnx::Constant' " f"for argument '{arg_name}' of node '{node_name}', got '{node.kind()}'.", value, ) @_beartype.beartype def _node_get(node: _C.Node, key: str): """Gets attributes of a node which is polymorphic over return type.""" assert isinstance(node, _C.Node) sel = node.kindOf(key) return getattr(node, sel)(key) @_beartype.beartype def _is_onnx_constant(value: _C.Value): """Whether a Value is an ONNX constant.""" return value.node().kind() == "onnx::Constant" @_beartype.beartype def _maybe_get_const( value: Optional[Union[_C.Value, torch.Tensor, Number, Sequence]], descriptor: _ValueDescriptor, ): # NOTE: prim::Constant at this stage usually means something not compatible in ONNX, # otherwise it'd be converted to onnx::Constant # TODO(justinchuby): Replace insinstance with _is_value once we figure out mypy if isinstance(value, _C.Value) and _is_onnx_constant(value): return _parse_arg(value, descriptor) return value @_beartype.beartype def _maybe_get_scalar(value): value_t = _maybe_get_const(value, "t") if isinstance(value_t, torch.Tensor) and value_t.shape == (): return value_t return value @_beartype.beartype def _get_const(value, desc, arg_name): if not _is_constant(value): raise errors.SymbolicValueError( f"ONNX symbolic expected a constant value of the '{arg_name}' argument, " f"got '{value}'", value, ) return _parse_arg(value, desc) @_beartype.beartype def _unpack_list(list_value: _C.Value) -> List[_C.Value]: list_node = list_value.node() if list_node.kind() != "prim::ListConstruct": raise errors.SymbolicValueError( f"ONNX symbolic expected node type prim::ListConstruct, " f"got '{list_node}'.", list_value, ) return list(list_node.inputs()) @_beartype.beartype def _unpack_tuple(tuple_value: _C.Value) -> Tuple[_C.Value, ...]: tuple_node = tuple_value.node() if not _is_tuple_construct(tuple_value): raise errors.SymbolicValueError( f"ONNX symbolic expected node type 'prim::TupleConstruct', " f"got '{tuple_node.kind()}'.", tuple_value, ) return tuple(tuple_node.inputs()) @_beartype.beartype def _unpack_quantized_tensor(tuple_value: _C.Value) -> Tuple[_C.Value, ...]: """Unpacks a quantized tensor into a tuple of tensor and scale/zero_point. Args: tuple_value: A tuple of tensor, scale, zero_point, and optionally axis. Returns: A tuple of tensor, scale, zero_point, and optionally axis. """ tuple_node = tuple_value.node() # A quantized tensor is represented as tuple of the form (tensor, scale, zero_point, ) if not _is_tuple_construct(tuple_value): raise errors.SymbolicValueError( f"ONNX symbolic expected the output of `{tuple_node}` to be a quantized " f"tensor. Is this likely due to missing support for quantized " f"`{tuple_node.kind()}`. Please create an issue on {_constants.PYTORCH_GITHUB_ISSUES_URL}", tuple_value, ) unpacked = tuple(tuple_node.inputs()) assert len(unpacked) == 3 or len(unpacked) == 4 return unpacked # Check if list_value is output from prim::ListConstruct # This is usually called before _unpack_list to ensure the list can be unpacked. @_beartype.beartype def _is_packed_list(list_value: Any) -> bool: return _is_value(list_value) and list_value.node().kind() == "prim::ListConstruct" @_beartype.beartype def parse_args(*arg_descriptors: _ValueDescriptor): """A decorator which converts args from torch._C.Value to built-in types. For example: ``` @parse_args('v', 'i', 'fs') foo(g, a, b, c): assert isinstance(a, torch._C.Value) assert isinstance(b, int) assert isinstance(c, list) assert isinstance(c[0], float) ``` Args: arg_descriptors: list of str, where each element is a string that specifies the type to convert to. Valid descriptors: "v": no conversion, keep torch._C.Value. "i": int "is": list of int "f": float "fs": list of float "b": bool "s": str "t": torch.Tensor "none": the variable is unused """ def decorator(fn): fn._arg_descriptors = arg_descriptors @functools.wraps(fn) def wrapper(g, *args, **kwargs): # some args may be optional, so the length may be smaller FILE_BUG_MSG = ( "If you believe this is not due to custom symbolic implementation within your code or " "an external library, please file an issue at " "https://github.com/pytorch/pytorch/issues/new?template=bug-report.yml to report this bug." ) assert len(arg_descriptors) >= len(args), ( f"A mismatch between the number of arguments ({len(args)}) and " f"their descriptors ({len(arg_descriptors)}) was found at symbolic function '{fn.__name__}'. " f"{FILE_BUG_MSG}" ) try: sig = inspect.signature(fn) arg_names = list(sig.parameters.keys())[1:] fn_name = fn.__name__ except Exception: # FIXME(justinchuby): Avoid catching Exception. # Catch a more specific exception instead. arg_names = [None] * len(args) # type: ignore[list-item] fn_name = None args = [ _parse_arg(arg, arg_desc, arg_name, fn_name) # type: ignore[method-assign] for arg, arg_desc, arg_name in zip(args, arg_descriptors, arg_names) ] # only support _outputs in kwargs assert len(kwargs) <= 1, ( f"Symbolic function {fn.__name__}'s '**kwargs' can contain a single " f"key/value entry. " f"{FILE_BUG_MSG}" ) if len(kwargs) == 1: assert "_outputs" in kwargs, ( f"Symbolic function {fn.__name__}'s '**kwargs' can only contain " f"'_outputs' key at '**kwargs'. " f"{FILE_BUG_MSG}" ) return fn(g, *args, **kwargs) return wrapper return decorator @_beartype.beartype def quantized_args( *arg_q_descriptors: bool, scale: Optional[float] = None, zero_point: Optional[int] = None, quantize_output: bool = True, ): """A decorator which extends support for quantized version of the base operator. Quantization is detected by examining the arguments that are annotated by `arg_q_descriptors`. If quantization is detected, the base operator symbolic function will be wrapped with argument de-quantization and output quantization. Otherwise, only the base symbolic function will be invoked. For example: ``` @quantized_args(True, False) def foo(g, x, y): return x + y ``` is equivalent to ``` def q_foo(g, x, y): if is_quantized_tensor(x): x = dequantize(x) out = foo(g, x, y) return quantize(out) else: return foo(g, x, y) ``` Args: arg_q_descriptors: A sequence of bool, where each element represents if the argument is QTensor for quantized version of this operator. It defaults to False for unspecified (variable length) arguments. scale: Quantized output scale. If None, derive from the first quantized input scale. zero_point: Quantized output zero point. If None, derive from the first quantized input zero point. quantize_output: If True, quantize the output of the base operator. Default is True """ def decorator(fn): @functools.wraps(fn) def wrapper(g, *args, **kwargs): nonlocal scale nonlocal zero_point if scale is not None: _scale = g.op("Constant", value_t=torch.tensor(scale)) else: _scale = None if zero_point is not None: _zero_point = g.op("Constant", value_t=torch.tensor(zero_point)) else: _zero_point = None # Support variable length arguments by marking unspecified ones as non-quantized arg_q_descriptors_extended = arg_q_descriptors + (False,) * ( len(args) - len(arg_q_descriptors) ) descriptor_args = tuple(zip(arg_q_descriptors_extended, args)) def _is_arg_quantized(descriptor, arg): return descriptor and _is_value(arg) and _is_tuple_construct(arg) # Run regular symbolic function if none of the argument is QTensor. is_quantized = list() for descriptor, arg in descriptor_args: # ListConstruct if _is_packed_list(arg): for arg_input in arg.node().inputs(): is_quantized.append(_is_arg_quantized(descriptor, arg_input)) else: is_quantized.append(_is_arg_quantized(descriptor, arg)) if not any(is_quantized): return fn(g, *args, **kwargs) # Dequantize arguments that are quantized non_quantized_args = [] for descriptor, arg in descriptor_args: if _is_arg_quantized(descriptor, arg): # Quantized arg is a tuple of (value, scale, zero_point) dequantized_arg, arg_scale, arg_zero_point, _ = dequantize_helper( g, arg ) non_quantized_args.append(dequantized_arg) # Set scale and zero_point to the first quantized input if not already set if _scale is None: _scale = arg_scale if _zero_point is None: _zero_point = arg_zero_point # ListConstruct elif _is_packed_list(arg): for arg_input in arg.node().inputs(): if _is_arg_quantized(descriptor, arg_input): # Quantized arg is a tuple of (value, scale, zero_point) ( dequantized_arg, arg_scale, arg_zero_point, _, ) = dequantize_helper(g, arg_input) # Set scale and zero_point to the first quantized input if not already set if _scale is None: _scale = arg_scale if _zero_point is None: _zero_point = arg_zero_point arg_input.replaceAllUsesWith(dequantized_arg) non_quantized_args.append(arg) else: # Non-quantized arg non_quantized_args.append(arg) # TODO(justinchuby): Only single output is supported for now. We may want to # support multiple outputs in the future. output = fn(g, *non_quantized_args, **kwargs) assert _scale is not None, "Bug: Scale must be set for quantized operator" assert ( _zero_point is not None ), "Bug: Zero point must be set for quantized operator" if quantize_output: return quantize_helper(g, output, _scale, _zero_point) return output return wrapper return decorator @_beartype.beartype def _scalar(x: Any) -> Optional[Number]: """Convert a scalar tensor into a Python value.""" if isinstance(x, torch.Tensor) and x.shape == (): return x.item() return None @_beartype.beartype def _if_scalar_type_as(self, tensor): """ Convert self into the same type of tensor, as necessary. We only support implicit casting for scalars, so we never actually need to insert an ONNX cast operator here; just fix up the scalar. """ if isinstance(self, _C.Value): return self scalar_type = _type_utils.JitScalarType.from_value( tensor, _type_utils.JitScalarType.UNDEFINED ) if scalar_type != _type_utils.JitScalarType.UNDEFINED: ty = scalar_type.scalar_name().lower() return getattr(self, ty)() return self @_beartype.beartype def _is_none(x: Any) -> bool: return x is None or (x.node().mustBeNone() if isinstance(x, _C.Value) else False) @_beartype.beartype def _is_value(x: Any) -> bool: return isinstance(x, _C.Value) @_beartype.beartype def _is_constant(value: Any) -> bool: return not _is_value(value) or value.node().kind() in { "onnx::Constant", "prim::Constant", } @_beartype.beartype def _is_tensor(x: _C.Value) -> bool: return x.type().isSubtypeOf(_C.TensorType.get()) # Note: _C.JitType is not exposed to Python and cannot be checked in runtime. def _as_list_type(jit_type: _C.JitType) -> Optional[_C.ListType]: if isinstance(jit_type, _C.ListType): return jit_type return None @_beartype.beartype def _is_list(x: _C.Value) -> bool: return _as_list_type(x.type()) is not None @_beartype.beartype def _is_tensor_list(x: _C.Value) -> bool: x_type = _as_list_type(x.type()) if x_type is None: return False return isinstance(x_type.getElementType(), _C.TensorType) @_beartype.beartype def _is_scalar_list(x: _C.Value) -> bool: """Checks if x is a scalar list, for example: List[float], List[int]. Besides checking the type is ListType, we also check if the data type is a valid ONNX data type. """ x_type = _as_list_type(x.type()) if x_type is None: return False scalar_type = _type_utils.JitScalarType.from_value(x) return scalar_type.onnx_compatible() @_beartype.beartype def _is_tuple_construct(x: _C.Value) -> bool: return x.node().kind() == "prim::TupleConstruct" @_beartype.beartype def is_complex_value(x: _C.Value) -> bool: assert _is_value(x) return _type_utils.JitScalarType.from_value( x, _type_utils.JitScalarType.UNDEFINED ) in { _type_utils.JitScalarType.COMPLEX32, _type_utils.JitScalarType.COMPLEX64, _type_utils.JitScalarType.COMPLEX128, } @_beartype.beartype def is_caffe2_aten_fallback() -> bool: return ( GLOBALS.operator_export_type == _C_onnx.OperatorExportTypes.ONNX_ATEN_FALLBACK and _C_onnx._CAFFE2_ATEN_FALLBACK ) @_beartype.beartype def _get_tensor_rank(x: _C.Value) -> Optional[int]: if not _is_tensor(x) or x.type() is None: return None x_type = x.type() x_type = typing.cast(_C.TensorType, x_type) return x_type.dim() @_beartype.beartype def _get_tensor_sizes(x: _C.Value, allow_nonstatic: bool = True): if not _is_tensor(x) or x.type() is None: return None x_type = x.type() x_type = typing.cast(_C.TensorType, x_type) if allow_nonstatic: # Each individual symbol is returned as None. # e.g. [1, "a", "b"] -> [1, None, None] return x_type.varyingSizes() # returns None, if exists any symbol in sizes. # e.g. [1, "a", "b"] -> None return x_type.sizes() @_beartype.beartype def _get_tensor_dim_size(x: _C.Value, dim: int) -> Optional[int]: sizes = _get_tensor_sizes(x) return sizes[dim] if sizes else None @_beartype.beartype def _get_dim_for_cross(x: _C.Value, dim: Optional[int]): if dim == -1: tensor_rank = _get_tensor_rank(x) assert tensor_rank is not None return dim + tensor_rank # If dim is not given, it defaults to the first dimension found with the size 3 if dim is None: sizes = _get_tensor_sizes(x) assert sizes is not None for index, size in enumerate(sizes): if size is not None and size == 3: return index return dim @_beartype.beartype def _unimplemented(op: str, msg: str, value: Optional[_C.Value] = None) -> None: # For BC reasons, the behavior for Caffe2 does not raise exception for unimplemented operators if _C_onnx._CAFFE2_ATEN_FALLBACK: warnings.warn(f"ONNX export failed on {op} because {msg} not supported") elif GLOBALS.operator_export_type == _C_onnx.OperatorExportTypes.ONNX: _onnx_unsupported(f"{op}, {msg}", value) @_beartype.beartype def _onnx_unsupported(op_name: str, value: Optional[_C.Value] = None) -> NoReturn: message = ( f"Unsupported: ONNX export of operator {op_name}. " f"Please feel free to request support or submit a pull request " f"on PyTorch GitHub: {_constants.PYTORCH_GITHUB_ISSUES_URL}" ) if isinstance(value, _C.Value): raise errors.SymbolicValueError( message, value, ) raise errors.OnnxExporterError(message) @_beartype.beartype def _onnx_opset_unsupported( op_name: str, current_opset: int, supported_opset: int, value: Optional[_C.Value] = None, ) -> NoReturn: message = ( f"Unsupported: ONNX export of {op_name} in opset {current_opset}. " f"Please try opset version {supported_opset}." ) if isinstance(value, _C.Value): raise errors.SymbolicValueError( message, value, ) raise errors.OnnxExporterError(message) @_beartype.beartype def _onnx_opset_unsupported_detailed( op_name: str, current_opset: int, supported_opset: int, reason: str, value: Optional[_C.Value] = None, ) -> NoReturn: message = ( f"Unsupported: ONNX export of {op_name} in " f"opset {current_opset}. {reason}. Please try opset version {supported_opset}." ) if isinstance(value, _C.Value): raise errors.SymbolicValueError( message, value, ) raise errors.OnnxExporterError(message) @_beartype.beartype def _block_list_in_opset(name: str): def symbolic_fn(*args, **kwargs): raise errors.OnnxExporterError( f"ONNX export failed on {name}, which is not implemented for opset " f"{GLOBALS.export_onnx_opset_version}. " "Try exporting with other opset versions." ) return symbolic_fn @_beartype.beartype def _try_get_scalar_type(*args) -> Optional[_type_utils.JitScalarType]: for arg in args: scalar_type = _type_utils.JitScalarType.from_value( arg, _type_utils.JitScalarType.UNDEFINED ) if scalar_type != _type_utils.JitScalarType.UNDEFINED: return scalar_type return None @_beartype.beartype def _select_helper(g: jit_utils.GraphContext, self, dim, index, apply_reshape=True): index_const = _maybe_get_scalar(index) index_dim = _get_tensor_rank(index) if not _is_value(index_const): # Index is a constant scalar. Make it a size 1 constant tensor. index = g.op("Constant", value_t=torch.LongTensor([index_const])) elif index_dim is not None and apply_reshape: if index_dim == 0: # Index is a scalar. Reshape it to a size 1 tensor. index = _reshape_helper( g, index, g.op("Constant", value_t=torch.LongTensor([1])) ) index_scalar_type = _type_utils.JitScalarType.from_value( index, _type_utils.JitScalarType.UNDEFINED ) if index_scalar_type not in { _type_utils.JitScalarType.INT64, _type_utils.JitScalarType.INT, }: index = g.op("Cast", index, to_i=_C_onnx.TensorProtoDataType.INT64) return g.op("Gather", self, index, axis_i=dim) @_beartype.beartype def _slice_helper( g: jit_utils.GraphContext, input, axes, starts, ends, steps=None, ): if g.opset <= 9: from torch.onnx.symbolic_opset9 import _slice as _slice9 return _slice9(g, input, axes, starts, ends) else: from torch.onnx.symbolic_opset10 import _slice as _slice10 return _slice10(g, input, axes, starts, ends, steps) @_beartype.beartype def _is_fp(value) -> bool: return _type_utils.JitScalarType.from_value( value, _type_utils.JitScalarType.UNDEFINED ) in { _type_utils.JitScalarType.FLOAT, _type_utils.JitScalarType.DOUBLE, _type_utils.JitScalarType.HALF, _type_utils.JitScalarType.BFLOAT16, } @_beartype.beartype def _is_bool(value) -> bool: return _type_utils.JitScalarType.from_value( value, _type_utils.JitScalarType.UNDEFINED ) in {_type_utils.JitScalarType.BOOL} @_beartype.beartype def _generate_wrapped_number(g: jit_utils.GraphContext, scalar): """Creates a wrapped number based on https://github.com/pytorch/pytorch/issues/9515. A Tensor is a considered a "wrapped number" if it is auto-wrapped from a C++ or Python number type. Integer types are wrapped as 0-dim int64 tensors and floating-point types are wrapped as 0-dim double tensors. The input to this function is constant value. If the data type is a floating point type, it is converted to a 0-dim double tensor, else it is converted to a 0-dim tensor of its original type """ assert not isinstance(scalar, torch.Tensor) if isinstance(scalar, float): return g.op("Constant", value_t=torch.tensor(scalar, dtype=torch.double)) return g.op("Constant", value_t=torch.tensor(scalar)) @_beartype.beartype def _sort_helper(g: jit_utils.GraphContext, input, dim, decending=True, out=None): if out is not None: _unimplemented("Sort", "Out parameter is not supported") shape_ = g.op("Shape", input) dim_size_ = g.op( "Gather", shape_, g.op("Constant", value_t=torch.tensor([dim], dtype=torch.int64)), ) if g.opset <= 10: if not decending: _unimplemented("Sort", "Ascending is not supported") return g.op("TopK", input, dim_size_, axis_i=dim, outputs=2) else: return g.op( "TopK", input, dim_size_, axis_i=dim, largest_i=decending, outputs=2 ) @_beartype.beartype def _topk_helper( g: jit_utils.GraphContext, input, k, dim, largest=True, sorted=False, out=None ): if out is not None: _unimplemented("TopK", "Out parameter is not supported") if not _is_value(k): k = g.op("Constant", value_t=torch.tensor([k], dtype=torch.int64)) else: k = _reshape_helper(g, k, g.op("Constant", value_t=torch.tensor([1]))) if _try_get_scalar_type(k) != _type_utils.JitScalarType.INT64: k = g.op("Cast", k, to_i=_C_onnx.TensorProtoDataType.INT64) if g.opset <= 10: if not largest: _unimplemented("TopK", "Ascending is not supported") return g.op("TopK", input, k, axis_i=dim, outputs=2) else: return g.op( "TopK", input, k, axis_i=dim, largest_i=largest, sorted_i=sorted, outputs=2 ) @_beartype.beartype def _lt_helper(g: jit_utils.GraphContext, input, other): if g.opset <= 8: from torch.onnx.symbolic_opset8 import lt as _lt8 return _lt8(g, input, other) else: from torch.onnx.symbolic_opset9 import lt as _lt9 return _lt9(g, input, other) @_beartype.beartype def _interpolate_warning(interpolate_mode): onnx_op = ( "onnx:Resize" if GLOBALS.export_onnx_opset_version >= 10 else "onnx:Upsample" ) warnings.warn( "You are trying to export the model with " + onnx_op + " for ONNX opset version " "" + str(GLOBALS.export_onnx_opset_version) + ". " "This operator might cause results to not match the expected results by PyTorch.\n" "ONNX's Upsample/Resize operator did not match Pytorch's Interpolation until opset 11. " "Attributes to determine how to transform the input were added in onnx:Resize in opset 11 " "to support Pytorch's behavior (like coordinate_transformation_mode and nearest_mode).\n" "We recommend using opset 11 and above for models using this operator." ) @_beartype.beartype def _unsqueeze_helper(g: jit_utils.GraphContext, input, axes_i): if _is_constant(axes_i[0]): if g.opset >= 13: axes = g.op("Constant", value_t=torch.tensor(axes_i, dtype=torch.long)) return g.op("Unsqueeze", input, axes) return g.op("Unsqueeze", input, axes_i=axes_i) # Tensor type if g.opset < 13: raise errors.SymbolicValueError( "Opset version must be >= 13 for Unsqueeze with dynamic axes.", input ) return g.op("Unsqueeze", input, axes_i[0]) @_beartype.beartype def _squeeze_helper(g: jit_utils.GraphContext, input, axes_i): if _is_constant(axes_i[0]): if g.opset >= 13: axes = g.op("Constant", value_t=torch.tensor(axes_i, dtype=torch.long)) return g.op("Squeeze", input, axes) return g.op("Squeeze", input, axes_i=axes_i) # Tensor type if g.opset < 13: raise errors.SymbolicValueError( "Opset version must be >= 13 for Squeeze with dynamic axes.", input ) axes_t = axes_i[0] axes_rank = _get_tensor_rank(axes_t) assert axes_rank is not None if axes_rank > 1: raise errors.SymbolicValueError( "For Squeeze axses as input, the axes rank must be one in ONNX spec.", input ) elif axes_rank == 0: # The axes is a scalar. Unsqueeze it to a rank 1 tensor. axes_t = _unsqueeze_helper(g, axes_t, [0]) return g.op("Squeeze", input, axes_t) return g.op("Squeeze", input, axes_t) @_beartype.beartype def _reducesum_helper( g: jit_utils.GraphContext, input, axes_i=None, keepdims_i=1, noop_with_empty_axes_i=0, ): keepdims_i = _maybe_get_const(keepdims_i, "i") if g.opset >= 13: if axes_i: if not _is_value(axes_i): axes_i = g.op( "Constant", value_t=torch.tensor(axes_i, dtype=torch.long) ) return g.op( "ReduceSum", input, axes_i, keepdims_i=keepdims_i, noop_with_empty_axes_i=noop_with_empty_axes_i, ) return g.op( "ReduceSum", input, keepdims_i=keepdims_i, noop_with_empty_axes_i=noop_with_empty_axes_i, ) else: return g.op("ReduceSum", input, axes_i=axes_i, keepdims_i=keepdims_i) @_beartype.beartype def _interpolate_size_to_scales(g: jit_utils.GraphContext, input, output_size, dim): output_size = _maybe_get_const(output_size, "is") if _is_value(output_size): offset = 2 offsets = g.op("Constant", value_t=torch.ones(offset, dtype=torch.float32)) dividend = g.op("Cast", output_size, to_i=_C_onnx.TensorProtoDataType.FLOAT) divisor = _slice_helper( g, g.op("Shape", input), axes=[0], ends=[sys.maxsize], starts=[offset] ) divisor = g.op("Cast", divisor, to_i=_C_onnx.TensorProtoDataType.FLOAT) scale_dims = g.op("Div", dividend, divisor) scales = g.op("Concat", offsets, scale_dims, axis_i=0) else: scales_constant = [ 1.0 if i < 2 else float(output_size[-(dim - i)]) / float(input.type().sizes()[-(dim - i)]) for i in range(0, dim) ] scales = g.op( "Constant", value_t=torch.tensor(scales_constant, dtype=torch.float32) ) return scales @_beartype.beartype def _interpolate_get_scales_if_available(g: jit_utils.GraphContext, scales): available_scales = _maybe_get_const(scales[0], "fs") != -1 and not _is_none( scales[0] ) if not available_scales: return None offsets = g.op("Constant", value_t=torch.ones(2, dtype=torch.float32)) scales_list = g.op( "Constant", value_t=torch.tensor(_maybe_get_const(scales[0], "fs")) ) scales = g.op("Concat", offsets, scales_list, axis_i=0) return scales @_beartype.beartype def _get_interpolate_attributes(g: jit_utils.GraphContext, mode, args): if mode == "nearest": align_corners = None scales = args[0:] else: align_corners = args[0] scales = args[1:] scales = _interpolate_get_scales_if_available(g, scales) return scales, align_corners @_beartype.beartype def _interpolate_get_scales(g: jit_utils.GraphContext, scale_factor, dim): offsets = g.op("Constant", value_t=torch.ones(2, dtype=torch.float32)) scale_factor_rank = _get_tensor_rank(scale_factor) if isinstance(scale_factor.type(), _C.ListType) or ( scale_factor_rank is not None and scale_factor_rank > 0 ): return g.op("Concat", offsets, scale_factor, axis_i=0) else: scale_factor = _unsqueeze_helper(g, scale_factor, [0]) scale_factor = g.op( "Cast", scale_factor, to_i=_C_onnx.TensorProtoDataType.FLOAT ) scales = [scale_factor for i in range(dim - 2)] scale_factor = g.op("Concat", offsets, *scales, axis_i=0) return scale_factor @_beartype.beartype def _interpolate_get_scales_and_mode( g: jit_utils.GraphContext, input, size, scale_factor, mode, align_corners ): mode = _maybe_get_const(mode, "s") if "linear" in mode: mode = "linear" if "cubic" in mode: mode = "cubic" _interpolate_warning(mode) align_corners = _maybe_get_const(align_corners, "b") if isinstance(align_corners, bool) and align_corners: return _unimplemented("interpolate", "align_corners == True") if not input.type().dim(): return _unimplemented("interpolate", "missing input shape") dim = input.type().dim() if not _is_none(scale_factor): scale_factor = _interpolate_get_scales(g, scale_factor, dim) elif not _is_none(size): if not _is_packed_list(size): is_scalar = _maybe_get_const(size, "t").dim() == 0 if is_scalar: size = _unsqueeze_helper(g, size, [0]) size = [size for i in range(dim - 2)] size = g.op("Concat", *size, axis_i=0) scale_factor = _interpolate_size_to_scales(g, input, size, dim) else: return _unimplemented( "interpolate", "Both size and scales are None in __interpolate" ) return scale_factor, mode @_beartype.beartype def _argmin_argmax_helper( g: jit_utils.GraphContext, input: torch._C.Value, dim: torch._C.Value, keepdim: bool, op_name: str, ): def op_wrapper(input, axis_i, keepdims_i): if g.opset >= 12: return g.op( op_name, input, axis_i=axis_i, keepdims_i=keepdims_i, select_last_index_i=False, ) return g.op(op_name, input, axis_i=axis_i, keepdims_i=keepdims_i) if _is_none(dim): flattened = _reshape_helper( g, input, g.op("Constant", value_t=torch.tensor([-1])) ) output = op_wrapper(flattened, axis_i=0, keepdims_i=False) if keepdim: input_shape = g.op("Shape", input) input_shape_shape = g.op("Shape", input_shape) new_shape = g.op( "ConstantOfShape", input_shape_shape, value_t=torch.tensor([1], dtype=torch.int64), ) output = g.op("Reshape", output, new_shape) return output dim = _parse_arg(dim, "i") return op_wrapper(input, axis_i=dim, keepdims_i=keepdim) @_beartype.beartype def _interpolate_helper(name, dim, interpolate_mode): @quantized_args(True, False, False) def symbolic_fn(g, input, output_size, *args): scales, align_corners = _get_interpolate_attributes(g, interpolate_mode, args) align_corners = _maybe_get_scalar(align_corners) coordinate_transformation_mode = ( "asymmetric" if interpolate_mode == "nearest" else "align_corners" if align_corners else "half_pixel" ) if scales is None: input_size = g.op("Shape", input) input_size_beg = _slice_helper( g, input_size, axes=[0], ends=[2], starts=[0] ) output_size = g.op( "Cast", output_size, to_i=_C_onnx.TensorProtoDataType.INT64 ) output_size = g.op("Concat", input_size_beg, output_size, axis_i=0) if g.opset >= 13: empty_roi = _optional_input_placeholder_tensor(g) empty_scales = _optional_input_placeholder_tensor(g) else: empty_roi = g.op( "Constant", value_t=torch.tensor([], dtype=torch.float32) ) empty_scales = g.op( "Constant", value_t=torch.tensor([], dtype=torch.float32) ) return g.op( "Resize", input, empty_roi, empty_scales, output_size, coordinate_transformation_mode_s=coordinate_transformation_mode, cubic_coeff_a_f=-0.75, # only valid when mode="cubic" mode_s=interpolate_mode, # nearest, linear, or cubic nearest_mode_s="floor", ) # only valid when mode="nearest" else: if g.opset >= 13: empty_roi = _optional_input_placeholder_tensor(g) else: empty_roi = g.op( "Constant", value_t=torch.tensor([], dtype=torch.float32) ) return g.op( "Resize", input, empty_roi, scales, coordinate_transformation_mode_s=coordinate_transformation_mode, cubic_coeff_a_f=-0.75, # only valid when mode="cubic" mode_s=interpolate_mode, # nearest, linear, or cubic nearest_mode_s="floor", ) # only valid when mode="nearest" return symbolic_fn @_beartype.beartype def __interpolate_helper( g: jit_utils.GraphContext, input, size, scale_factor, mode, align_corners, recompute_scale_factor, ): mode = _maybe_get_const(mode, "s") if "linear" in mode: mode = "linear" if "cubic" in mode: mode = "cubic" align_corners = _maybe_get_const(align_corners, "b") align_corners = False if not isinstance(align_corners, bool) else align_corners coordinate_transformation_mode = ( "asymmetric" if mode == "nearest" else "align_corners" if align_corners else "half_pixel" ) if not _is_none(size): input_size = g.op("Shape", input) input_size = _slice_helper(g, input_size, axes=[0], ends=[2], starts=[0]) # in some cases size is not a packed list but size is a scalar # We need to also verify that (_maybe_get_const(size, "t").dim() == 0) # but this information is not always available. Try to get the dim, # and if not assume that it is not a scalar. try: is_scalar = not _is_packed_list(size) and ( _maybe_get_const(size, "t").dim() == 0 ) except AttributeError: is_scalar = not _is_packed_list(size) if not is_scalar: warnings.warn( "Cannot verify if the output_size is a scalar " "while exporting interpolate. Assuming that it is not a scalar." ) if is_scalar: rank = _get_tensor_rank(input) if rank is None: return _unimplemented( "interpolate (with a scalar output_size)", "missing input shape (try giving an array of output_size values)", ) size = _unsqueeze_helper(g, size, [0]) size = [size for i in range(rank - 2)] size = g.op("Concat", *size, axis_i=0) size = g.op("Cast", size, to_i=_C_onnx.TensorProtoDataType.INT64) size = g.op("Concat", input_size, size, axis_i=0) if g.opset >= 13: empty_roi = _optional_input_placeholder_tensor(g) empty_scales = _optional_input_placeholder_tensor(g) else: empty_roi = g.op("Constant", value_t=torch.tensor([], dtype=torch.float32)) empty_scales = g.op( "Constant", value_t=torch.tensor([], dtype=torch.float32) ) return g.op( "Resize", input, empty_roi, empty_scales, size, coordinate_transformation_mode_s=coordinate_transformation_mode, cubic_coeff_a_f=-0.75, # only valid when mode="cubic" mode_s=mode, # nearest, linear, or cubic nearest_mode_s="floor", ) else: # if not _is_none(scales) rank = _get_tensor_rank(input) if rank is None: return _unimplemented("interpolate (with scales)", "missing input shape") if g.opset >= 13: empty_roi = _optional_input_placeholder_tensor(g) else: empty_roi = g.op("Constant", value_t=torch.tensor([], dtype=torch.float32)) scales = _interpolate_get_scales(g, scale_factor, rank) return g.op( "Resize", input, empty_roi, scales, coordinate_transformation_mode_s=coordinate_transformation_mode, cubic_coeff_a_f=-0.75, # only valid when mode="cubic" mode_s=mode, # nearest, linear, or cubic nearest_mode_s="floor", ) # only valid when mode="nearest" @_beartype.beartype def _unbind_helper(g: jit_utils.GraphContext, self, dim, _outputs): if g.opset < 11: from torch.onnx.symbolic_opset9 import unbind elif g.opset <= 12: from torch.onnx.symbolic_opset11 import unbind # type: ignore[no-redef] else: from torch.onnx.symbolic_opset13 import unbind # type: ignore[no-redef] return unbind(g, self, dim, _outputs) @_beartype.beartype def _scatter_helper(g: jit_utils.GraphContext, self, dim, index, src): if g.opset <= 10: from torch.onnx.symbolic_opset9 import scatter else: # for mypy, scatter was imported two lines above from torch.onnx.symbolic_opset11 import scatter # type: ignore[no-redef] return scatter(g, self, dim, index, src) @_beartype.beartype def _repeat_interleave_split_helper(g: jit_utils.GraphContext, self, reps, dim): if g.opset <= 12: split_out = g.op("Split", self, split_i=[1] * reps, axis_i=dim, outputs=reps) else: from torch.onnx.symbolic_opset13 import split repeats = g.op("Constant", value_t=torch.tensor([1] * reps)) split_out = split(g, self, repeats, dim, _outputs=reps) return split_out if reps > 1 else [split_out] @_beartype.beartype def _repeat_interleave_single_value_repeat_helper( g: jit_utils.GraphContext, self, repeats, dim ): from torch.onnx.symbolic_opset9 import flatten, unsqueeze if not _is_tensor(repeats): repeats = g.op("Constant", value_t=torch.LongTensor(repeats)) const_repeats: bool = _is_constant(repeats) reps = _maybe_get_const(repeats, "t") # Convert 'repeats' to 1-d if it is 0-d. if _get_tensor_rank(repeats) == 0: repeats = g.op("Reshape", repeats, g.op("Constant", value_t=torch.tensor([1]))) # Create a new dim of size 1, then expand it to be 'repeats' long, and finally collapse it. unsqueezed = unsqueeze(g, self, dim + 1) # repeats_per_dim is 1 for all dims except for the new unsqueezed dim, where it has value 'repeats'. if const_repeats: # 'Repeats' is a constant, 'repeats_per_dim' can be a constant. onehot = torch.ones(_get_tensor_rank(unsqueezed), dtype=torch.int64) onehot[dim + 1] = reps repeats_per_dim = g.op("Constant", value_t=onehot) else: # 'Repeats' is a variable, 'repeats_per_dim' cannot be a constant. onehot = g.op( "OneHot", unsqueeze(g, dim + 1, 0), # indices, must be >= 1-dimensional g.op( "Constant", value_t=torch.tensor(_get_tensor_rank(unsqueezed)) ), # depth g.op( "Concat", g.op("Constant", value_t=torch.tensor([1])), repeats, axis_i=0 ), # on/off values ) repeats_per_dim = flatten(g, onehot, 0, 1) tiled = g.op("Tile", unsqueezed, repeats_per_dim) return flatten(g, tiled, dim, dim + 1) @_beartype.beartype def _arange_cast_helper( g: jit_utils.GraphContext, end, start=None, step=None, dtype=None ) -> Tuple[ _type_utils.JitScalarType, Optional[_C.Value], Optional[_C.Value], Optional[_C.Value], ]: def _is_all_integral(scalars): for scalar in scalars: scalar_type = _type_utils.JitScalarType.from_value( scalar, _type_utils.JitScalarType.UNDEFINED ) if ( scalar_type != _type_utils.JitScalarType.INT64 and scalar_type != _type_utils.JitScalarType.UNDEFINED ): return False return True # This logic is based on torch.arange docs. If "dtype" is provided, # infer input types from dtype. If not, then check if any of start, stop, # or step are floating point, and infer the type from get_default. # Otherwise, the dtype is inferred to be torch.int64. if dtype is None or (_is_value(dtype) and _is_none(dtype)): if _is_all_integral([start, end, step]): scalar_type = _type_utils.JitScalarType.INT64 else: scalar_type = _type_utils.JitScalarType.from_dtype( torch.get_default_dtype() ) else: assert isinstance(dtype, int) # TODO(justinchuby): Check if dtype is indeed a int. scalar_type = _type_utils.JitScalarType(dtype) start = g.op("Cast", start, to_i=scalar_type.onnx_type()) if start else None end = g.op("Cast", end, to_i=scalar_type.onnx_type()) if end else None step = g.op("Cast", step, to_i=scalar_type.onnx_type()) if step else None return scalar_type, end, start, step @_beartype.beartype def _arange_helper(g: jit_utils.GraphContext, *args): if g.opset <= 10: from torch.onnx.symbolic_opset9 import arange else: from torch.onnx.symbolic_opset11 import arange # type: ignore[no-redef] return arange(g, *args) @_beartype.beartype def _size_helper(g: jit_utils.GraphContext, self, dim): full_shape = g.op("Shape", self) from torch.onnx.symbolic_opset9 import select return select(g, full_shape, g.op("Constant", value_t=torch.tensor([0])), dim) @_beartype.beartype def _index_fill_reshape_helper(g: jit_utils.GraphContext, self, dim, index): # 1. reshape index => [1, ..., 1, dim, 1, ..., 1] # 2. expand index => [..., dim, ...], same shape as self except for dim. # 3. expand value as well. # 4. apply onnx::scatter. from torch.onnx.symbolic_opset9 import expand if g.opset <= 10: from torch.onnx.symbolic_opset9 import scatter else: # for mypy, scatter was imported two lines above from torch.onnx.symbolic_opset11 import scatter # type: ignore[no-redef] if self.type().dim() is None: return _unimplemented("index_fill", "input rank not accessible") self_dim = self.type().dim() dim_value = _parse_arg(dim, "i") if dim_value < 0: dim_value += self_dim unsqueezed_index = _unsqueeze_helper( g, index, [i for i in range(self_dim) if i != dim_value] ) expanded_index_shape = scatter( g, g.op("Shape", self), 0, _unsqueeze_helper(g, dim, [0]), g.op("Shape", index) ) expanded_index = expand(g, unsqueezed_index, expanded_index_shape, None) return expanded_index_shape, expanded_index # By default, when any value in the 'shape' input is equal to zero # the corresponding dimension value is copied from the input tensor dynamically. # allowzero=1 indicates that if any value in the 'shape' input is set to zero, # the zero value is honored, similar to NumPy. # allowzero=1 is only supported for opset version >= 14. @_beartype.beartype def _reshape_helper(g: jit_utils.GraphContext, input, shape, allowzero=0): shape = _maybe_get_const(shape, "is") if not _is_value(shape): shape = g.op("Constant", value_t=torch.LongTensor(shape)) if g.opset <= 13: if allowzero == 1: _onnx_opset_unsupported( "Reshape with allowzero=1", GLOBALS.export_onnx_opset_version, 14, input ) return g.op("Reshape", input, shape) else: return g.op("Reshape", input, shape, allowzero_i=allowzero) @_beartype.beartype def _batchnorm_helper( g: jit_utils.GraphContext, input, weight, bias, running_mean, running_var ): from torch.onnx.symbolic_opset9 import _var_mean batch_size = _get_tensor_dim_size(input, 0) channel_size = _get_tensor_dim_size(input, 1) if weight is None or _is_none(weight): if channel_size is None: raise errors.SymbolicValueError( "Unsupported: ONNX export of batch_norm for unknown channel size.", input, ) weight_value = torch.tensor( [1.0] * channel_size, dtype=_type_utils.JitScalarType.from_value(input).dtype(), ) weight = g.op("Constant", value_t=weight_value) if bias is None or _is_none(bias): if channel_size is None: raise errors.SymbolicValueError( "Unsupported: ONNX export of batch_norm for unknown channel size.", input, ) bias_value = torch.tensor( [0.0] * channel_size, dtype=_type_utils.JitScalarType.from_value(input).dtype(), ) bias = g.op("Constant", value_t=bias_value) # If track_running_stats is set to False batch statistics are instead used during evaluation time if ( running_mean is None or _is_none(running_mean) or running_var is None or _is_none(running_var) ): assert batch_size is not None and channel_size is not None reshape_in = _reshape_helper( g, input, g.op( "Constant", value_t=torch.tensor([batch_size, channel_size, -1], dtype=torch.int64), ), ) trans_in = g.op("Transpose", reshape_in, perm_i=[0, 2, 1]) running_var, running_mean = _var_mean( g, trans_in, g.op("Constant", value_t=torch.tensor([0, 1], dtype=torch.int64)), False, False, ) return weight, bias, running_mean, running_var @_beartype.beartype def _avgpool_helper( tuple_fn: Callable[[Any], Sequence[int]], padding: Union[int, Sequence[int]], kernel_size, stride, divisor_override, name, ) -> Tuple[int, ...]: if divisor_override and divisor_override.node().kind() != "prim::Constant": _unimplemented(name, "divisor_override") return tuple(tuple_fn(padding)) @_beartype.beartype def check_training_mode(op_train_mode: int, op_name: str) -> None: """Warns the user if the model's training mode and the export mode do not agree.""" if GLOBALS.training_mode == _C_onnx.TrainingMode.PRESERVE: return if op_train_mode: op_mode_enum = _C_onnx.TrainingMode.TRAINING else: op_mode_enum = _C_onnx.TrainingMode.EVAL if op_mode_enum == GLOBALS.training_mode: # The modes agree. Do nothing return op_mode_text = f"train={bool(op_train_mode)}" # Setting the model mode could result in op_mode != GLOBALS.training_mode # if the model is a FuncModule. In this case we warn the user of # the state and export depending on op_mode # This is to support use-cases of fixing certain layer weights # in training. warnings.warn( f"ONNX export mode is set to {GLOBALS.training_mode}, but operator '{op_name}' " f"is set to {op_mode_text}. Exporting with {op_mode_text}." ) @_beartype.beartype def _flatten_helper(g: jit_utils.GraphContext, input, start_dim, end_dim, dim): input_size = g.op("Shape", input) slice1 = _slice_helper(g, input_size, axes=[0], starts=[0], ends=[start_dim]) slices = [slice1, g.op("Constant", value_t=torch.tensor([-1], dtype=torch.long))] if end_dim < dim - 1: slice3 = _slice_helper( g, input_size, axes=[0], starts=[end_dim + 1], ends=[dim] ) slices = [ slice1, g.op("Constant", value_t=torch.tensor([-1], dtype=torch.long)), slice3, ] final_shape = g.op("Concat", *slices, axis_i=0) from torch.onnx.symbolic_opset9 import _reshape_from_tensor return _reshape_from_tensor(g, input, final_shape) @_beartype.beartype def _is_split_static(split_size_or_sizes, _outputs): if _outputs is None: return False if ( _is_value(split_size_or_sizes) and split_size_or_sizes.node().kind() != "onnx::Constant" ): return False return True @_beartype.beartype def _optional_input_placeholder_tensor(g): n = g.op("prim::Constant") n.setType(_C.OptionalType.ofTensor()) return n @_beartype.beartype def _handle_reduce_dim_none(g: jit_utils.GraphContext, self, op_name): rank = _get_tensor_rank(self) if rank is not None and any( _get_tensor_dim_size(self, i) == 0 for i in range(rank) ): # If input tensor is empty, according to ONNX ReduceSum definition, # set keepdims=1 so that the resulted tensor has the same rank as the input. return g.op(op_name, self, keepdims_i=1) return g.op(op_name, self, keepdims_i=0) @_beartype.beartype def dequantize_helper( g: jit_utils.GraphContext, qtensor: _C.Value, qdtype: Optional[_C_onnx.TensorProtoDataType] = None, ) -> Tuple[_C.Value, _C.Value, _C.Value, Optional[_C.Value]]: """Appends to graph `g` ONNX nodes that dequantizes `qtensor` into `tensor`. Args: g: Graph, the ONNX IR graph that is under construction. qtensor: torch._C.Value, either a tuple of (quantized_tensor, scale, zero_point) for per tensor quantization, or (quantized_tensor, scale, zero_point, axis) for per channel quantization, representing the quantized tensor. qdtype: torch.onnx.TensorProtoDataType default None, if not None, represents the data type of quantized tensor. It must be either torch.onnx.TensorProtoDataType.UINT8 or torch.onnx.TensorProtoDataType.INT8. """ unpacked_qtensors = _unpack_quantized_tensor(qtensor) tensor, scale, zero_point = unpacked_qtensors[:3] axis = unpacked_qtensors[3] if len(unpacked_qtensors) >= 4 else None axis_i = _get_const(axis, "i", "axis") input_qdtype = _type_utils.JitScalarType.from_value(tensor) if qdtype is None: if input_qdtype is not None: qdtype = input_qdtype.onnx_type() else: qdtype = _C_onnx.TensorProtoDataType.UINT8 value = g.op("Cast", tensor, to_i=qdtype) scale = g.op("Cast", scale, to_i=_C_onnx.TensorProtoDataType.FLOAT) zero_point = g.op("Cast", zero_point, to_i=qdtype) if axis_i is not None and GLOBALS.export_onnx_opset_version < 13: _onnx_opset_unsupported_detailed( "DequantizeLinear", GLOBALS.export_onnx_opset_version, 13, "Attribute axis is not supported.", qtensor, ) return ( g.op("DequantizeLinear", value, scale, zero_point, axis_i=axis_i), scale, zero_point, axis, ) @_beartype.beartype def quantize_helper( g: jit_utils.GraphContext, tensor: _C.Value, scale: _C.Value, zero_point: _C.Value, axis: Optional[_C.Value] = None, ) -> _C.Value: """Appends to graph `g` ONNX nodes that quantizes `tensor` based on `scale`, `zero_point` and `axis`. Args: g: Graph, the ONNX IR graph that is under construction. tensor: torch._C.Value, representing the tensor to be quantized. scale: torch._C.Value, quantized scale. zero_point: torch._C.Value, quantized zero point. axis: Optional[torch._C.Value] default None, if None, represents per tensor quantization. Otherwise, represents per channel quantization, along given axis. Returns: A TupleConstruct storing information of the quantized tensor. """ if ( axis is not None and not _is_none(axis) and GLOBALS.export_onnx_opset_version < 13 ): _onnx_opset_unsupported_detailed( "QuantizeLinear", GLOBALS.export_onnx_opset_version, 13, "Attribute axis is not supported.", tensor, ) assert scale is not None if ( _type_utils.JitScalarType.from_value(scale, _type_utils.JitScalarType.UNDEFINED) != _type_utils.JitScalarType.FLOAT ): scale = g.op("Cast", scale, to_i=_C_onnx.TensorProtoDataType.FLOAT) assert zero_point is not None if _type_utils.JitScalarType.from_value( zero_point, _type_utils.JitScalarType.UNDEFINED ) not in { _type_utils.JitScalarType.UINT8, _type_utils.JitScalarType.INT8, }: zero_point = g.op("Cast", zero_point, to_i=_C_onnx.TensorProtoDataType.UINT8) output = g.op( "QuantizeLinear", tensor, scale, zero_point, axis_i=_get_const(axis, "i", "axis"), ) args = [output, scale, zero_point] if axis is not None and not _is_none(axis): args.append(axis) return g.op("prim::TupleConstruct", *args) @_beartype.beartype def requantize_bias_helper( g: jit_utils.GraphContext, bias, input_scale, weight_scale, axis=None ): """In PyTorch, bias is float and is quantized to int32 implicitly inside the quantized ATen op kernel. In ONNX we need to make the quantization explicit because operators expect all of their inputs to be quantized. Since int32 is not a supported output type by ONNX operator `QuantizeLinear`, quantization is exported using regular operators. """ bias_scale = g.op("Mul", weight_scale, input_scale) bias_scale_shape = g.op("Shape", bias_scale) bias_zero_point = g.op( "ConstantOfShape", bias_scale_shape, value_t=torch.tensor([0], dtype=torch.int) ) q_bias = g.op( "Cast", g.op("Div", bias, bias_scale), to_i=_C_onnx.TensorProtoDataType.INT32 ) axis_args = [] if axis is not None and not _is_none(axis): axis_args.append(axis) return g.op("prim::TupleConstruct", q_bias, bias_scale, bias_zero_point, *axis_args) @_beartype.beartype def args_have_same_dtype(args): assert args base_dtype = _type_utils.JitScalarType.from_value(args[0]) has_same_dtype = all( _type_utils.JitScalarType.from_value(elem) == base_dtype for elem in args ) return has_same_dtype # Deprecated. Internally use _type_utils.ScalarType # TODO: remove these once we support Type's in the JIT IR and we can once again # use the unified toType operator cast_pytorch_to_onnx = { "Byte": _C_onnx.TensorProtoDataType.UINT8, "Char": _C_onnx.TensorProtoDataType.INT8, "Double": _C_onnx.TensorProtoDataType.DOUBLE, "Float": _C_onnx.TensorProtoDataType.FLOAT, "Half": _C_onnx.TensorProtoDataType.FLOAT16, "Int": _C_onnx.TensorProtoDataType.INT32, "Long": _C_onnx.TensorProtoDataType.INT64, "Short": _C_onnx.TensorProtoDataType.INT16, "Bool": _C_onnx.TensorProtoDataType.BOOL, "ComplexFloat": _C_onnx.TensorProtoDataType.COMPLEX64, "ComplexDouble": _C_onnx.TensorProtoDataType.COMPLEX128, "BFloat16": _C_onnx.TensorProtoDataType.BFLOAT16, "Undefined": _C_onnx.TensorProtoDataType.UNDEFINED, } # Deprecated. Internally use _type_utils.ScalarType scalar_name_to_pytorch = { "uint8_t": "Byte", "int8_t": "Char", "double": "Double", "float": "Float", "half": "Half", "int": "Int", "int64_t": "Long", "int16_t": "Short", "bool": "Bool", "complex64": "ComplexFloat", "complex128": "ComplexDouble", "qint8": "QInt8", "quint8": "QUInt8", "qint32": "QInt32", "bfloat16": "BFloat16", } # Deprecated. Internally use _type_utils.ScalarType # This indicates each scalar type's corresponding # torch type. Related source: # https://github.com/pytorch/pytorch/blob/344defc9733a45fee8d0c4d3f5530f631e823196/c10/core/ScalarType.h scalar_type_to_pytorch_type = [ torch.uint8, # 0 torch.int8, # 1 torch.short, # 2 torch.int, # 3 torch.int64, # 4 torch.half, # 5 torch.float, # 6 torch.double, # 7 torch.complex32, # 8 torch.complex64, # 9 torch.complex128, # 10 torch.bool, # 11 torch.qint8, # 12 torch.quint8, # 13 torch.qint32, # 14 torch.bfloat16, # 15 ] # Deprecated. Internally use _type_utils.ScalarType # source of truth is # https://github.com/pytorch/pytorch/blob/master/torch/csrc/utils/tensor_dtypes.cpp pytorch_name_to_type = { "Byte": torch.uint8, "Char": torch.int8, "Double": torch.double, "Float": torch.float, "Half": torch.half, "Int": torch.int, "Long": torch.int64, "Short": torch.short, "Bool": torch.bool, "ComplexFloat": torch.complex64, "ComplexDouble": torch.complex128, "QInt8": torch.qint8, "QUInt8": torch.quint8, "QInt32": torch.qint32, "BFloat16": torch.bfloat16, } # Deprecated. Internally use _type_utils.ScalarType scalar_type_to_onnx = [ cast_pytorch_to_onnx["Byte"], # 0 cast_pytorch_to_onnx["Char"], # 1 cast_pytorch_to_onnx["Short"], # 2 cast_pytorch_to_onnx["Int"], # 3 cast_pytorch_to_onnx["Long"], # 4 cast_pytorch_to_onnx["Half"], # 5 cast_pytorch_to_onnx["Float"], # 6 cast_pytorch_to_onnx["Double"], # 7 cast_pytorch_to_onnx["Undefined"], # 8 cast_pytorch_to_onnx["ComplexFloat"], # 9 cast_pytorch_to_onnx["ComplexDouble"], # 10 cast_pytorch_to_onnx["Bool"], # 11 cast_pytorch_to_onnx["Char"], # 12 cast_pytorch_to_onnx["Byte"], # 13 cast_pytorch_to_onnx["Int"], # 14 cast_pytorch_to_onnx["BFloat16"], # 15 ] # Global set to store the list of quantized operators in the network. # This is currently only used in the conversion of quantized ops from PT -> C2 via ONNX. _quantized_ops: Set[int] = set()