Upload 5875 files

9dd3461 over 3 years ago

9.29 kB

"""Importing this patches torch._C classes to add ONNX conveniences."""
import numbers
import re
from typing import Any, Iterable, Tuple, Union

import torch
from torch import _C
from torch._C import _onnx as _C_onnx

# Import utils to get _params_dict because it is a global that is accessed by c++ code
from torch.onnx import _deprecation, utils
from torch.onnx._globals import GLOBALS
from torch.onnx._internal import _beartype

_ATTR_PATTERN = re.compile("^(.+)_(([ifstgz])|(ty))$")


# TODO(#78694): Remove this file after PyTorch 1.14.
# All functions in this file are deprecated and should not be used


@_deprecation.deprecated(
    "1.13",
    "1.14",
    "note 'g.op()' is to be removed from torch.Graph. Please open a"
    " GitHub issue if you need this functionality.",
)
@_beartype.beartype
def _graph_op(
    g: _C.Graph,
    opname: str,
    *raw_args: Union[torch.Tensor, _C.Value],
    outputs: int = 1,
    **kwargs,
) -> Union[_C.Value, Tuple[_C.Value, ...]]:
    r"""Creates an ONNX operator "opname", taking "args" as inputs and attributes "kwargs".
    The set of operators and the inputs/attributes they take
    is documented at https://github.com/onnx/onnx/blob/master/docs/Operators.md
    This function is monkey-patched onto Graph.
    Args:
        g: The Torch graph.
        opname: The ONNX operator name, e.g., `Abs` or `Add`, or an operator qualified
            with a namespace, e.g., `aten::add`.
        raw_args: The inputs to the operator; usually provided
            as arguments to the `symbolic` definition.
        outputs: The number of outputs this operator returns.
            By default an operator is assumed to return a single output.
            If `outputs` is greater than one, this functions returns a tuple
            of output `Node`, representing each output of the ONNX operator
            in positional.
        kwargs: The attributes of the ONNX operator, whose keys are named
            according to the following convention: `alpha_f` indicates
            the `alpha` attribute with type `f`.  The valid type specifiers are
            `f` (float), `i` (int), `s` (string) or `t` (Tensor).  An attribute
            specified with type float accepts either a single float, or a
            list of floats (e.g., you would say `dims_i` for a `dims` attribute
            that takes a list of integers).
    Returns:
        The node representing the single output of this operator (see the `outputs`
        keyword argument for multi-return nodes).
    """
    # Filter out None attributes, this can be convenient client side because
    # now they can pass through None attributes, and have them not show up
    kwargs = {k: v for k, v in kwargs.items() if v is not None}

    args = [_const_if_tensor(g, arg) for arg in raw_args]

    if "::" in opname:
        namespace, op = opname.split("::")
    else:
        namespace = "onnx"
        op = opname

    n = g.insertNode(_new_node(g, namespace, op, outputs, *args, **kwargs))

    if GLOBALS.onnx_shape_inference:
        _C._jit_pass_onnx_node_shape_type_inference(
            n, utils._params_dict, GLOBALS.export_onnx_opset_version
        )

    if outputs == 1:
        return n.output()
    return tuple(n.outputs())


@_beartype.beartype
def _const_if_tensor(g: _C.Graph, arg):
    if arg is None:
        return arg
    if isinstance(arg, _C.Value):
        return arg
    return _graph_op(g, "Constant", value_z=arg)


@_deprecation.deprecated(
    "1.13",
    "1.14",
    "note 'g.at()' is to be removed from torch.Graph. Please open a"
    " GitHub issue if you need this functionality.",
)
# Generate an ONNX ATen op node.
@_beartype.beartype
def _aten_op(g: _C.Graph, operator: str, *args, overload_name: str = "", **kwargs):
    return _graph_op(
        g,
        "aten::ATen",
        *args,
        operator_s=operator,
        overload_name_s=overload_name,
        **kwargs,
    )


@_deprecation.deprecated(
    "1.13",
    "1.14",
    "note 'b.op()' is to be removed from torch.Block. Please open a"
    " GitHub issue if you need this functionality.",
)
@_beartype.beartype
def _block_op(block: _C.Block, opname: str, *args: _C.Value, **kwargs):
    if "::" in opname:
        namespace, op = opname.split("::")
    else:
        namespace = "onnx"
        op = opname

    n = block.addNode(f"{namespace}::{op}", args)
    aten = namespace == "aten"
    skip_attrs = {"inplace", "aten"}
    for k, v in sorted(kwargs.items()):
        if k in skip_attrs:
            continue
        _add_attribute(n, k, v, aten=aten)
    outputs = tuple(n.outputs())
    if len(outputs) == 1:
        return n.output()
    return outputs


