def mm(g, self, other):
# Create a dummy C tensor. Only needed for API purposes, the value is
# since beta = 0
scalar_type = symbolic_helper._try_get_scalar_type(self, other)
if scalar_type is None:
raise errors.SymbolicValueError(
"mm can only operate on tensors with known types", self
)
zero_constant = g.op(
"Constant",
value_t=torch.tensor(
[0], dtype=_type_utils.JitScalarType.from_name(scalar_type).dtype()
),
)
if symbolic_helper._try_get_scalar_type(self):
old_type, self, other, zero_constant = _try_cast_integer_to_float(
g, self, other, zero_constant
)
return _cast_to_type(
g,
g.op("Gemm", self, other, zero_constant, beta_f=0.0, alpha_f=1.0),
old_type,
)
return g.op("Gemm", self, other, zero_constant, beta_f=0.0, alpha_f=1.0)
def fake_quantize_per_tensor_affine(
g, inputs, scale, zero_point, quant_min=-128, quant_max=127
):
# NOTE: (0, 127) is allowed as special case. PyTorch restricts activations to be in the range (0, 127).
# https://github.com/pytorch/pytorch/blob/b34b192d6b97325c9f78e5995c48c8498ede34bd/torch/ao/quantization/observer.py#L1422
if (quant_min, quant_max) not in [(0, 255), (-128, 127), (0, 127)]:
raise errors.SymbolicValueError(
"For (quant_min, quant_max), ONNX allows only (0, 127), (0, 255) and (-128, 127). "
f"Got ({quant_min}, {quant_max})",
inputs,
)
if quant_min == 0:
zero_point = g.op("Cast", zero_point, to_i=_C_onnx.TensorProtoDataType.UINT8)
else:
zero_point = g.op("Cast", zero_point, to_i=_C_onnx.TensorProtoDataType.INT8)
if scale.type().scalarType() != "Float":
scale = g.op("Cast", scale, to_i=_C_onnx.TensorProtoDataType.FLOAT)
quantized = g.op("QuantizeLinear", inputs, scale, zero_point)
if (quant_min, quant_max) == (0, 127):
quantized = g.op(
"Clip",
quantized,
opset9.unused(g),
g.op("Constant", value_t=torch.tensor(127, dtype=torch.uint8)),
)
return g.op("DequantizeLinear", quantized, scale, zero_point)
def split(g, self, split_size_or_sizes, dim, _outputs=None):
if not symbolic_helper._is_split_static(split_size_or_sizes, _outputs):
split_out = g.op("SplitToSequence", self, split_size_or_sizes, axis_i=dim)
if _outputs is None:
return split_out
# Convert to multiple slice nodes iff number of splits and number of outputs are statically known.
if (
symbolic_helper._is_packed_list(split_size_or_sizes)
and len(symbolic_helper._unpack_list(split_size_or_sizes)) == _outputs
):
split_sizes = [
symbolic_helper._unsqueeze_helper(g, v, [0])
for v in symbolic_helper._unpack_list(split_size_or_sizes)
]
start = g.op("Constant", value_t=torch.tensor([0], dtype=torch.long))
axis = g.op("Constant", value_t=torch.tensor([dim], dtype=torch.long))
res = []
for i in range(_outputs):
end = g.op(
"Add", start, split_sizes[i]
) # split_sizes is a list of same length as _outputs
res.append(g.op("Slice", self, start, end, axis))
start = end
return res
return [
g.op(
"SequenceAt",
split_out,
g.op("Constant", value_t=torch.tensor([i], dtype=torch.long)),
)
for i in range(_outputs)
]
split_val = symbolic_helper._node_get(split_size_or_sizes.node(), "value")
if split_val.dim() > 0:
return g.op("Split", self, split_size_or_sizes, axis_i=dim, outputs=_outputs)
split_size = symbolic_helper._get_const(split_size_or_sizes, "i", "split_size")
size = symbolic_helper._get_tensor_dim_size(self, dim)
if size is None:
if _outputs is not None:
size = split_size * _outputs
else:
raise errors.SymbolicValueError(
"Unknown dimension size not supported", self
)
splits = [split_size] * (size // split_size)
leftover = size % split_size
if leftover:
splits.append(leftover)
splits = g.op("Constant", value_t=torch.tensor(splits))
return g.op("Split", self, splits, axis_i=dim, outputs=_outputs)
def fake_quantize_per_tensor_affine(g,
inputs,
scale,
zero_point,
quant_min=-128,
quant_max=127):
# NOTE: (0, 127) is a special case. PyTorch restricts activations to be in the range (0, 127).
# https://github.com/pytorch/pytorch/blob/b34b192d6b97325c9f78e5995c48c8498ede34bd/torch/ao/quantization/observer.py#L1422
if (quant_min, quant_max) == (0, 127):
symbolic_helper._onnx_opset_unsupported_detailed(
"fake_quantize_per_tensor_affine",
10,
13,
"Quantize range (0, 127) not supported, requires opset 13 Clip",
inputs,
)
if (quant_min, quant_max) not in [(0, 255), (-128, 127)]:
raise errors.SymbolicValueError(
f"For (quant_min, quant_max), ONNX allows only (0, 255) and (-128, 127). "
f"Got ({quant_min}, {quant_max})",
inputs,
)
scale = symbolic_helper._maybe_get_scalar(scale)
if scale is None:
symbolic_helper._onnx_opset_unsupported_detailed(
"fake_quantize_per_tensor_affine",
10,
13,
"Non-constant scale not supported",
inputs,
)
scale = scale.float().data # Avoid exporter generating double type
if quant_min == 0:
zero_point = g.op("Cast",
zero_point,
to_i=_C_onnx.TensorProtoDataType.UINT8)
else:
zero_point = g.op("Cast",
zero_point,
to_i=_C_onnx.TensorProtoDataType.INT8)
return g.op(
"DequantizeLinear",
g.op("QuantizeLinear", inputs, scale, zero_point),
scale,
zero_point,
)
def slice(g, self, *args):
if len(args) == 4:
# aten::slice(Tensor self, int dim, int? start=None, int? end=None, int step=1) -> Tensor
dim, start, end, step = args
elif len(args) == 3:
# aten::slice(t[] l, int? start=None, int? end=None, int step=1) -> t[]
start, end, step = args
dim = 0
else:
raise errors.SymbolicValueError("Unknown aten::slice signature", self)
is_start_none = start.node().kind() == "prim::Constant" and isinstance(
start.type(), _C.NoneType)
is_end_none = end.node().kind() == "prim::Constant" and isinstance(
end.type(), _C.NoneType)
is_start_onnx_const = start.node().kind() == "onnx::Constant"
is_end_onnx_const = end.node().kind() == "onnx::Constant"
step = symbolic_helper._parse_arg(step, "i")
if ((not is_start_none and not is_start_onnx_const)
or (not isinstance(end, int) and not is_end_none
and not is_end_onnx_const) or
(not isinstance(dim, int) and dim.node().kind() != "onnx::Constant")):
dynamic_slice = True
if is_start_none:
start = g.op("Constant", value_t=torch.tensor(0))
if is_end_none:
end = g.op("Constant", value_t=torch.tensor(9223372036854775807))
else:
start = [
0 if is_start_none else symbolic_helper._parse_arg(start, "i")
]
end = [
9223372036854775807 if is_end_none else symbolic_helper._parse_arg(
end, "i")
]
dim = [symbolic_helper._parse_arg(dim, "i")]
dynamic_slice = False
return symbolic_helper._slice_helper(
g,
self,
axes=dim,
starts=start,
ends=end,
steps=[step],
dynamic_slice=dynamic_slice,
)
def cross_entropy_loss(g, self, target, weight, reduction, ignore_index,
label_smoothing):
# none reduction : onnx::Constant[value={0}]
# mean reduction : onnx::Constant[value={1}]
# sum reduction : onnx::Constant[value={2}]
reduction = symbolic_helper._maybe_get_const(reduction, "i")
reduction_vals = ["none", "mean", "sum"]
reduction = reduction_vals[reduction]
label_smoothing = symbolic_helper._maybe_get_const(label_smoothing, "f")
if label_smoothing is not None and label_smoothing > 0.0:
raise errors.SymbolicValueError(
"Unsupported: ONNX does not support label_smoothing", self)
# in onnx SoftmaxCrossEntropyLoss specification, ignore_index is optional without default value.
# therefore we need to set ignore_index attribute even if it is not specified (e.g. ignore_index=-100).
ignore_index = symbolic_helper._maybe_get_const(ignore_index, "i")
if weight.node().mustBeNone():
celoss = g.op(
"SoftmaxCrossEntropyLoss",
self,
target,
reduction_s=reduction,
ignore_index_i=ignore_index,
)
else:
celoss = g.op(
"SoftmaxCrossEntropyLoss",
self,
target,
weight,
reduction_s=reduction,
ignore_index_i=ignore_index,
)
return celoss
def repeat_interleave(g, self, repeats, dim=None, output_size=None):
input = self
final_dim = dim
# if dim is None flatten
# By default, use the flattened input array, and return a flat output array
if symbolic_helper._is_none(dim):
input = symbolic_helper._reshape_helper(
g, self, g.op("Constant", value_t=torch.tensor([-1]))
)
dim = 0
else:
dim = symbolic_helper._maybe_get_scalar(dim)
repeats_dim = symbolic_helper._get_tensor_rank(repeats)
repeats_sizes = symbolic_helper._get_tensor_sizes(repeats)
input_sizes = symbolic_helper._get_tensor_sizes(input)
if repeats_dim is None:
raise errors.SymbolicValueError(
"Unsupported: ONNX export of repeat_interleave for unknown repeats rank.",
self,
)
if repeats_sizes is None:
raise errors.SymbolicValueError(
"Unsupported: ONNX export of repeat_interleave for unknown repeats size.",
self,
)
if input_sizes is None:
raise errors.SymbolicValueError(
"Unsupported: ONNX export of repeat_interleave for unknown input size.",
self,
)
# Handle cases where dim is negative
if dim < 0:
dim += len(input_sizes)
output_sizes = input_sizes.copy()
for idx, input_size in enumerate(input_sizes):
if input_size is None:
output_sizes[idx], input_sizes[idx] = 0, -1
cond_dynamic_repeats = repeats_dim == 1 and repeats_sizes[0] is None
# If input size is dynamic or repeats vector is dynamic
if output_sizes[dim] == 0 or cond_dynamic_repeats:
reps = symbolic_helper._size_helper(g, input, dim)
reps = opset11.unsqueeze(g, reps, 0)
# Check if repeats vector is a single integer value
# or a single dimension tensor with non-dynamic values
if repeats_dim == 0 or (repeats_dim == 1 and repeats_sizes[0] == 1):
if not symbolic_helper._is_tensor(repeats):
repeats = g.op("Constant", value_t=torch.LongTensor(repeats))
repeats = g.op("Expand", repeats, reps)
# Check if repeats is dynamic
# As repeats is dynamic, we use a where node as a substitute for the if statement
# If repests_dim = 1, expand repeats otherwise use original tensor
elif cond_dynamic_repeats:
repeat_dim = symbolic_helper._size_helper(
g, repeats, g.op("Constant", value_t=torch.LongTensor([0]))
)
repeat_cond = g.op(
"Equal", repeat_dim, g.op("Constant", value_t=torch.LongTensor([1]))
)
repeats = where(g, repeat_cond, g.op("Expand", repeats, reps), repeats)
# There are cases when the repeats are 1-d tensor with multiple repeats, but dim
# provided along one of the dynamic axes provided. A simple example would be
# input.shape -> [1, 1, *] where * represents the dynamic axes, and dim = 2
# Now, repeat interleaving can be performed in pytorch when the value of * matches
# with the number of elements in repeat, for example if * -> 2, number of repeats
# should be 2 as well.
else:
return opset9.repeat_interleave(g, self, repeats, final_dim)
reps_like = g.op(
"ConstantOfShape",
g.op("Shape", repeats),
value_t=torch.tensor([1], dtype=torch.long),
)
r_splits = split(g, repeats, reps_like, 0)
i_splits = split(g, input, reps_like, dim)
output_sizes[dim], input_sizes[dim] = -1, 1
# Create a loop to iterate over each value along the dimension
# and perform individual interleaving using the repeats tensor
# Loop is of the following pattern
# input (trip_count, cond)
# int trip_count = ...;
# bool cond = ...;
# for (int i=0; i < trip_count && cond; ++i) {
# cond = ...;
# }
# Loop conditions
loop_condition = g.op("Constant", value_t=torch.tensor(1))
loop_condition = g.op("Cast", loop_condition, to_i=9)
loop_len = reps
# Create an empty sequence to store final expansions
final_splits = g.op("SequenceEmpty")
loop = g.op("Loop", loop_len, loop_condition, final_splits)
# Loop inputs
loop_block = utils._add_block(loop.node())
block_input_iter = utils._add_input_to_block(loop_block)
cond = utils._add_input_to_block(loop_block)
final_splits = utils._add_input_to_block(loop_block)
r_split = loop_block.op("SequenceAt", r_splits, block_input_iter)
i_split = loop_block.op("SequenceAt", i_splits, block_input_iter)
i_split = opset11.unsqueeze(loop_block, i_split, dim + 1)
r_concat = [
loop_block.op("Constant", value_t=torch.LongTensor(input_sizes[: dim + 1])),
r_split,
loop_block.op("Constant", value_t=torch.LongTensor(input_sizes[dim + 1 :])),
]
r_concat = loop_block.op("Concat", *r_concat, axis_i=0)
i_split = opset9.expand(loop_block, i_split, r_concat, None)
i_split = symbolic_helper._reshape_helper(
loop_block, i_split, g.op("Constant", value_t=torch.LongTensor(output_sizes))
)
final_splits = loop_block.op("SequenceInsert", final_splits, i_split)
# Loop outputs
cond_out = loop_block.op("Cast", loop_condition, to_i=9)
utils._add_output_to_block(loop_block, cond_out)
utils._add_output_to_block(loop_block, final_splits)
loop_out = loop.node().output()
loop_out = g.op("ConcatFromSequence", loop_out, axis_i=dim)
return loop_out
def tensor_split(g, self, indices_or_sections, dim, _outputs=None):
axis = g.op("Constant", value_t=torch.tensor(dim, dtype=torch.long))
axis = opset11.unsqueeze(g, axis, 0)
const_1 = g.op("Constant", value_t=torch.tensor(1, dtype=torch.long))
if symbolic_helper._is_split_static(indices_or_sections, _outputs):
split_val = symbolic_helper._node_get(indices_or_sections.node(), "value")
if split_val.dim() > 0:
start = g.op("Constant", value_t=torch.tensor([0], dtype=torch.long))
res = []
assert _outputs is not None
for i in range(_outputs - 1):
end = g.op(
"Gather",
indices_or_sections,
g.op("Constant", value_t=torch.tensor([i], dtype=torch.long)),
axis_i=0,
)
res.append(g.op("Slice", self, start, end, axis))
start = end
end = symbolic_helper._size_helper(g, self, axis)
res.append(g.op("Slice", self, start, end, axis))
return res
split_size = symbolic_helper._get_const(
indices_or_sections, "i", "indices_or_sections"
)
size = symbolic_helper._get_tensor_dim_size(self, dim)
if size is None:
if _outputs is not None:
size = split_size * _outputs
else:
raise errors.SymbolicValueError(
"Unknown dimension size not supported", self
)
min_split_size = size // split_size
num_splits_one_extra = size % split_size
splits = num_splits_one_extra * [min_split_size + 1]
leftover = (split_size - num_splits_one_extra) * [min_split_size]
splits = g.op(
"Constant", value_t=torch.tensor(splits + leftover, dtype=torch.long)
)
return g.op("Split", self, splits, axis_i=dim, outputs=_outputs)
if (
symbolic_helper._is_tensor(indices_or_sections)
and symbolic_helper._get_tensor_rank(indices_or_sections) == 1
):
loop_len = symbolic_helper._size_helper(
g, indices_or_sections, g.op("Constant", value_t=torch.tensor(0))
)
loop_len = opset11.unsqueeze(g, loop_len, 0)
loop_condition = g.op("Cast", const_1, to_i=_C_onnx.TensorProtoDataType.BOOL)
# To make the first slice in the below loop work,
# we pad a zero to the first position so that it will be the initial start of slice.
padding_0 = g.op("Constant", value_t=torch.tensor([0], dtype=torch.long))
indices_or_sections = g.op("Concat", padding_0, indices_or_sections, axis_i=0)
final_splits = g.op("SequenceEmpty")
loop = g.op("Loop", loop_len, loop_condition, final_splits)
# Loop inputs
loop_block = utils._add_block(loop.node())
block_input_iter = utils._add_input_to_block(loop_block)
cond = utils._add_input_to_block(loop_block)
final_splits = utils._add_input_to_block(loop_block)
start = loop_block.op("Gather", indices_or_sections, block_input_iter, axis_i=0)
end = loop_block.op(
"Gather",
indices_or_sections,
loop_block.op("Add", block_input_iter, const_1),
axis_i=0,
)
slice = loop_block.op("Slice", self, start, end, axis)
final_splits = loop_block.op("SequenceInsert", final_splits, slice)
# Loop outputs
cond_out = loop_block.op("Identity", loop_condition)
utils._add_output_to_block(loop_block, cond_out)
utils._add_output_to_block(loop_block, final_splits)
loop_out = loop.node().output()
start = g.op(
"Gather",
indices_or_sections,
g.op("Constant", value_t=torch.tensor(-1, dtype=torch.long)),
axis_i=0,
)
start = opset11.unsqueeze(g, start, 0)
end = symbolic_helper._size_helper(g, self, axis)
last_slice = g.op("Slice", self, start, end, axis)
return g.op("SequenceInsert", loop_out, last_slice)
else: # scalar tensor
dim_size = symbolic_helper._size_helper(g, self, axis)
min_split_size = g.op("Div", dim_size, indices_or_sections)
min_split_size_plus_1 = g.op(
"Add",
min_split_size,
const_1,
)
num_splits_one_extra = g.op("Mod", dim_size, indices_or_sections)
splits = g.op("Tile", min_split_size_plus_1, num_splits_one_extra)
leftover = g.op(
"Tile",
min_split_size,
g.op(
"Sub",
opset11.unsqueeze(g, indices_or_sections, 0),
num_splits_one_extra,
),
)
splits = g.op("Concat", splits, leftover, axis_i=0)
if _outputs is None:
return g.op("SplitToSequence", self, splits, axis_i=dim)
return g.op("Split", self, splits, axis_i=dim, outputs=_outputs)
def tensordot(g, input_a, input_b, dims_a, dims_b, out=None):
if out is not None:
symbolic_helper._unimplemented(
"Tensordot", "Out parameter is not supported for tensordot.")
dim_count_a = symbolic_helper._get_tensor_rank(input_a)
if dim_count_a is None:
raise errors.SymbolicValueError(
"Unsupported: ONNX export of tensordot for tensor(input_a) of unknown rank.",
input_a,
)
dim_count_b = symbolic_helper._get_tensor_rank(input_b)
if dim_count_b is None:
raise errors.SymbolicValueError(
"Unsupported: ONNX export of tensordot for tensor(input_b) of unknown rank.",
input_b,
)
dims_a = [(dims_a[i] + dim_count_a) if (dims_a[i] < 0) else dims_a[i]
for i in range(len(dims_a))]
dims_b = [(dims_b[i] + dim_count_b) if (dims_b[i] < 0) else dims_b[i]
for i in range(len(dims_b))]
left_dims_a = [i for i in range(dim_count_a) if (i not in dims_a)]
left_dims_b = [i for i in range(dim_count_b) if (i not in dims_b)]
new_input_a = opset9.permute(g, input_a, left_dims_a + dims_a)
new_input_b = opset9.permute(g, input_b, dims_b + left_dims_b)
input_shape = g.op("Shape", new_input_a)
left_sizes_a = symbolic_helper._slice_helper(g,
input_shape,
axes=[0],
starts=[0],
ends=[len(left_dims_a)])
shape_sizes = [
left_sizes_a,
g.op("Constant", value_t=torch.tensor([-1], dtype=torch.long)),
]
output_a = opset9._reshape_from_tensor(g, new_input_a, shape_sizes)
input_shape = g.op("Shape", output_a)
slices = symbolic_helper._slice_helper(g,
input_shape,
axes=[0],
starts=[-1],
ends=[sys.maxsize])
shape_sizes = [
g.op("Constant", value_t=torch.tensor([-1], dtype=torch.long)),
slices,
]
output_a = opset9._reshape_from_tensor(g, new_input_a, shape_sizes)
input_shape = g.op("Shape", new_input_b)
left_sizes_b = symbolic_helper._slice_helper(g,
input_shape,
axes=[0],
starts=[len(dims_b)],
ends=[sys.maxsize])
slices = symbolic_helper._slice_helper(g,
input_shape,
axes=[0],
starts=[0],
ends=[len(dims_b)])
shape_sizes = [
slices,
g.op("Constant", value_t=torch.tensor([-1], dtype=torch.long)),
]
output_b = opset9._reshape_from_tensor(g, new_input_b, shape_sizes)
input_shape = g.op("Shape", output_b)
slices = symbolic_helper._slice_helper(g,
input_shape,
axes=[0],
starts=[-1],
ends=[sys.maxsize])
shape_sizes = [
g.op("Constant", value_t=torch.tensor([-1], dtype=torch.long)),
slices,
]
output_b = opset9._reshape_from_tensor(g, new_input_b, shape_sizes)
output = einsum(g, "ij,jk->ik",
g.op("prim::ListConstruct", *[output_a, output_b]))
shape_sizes = [left_sizes_a, left_sizes_b]
return opset9._reshape_from_tensor(g, output, shape_sizes)