@_beartype.beartype
def _new_node(
    g: _C.Graph, namespace: str, op: str, outputs: int, *args: _C.Value, **kwargs
) -> _C.Node:
    """Creates a new node in the graph.
    Args:
        g: The graph to create the operator on.
        namespace: The namespace of the operator. E.g., "aten", "onnx".
        op: The name of the operator to create.
        outputs: The number of the outputs of the node.
    Returns:
        The new node.
    """
    aten = namespace == "aten"
    node = g.create(f"{namespace}::{op}", args, outputs)
    skip_attrs = {"inplace", "aten"}
    for k, v in sorted(kwargs.items()):
        if k in skip_attrs:
            continue
        _add_attribute(node, k, v, aten=aten)
    return node


@_beartype.beartype
def _is_onnx_list(value):
    return (
        not isinstance(value, torch._six.string_classes)
        and not isinstance(value, torch.Tensor)
        and isinstance(value, Iterable)
    )


@_beartype.beartype
def _scalar(x: torch.Tensor):
    """Convert a scalar tensor into a Python value."""
    assert x.numel() == 1
    return x[0]


@_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 _add_attribute(node: _C.Node, key: str, value: Any, aten: bool):
    r"""Initializes the right attribute based on type of value."""
    m = _ATTR_PATTERN.match(key)
    if m is None:
        raise ValueError(
            f"Invalid attribute specifier '{key}' names "
            "must be suffixed with type, e.g. 'dim_i' or 'dims_i'"
        )
    name, kind = m.group(1), m.group(2)
    if _is_onnx_list(value):
        kind += "s"

    if aten and _is_caffe2_aten_fallback():
        if isinstance(value, torch.Tensor):
            # Caffe2 proto does not support tensor attribute.
            if value.numel() > 1:
                raise ValueError("Should not pass tensor attribute")
            value = _scalar(value)
            if isinstance(value, float):
                kind = "f"
            else:
                kind = "i"
    return getattr(node, f"{kind}_")(name, value)


# TODO(#76254): Remove the deprecated function.
@_deprecation.deprecated(
    "1.13", "1.14", "Use 'g.op()' to create a constant node instead."
)
@_beartype.beartype
def _graph_constant(
    g,
    value,
    dims,
    type_: str,
    *args,
    **kwargs,
):
    """This helper function can create either constant tensor or constant scalar.
    If dims is None or 0 or [0], generate a 0-d tensor (scalar).
    """
    assert isinstance(value, numbers.Number)
    assert type_ is not None
    isscalar = False
    if dims is None or dims == 0 or set(dims) == {0}:
        dims = [1]
        isscalar = True
    type_ = type_.lower()
    tensor: Union[
        torch.CharTensor,
        torch.ShortTensor,
        torch.IntTensor,
        torch.LongTensor,
        torch.HalfTensor,
        torch.FloatTensor,
        torch.DoubleTensor,
    ]
    if type_ == "char":
        tensor = torch.CharTensor(*dims)
    elif type_ == "short":
        tensor = torch.ShortTensor(*dims)
    elif type_ == "int":
        tensor = torch.IntTensor(*dims)
    elif type_ == "long":
        tensor = torch.LongTensor(*dims)
    elif type_ == "half":
        tensor = torch.HalfTensor(*dims)
    elif type_ == "float":
        tensor = torch.FloatTensor(*dims)
    elif type_ == "double":
        tensor = torch.DoubleTensor(*dims)
    else:
        raise ValueError(
            "Unknown type, type should be one of the following strings: "
            "char, short, int, long, half, float, double"
        )
    tensor.fill_(value)  # type: ignore[call-overload]
    if isscalar:
        return g.op("Constant", *args, value_z=tensor, **kwargs)
    return g.op("Constant", *args, value_t=tensor, **kwargs)


# TODO(#76254): Remove the deprecated function.
@_deprecation.deprecated(
    "1.13",
    "1.14",
    "Internally use '_node_get' in symbolic_helper instead.",
)
def _node_getitem(self, k):
    """Gets attributes of a node which is polymorphic over return type.
    This is monkey-patched onto Node.
    """
    sel = self.kindOf(k)
    return getattr(self, sel)(k)


torch._C.Graph.op = _graph_op  # type: ignore[attr-defined]
torch._C.Graph.at = _aten_op  # type: ignore[attr-defined]
torch._C.Block.op = _block_op  # type: ignore[attr-defined]
torch._C.Graph.constant = _graph_constant  # type: ignore[attr-defined]
torch._C.Node.__getitem__ = _node_getitem  # type: ignore[attr-defined, misc, assignment]

	"""Importing this patches torch._C classes to add ONNX conveniences."""
	import numbers
	import re
	from typing import Any, Iterable, Tuple, Union

	import torch
	from torch import _C
	from torch._C import _onnx as _C_onnx

	# Import utils to get _params_dict because it is a global that is accessed by c++ code
	from torch.onnx import _deprecation, utils
	from torch.onnx._globals import GLOBALS
	from torch.onnx._internal import _beartype

	_ATTR_PATTERN = re.compile("^(.+)_(([ifstgz])\|(ty))$")


	# TODO(#78694): Remove this file after PyTorch 1.14.
	# All functions in this file are deprecated and should not be used


	@_deprecation.deprecated(
	"1.13",
	"1.14",
	"note 'g.op()' is to be removed from torch.Graph. Please open a"
	" GitHub issue if you need this functionality.",
	)
	@_beartype.beartype
	def _graph_op(
	g: _C.Graph,
	opname: str,
	*raw_args: Union[torch.Tensor, _C.Value],
	outputs: int = 1,
	**kwargs,
	) -> Union[_C.Value, Tuple[_C.Value, ...]]:
	r"""Creates an ONNX operator "opname", taking "args" as inputs and attributes "kwargs".

	The set of operators and the inputs/attributes they take
	is documented at https://github.com/onnx/onnx/blob/master/docs/Operators.md

	This function is monkey-patched onto Graph.

	Args:
	g: The Torch graph.
	opname: The ONNX operator name, e.g., `Abs` or `Add`, or an operator qualified
	with a namespace, e.g., `aten::add`.
	raw_args: The inputs to the operator; usually provided
	as arguments to the `symbolic` definition.
	outputs: The number of outputs this operator returns.
	By default an operator is assumed to return a single output.
	If `outputs` is greater than one, this functions returns a tuple
	of output `Node`, representing each output of the ONNX operator
	in positional.
	kwargs: The attributes of the ONNX operator, whose keys are named
	according to the following convention: `alpha_f` indicates
	the `alpha` attribute with type `f`. The valid type specifiers are
	`f` (float), `i` (int), `s` (string) or `t` (Tensor). An attribute
	specified with type float accepts either a single float, or a
	list of floats (e.g., you would say `dims_i` for a `dims` attribute
	that takes a list of integers).

	Returns:
	The node representing the single output of this operator (see the `outputs`
	keyword argument for multi-return nodes).
	"""
	# Filter out None attributes, this can be convenient client side because
	# now they can pass through None attributes, and have them not show up
	kwargs = {k: v for k, v in kwargs.items() if v is not None}

	args = [_const_if_tensor(g, arg) for arg in raw_args]

	if "::" in opname:
	namespace, op = opname.split("::")
	else:
	namespace = "onnx"
	op = opname

	n = g.insertNode(_new_node(g, namespace, op, outputs, args, *kwargs))

	if GLOBALS.onnx_shape_inference:
	_C._jit_pass_onnx_node_shape_type_inference(
	n, utils._params_dict, GLOBALS.export_onnx_opset_version
	)

	if outputs == 1:
	return n.output()
	return tuple(n.outputs())


	@_beartype.beartype
	def _const_if_tensor(g: _C.Graph, arg):
	if arg is None:
	return arg
	if isinstance(arg, _C.Value):
	return arg
	return _graph_op(g, "Constant", value_z=arg)


	@_deprecation.deprecated(
	"1.13",
	"1.14",
	"note 'g.at()' is to be removed from torch.Graph. Please open a"
	" GitHub issue if you need this functionality.",
	)
	# Generate an ONNX ATen op node.
	@_beartype.beartype
	def _aten_op(g: _C.Graph, operator: str, args, overload_name: str = "", *kwargs):
	return _graph_op(
	g,
	"aten::ATen",
	*args,
	operator_s=operator,
	overload_name_s=overload_name,
	**kwargs,
	)


	@_deprecation.deprecated(
	"1.13",
	"1.14",
	"note 'b.op()' is to be removed from torch.Block. Please open a"
	" GitHub issue if you need this functionality.",
	)
	@_beartype.beartype
	def _block_op(block: _C.Block, opname: str, args: _C.Value, *kwargs):
	if "::" in opname:
	namespace, op = opname.split("::")
	else:
	namespace = "onnx"
	op = opname

	n = block.addNode(f"{namespace}::{op}", args)
	aten = namespace == "aten"
	skip_attrs = {"inplace", "aten"}
	for k, v in sorted(kwargs.items()):
	if k in skip_attrs:
	continue
	_add_attribute(n, k, v, aten=aten)
	outputs = tuple(n.outputs())
	if len(outputs) == 1:
	return n.output()
	return outputs


	@_beartype.beartype
	def _new_node(
	g: _C.Graph, namespace: str, op: str, outputs: int, args: _C.Value, *kwargs
	) -> _C.Node:
	"""Creates a new node in the graph.

	Args:
	g: The graph to create the operator on.
	namespace: The namespace of the operator. E.g., "aten", "onnx".
	op: The name of the operator to create.
	outputs: The number of the outputs of the node.

	Returns:
	The new node.
	"""
	aten = namespace == "aten"
	node = g.create(f"{namespace}::{op}", args, outputs)
	skip_attrs = {"inplace", "aten"}
	for k, v in sorted(kwargs.items()):
	if k in skip_attrs:
	continue
	_add_attribute(node, k, v, aten=aten)
	return node


	@_beartype.beartype
	def _is_onnx_list(value):
	return (
	not isinstance(value, torch._six.string_classes)
	and not isinstance(value, torch.Tensor)
	and isinstance(value, Iterable)
	)


	@_beartype.beartype
	def _scalar(x: torch.Tensor):
	"""Convert a scalar tensor into a Python value."""
	assert x.numel() == 1
	return x[0]


	@_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 _add_attribute(node: _C.Node, key: str, value: Any, aten: bool):
	r"""Initializes the right attribute based on type of value."""
	m = _ATTR_PATTERN.match(key)
	if m is None:
	raise ValueError(
	f"Invalid attribute specifier '{key}' names "
	"must be suffixed with type, e.g. 'dim_i' or 'dims_i'"
	)
	name, kind = m.group(1), m.group(2)
	if _is_onnx_list(value):
	kind += "s"

	if aten and _is_caffe2_aten_fallback():
	if isinstance(value, torch.Tensor):
	# Caffe2 proto does not support tensor attribute.
	if value.numel() > 1:
	raise ValueError("Should not pass tensor attribute")
	value = _scalar(value)
	if isinstance(value, float):
	kind = "f"
	else:
	kind = "i"
	return getattr(node, f"{kind}_")(name, value)


	# TODO(#76254): Remove the deprecated function.
	@_deprecation.deprecated(
	"1.13", "1.14", "Use 'g.op()' to create a constant node instead."
	)
	@_beartype.beartype
	def _graph_constant(
	g,
	value,
	dims,
	type_: str,
	*args,
	**kwargs,
	):
	"""This helper function can create either constant tensor or constant scalar.

	If dims is None or 0 or [0], generate a 0-d tensor (scalar).
	"""
	assert isinstance(value, numbers.Number)
	assert type_ is not None
	isscalar = False
	if dims is None or dims == 0 or set(dims) == {0}:
	dims = [1]
	isscalar = True
	type_ = type_.lower()
	tensor: Union[
	torch.CharTensor,
	torch.ShortTensor,
	torch.IntTensor,
	torch.LongTensor,
	torch.HalfTensor,
	torch.FloatTensor,
	torch.DoubleTensor,
	]
	if type_ == "char":
	tensor = torch.CharTensor(*dims)
	elif type_ == "short":
	tensor = torch.ShortTensor(*dims)
	elif type_ == "int":
	tensor = torch.IntTensor(*dims)
	elif type_ == "long":
	tensor = torch.LongTensor(*dims)
	elif type_ == "half":
	tensor = torch.HalfTensor(*dims)
	elif type_ == "float":
	tensor = torch.FloatTensor(*dims)
	elif type_ == "double":
	tensor = torch.DoubleTensor(*dims)
	else:
	raise ValueError(
	"Unknown type, type should be one of the following strings: "
	"char, short, int, long, half, float, double"
	)
	tensor.fill_(value) # type: ignore[call-overload]
	if isscalar:
	return g.op("Constant", args, value_z=tensor, *kwargs)
	return g.op("Constant", args, value_t=tensor, *kwargs)


	# TODO(#76254): Remove the deprecated function.
	@_deprecation.deprecated(
	"1.13",
	"1.14",
	"Internally use '_node_get' in symbolic_helper instead.",
	)
	def _node_getitem(self, k):
	"""Gets attributes of a node which is polymorphic over return type.

	This is monkey-patched onto Node.
	"""
	sel = self.kindOf(k)
	return getattr(self, sel)(k)


	torch._C.Graph.op = _graph_op # type: ignore[attr-defined]
	torch._C.Graph.at = _aten_op # type: ignore[attr-defined]
	torch._C.Block.op = _block_op # type: ignore[attr-defined]
	torch._C.Graph.constant = _graph_constant # type: ignore[attr-defined]
	torch._C.Node.__getitem__ = _node_getitem # type: ignore[attr-defined, misc, assignment]