def gather(g, self, dim, index, sparse_grad=False):
if sym_help._maybe_get_const(sparse_grad, "i"):
return _unimplemented("gather", "sparse_grad == True")
if sym_help._operator_export_type == torch.onnx.OperatorExportTypes.ONNX_ATEN_FALLBACK:
return g.op("ATen", self, dim, index, sparse_grad, operator_s="gather")
return g.op("GatherElements", self, index, axis_i=dim)
def __interpolate(g, input, size, scale_factor, mode, align_corners,
recompute_scale_factor):
mode = sym_help._maybe_get_const(mode, 's')
if 'linear' in mode:
mode = 'linear'
if 'cubic' in mode:
mode = 'cubic'
align_corners = sym_help._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 "pytorch_half_pixel"
# roi only takes effect with coordinate_transformation_mode="tf_crop_and_resize"
roi = g.op("Constant", value_t=torch.tensor([], dtype=torch.float32))
if not sym_help._is_none(size):
input_size = g.op("Shape", input)
input_size = sym_help._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 (sym_help._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 sym_help._is_packed_list(size) and (
(sym_help._maybe_get_const(size, 't').dim() == 0))
except AttributeError:
is_scalar = not sym_help._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:
if not input.type().dim():
return sym_help._unimplemented(
"interpolate (with a scalar output_size)",
"missing input shape (try giving an array of output_size values)"
)
size = unsqueeze(g, size, 0)
size = [size for i in range(input.type().dim() - 2)]
size = g.op("Concat", *size, axis_i=0)
size = g.op("Cast", size, to_i=sym_help.cast_pytorch_to_onnx['Long'])
size = g.op("Concat", input_size, size, axis_i=0)
scales = g.op("Constant",
value_t=torch.tensor([], dtype=torch.float32))
return g.op(
"Resize",
input,
roi,
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 sym_help._is_none(scales)
if not input.type().dim():
return sym_help._unimplemented("interpolate (with scales)",
"missing input shape")
scales = sym_help._interpolate_get_scales(g, scale_factor,
input.type().dim())
return g.op(
"Resize",
input,
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"
def pixel_shuffle(g, self, upscale_factor):
rank = sym_help._get_tensor_rank(self)
if rank is not None and rank != 4:
return _unimplemented("pixel_shuffle", "only support 4d input")
return g.op("DepthToSpace", self, blocksize_i=upscale_factor, mode_s="CRD")
def gather(g, self, dim, index, sparse_grad=False):
if sym_help._maybe_get_const(sparse_grad, "i"):
return _unimplemented("gather", "sparse_grad == True")
if sym_help.is_caffe2_aten_fallback():
return g.at("gather", self, dim, index, sparse_grad)
return g.op("GatherElements", self, index, axis_i=dim)
def pixel_shuffle(g, self, upscale_factor):
dims = self.type().sizes()
if len(dims) != 4:
return _unimplemented("pixel_shuffle", "only support 4d input")
return g.op("DepthToSpace", self, blocksize_i=upscale_factor, mode_s="CRD")
def tensordot(g, input_a, input_b, dims_a, dims_b, out=None):
if out is not None:
_unimplemented("Tensordot", "Out parameter is not supported for tensordot.")
dim_count_a = sym_help._get_tensor_rank(input_a)
if dim_count_a is None:
raise RuntimeError(
"Unsupported: ONNX export of tensordot for tensor(input_a) of unknown rank."
)
dim_count_b = sym_help._get_tensor_rank(input_b)
if dim_count_b is None:
raise RuntimeError(
"Unsupported: ONNX export of tensordot for tensor(input_b) of unknown rank."
)
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 = permute(g, input_a, left_dims_a + dims_a)
new_input_b = permute(g, input_b, dims_b + left_dims_b)
input_shape = g.op("Shape", new_input_a)
left_sizes_a = sym_help._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 = _reshape_from_tensor(g, new_input_a, shape_sizes)
input_shape = g.op("Shape", output_a)
slices = sym_help._slice_helper(
g, input_shape, axes=[0], starts=[-1], ends=[maxsize]
)
shape_sizes = [
g.op("Constant", value_t=torch.tensor([-1], dtype=torch.long)),
slices,
]
output_a = _reshape_from_tensor(g, new_input_a, shape_sizes)
input_shape = g.op("Shape", new_input_b)
left_sizes_b = sym_help._slice_helper(
g, input_shape, axes=[0], starts=[len(dims_b)], ends=[maxsize]
)
slices = sym_help._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 = _reshape_from_tensor(g, new_input_b, shape_sizes)
input_shape = g.op("Shape", output_b)
slices = sym_help._slice_helper(
g, input_shape, axes=[0], starts=[-1], ends=[maxsize]
)
shape_sizes = [
g.op("Constant", value_t=torch.tensor([-1], dtype=torch.long)),
slices,
]
output_b = _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 _reshape_from_tensor(g, output, shape_sizes)
def unfold(g, input, dimension, size, step):
const_size = sym_help._maybe_get_const(size, "i")
const_step = sym_help._maybe_get_const(step, "i")
if not sym_help._is_value(const_size) and not sym_help._is_value(const_step):
from torch.onnx.symbolic_opset9 import unfold as _unfold
return _unfold(g, input, dimension, const_size, const_step)
if sym_help.is_caffe2_aten_fallback():
return g.at("unfold", input, dimension_i=dimension, size_i=size, step_i=step)
sizedim = sym_help._get_tensor_dim_size(input, dimension)
if sizedim is not None:
low_start = g.op("Constant", value_t=torch.tensor(0))
low_end = g.op("Constant", value_t=torch.tensor(sizedim))
hi_end = g.op("Constant", value_t=torch.tensor(sizedim + 1))
low_indices = g.op("Range", low_start, low_end, step)
hi_indices = g.op("Range", size, hi_end, step)
low_size = sym_help._size_helper(
g, low_indices, g.op("Constant", value_t=torch.tensor(0))
)
hi_size = sym_help._size_helper(
g, hi_indices, g.op("Constant", value_t=torch.tensor(0))
)
ndim = sym_help._get_tensor_rank(input)
perm = list(range(0, ndim))
perm.append(perm.pop(dimension))
unsqueeze_list = []
loop_condition = g.op("Constant", value_t=torch.tensor(1))
loop_condition = g.op("Cast", loop_condition, to_i=9)
loop_len = g.op("Min", low_size, hi_size)
loop = g.op("Loop", loop_len, loop_condition)
loop_block = torch.onnx.utils._add_block(loop.node())
block_input_iter = torch.onnx.utils._add_input_to_block(loop_block)
cond = torch.onnx.utils._add_input_to_block(loop_block)
starts = loop_block.op("Gather", low_indices, block_input_iter)
ends = loop_block.op("Gather", hi_indices, block_input_iter)
axes = loop_block.op("Constant", value_t=torch.tensor([2]))
starts = sym_help._unsqueeze_helper(loop_block, starts, [0])
ends = sym_help._unsqueeze_helper(loop_block, ends, [0])
stack = loop_block.op("Slice", input, starts, ends, axes)
unsqueeze = sym_help._unsqueeze_helper(
loop_block, loop_block.op("Transpose", stack, perm_i=perm), [dimension]
)
unsqueeze_list.append(unsqueeze)
concat = loop_block.op("Concat", *unsqueeze_list, axis_i=0)
cond_out = loop_block.op("Cast", loop_condition, to_i=9)
torch.onnx.utils._add_output_to_block(loop_block, cond_out)
torch.onnx.utils._add_output_to_block(loop_block, concat)
loop_output = loop.node().output()
perm = [0, 1, 2, 3, 4]
perm[0], perm[dimension + 1] = perm[dimension + 1], perm[0]
transpose = g.op("Transpose", loop_output, perm_i=perm)
squeeze = sym_help._squeeze_helper(g, transpose, [0])
return squeeze
else:
return _unimplemented("Unfold", "input size not accessible")
def reduce_dim(g, self, dim, keepdim, dtype):
if dtype.node().kind() != 'prim::Constant':
return _unimplemented(name, "dtype")
return symbolic(g, self, dim, keepdim)
def pixel_unshuffle(g, self, downscale_factor):
rank = sym_help._get_tensor_rank(self)
if rank is not None and rank != 4:
return _unimplemented("pixel_unshuffle", "only support 4d input")
return g.op("SpaceToDepth", self, blocksize_i=downscale_factor)
def diagonal(g, self, offset, dim1, dim2):
dim1_size = size(g,
self,
dim=g.op("Constant", value_t=torch.LongTensor([dim1])))
dim2_size = size(g,
self,
dim=g.op("Constant", value_t=torch.LongTensor([dim2])))
# Create appropriate mask
mask_shape = g.op("Concat", dim1_size, dim2_size, axis_i=0)
mask = zeros(g, mask_shape, None, None, None)
mask = g.op("EyeLike", mask, k_i=offset)
# dim1 and dim2 appended as a dimension at the end of the shape
rank = sym_help._get_tensor_rank(self)
if rank is not None:
axes = list(range(rank))
axes.remove(dim1)
axes.remove(dim2)
self = g.op("Transpose", self, perm_i=axes + [dim1, dim2])
else:
return _unimplemented("diagonal", "unknown input rank")
# Multiply input and mask to calculate values along diagonal
# The mask consists of one values where diagonal values are to be calculated
# For example:
# [[1.1, 1.2, 1.3], * [[1, 0, 0] = [[1.1, 0, 0],
# [2.1, 2.2, 2.3], [0, 1, 0] [0, 2.2, 0],
# [3.1, 3.2, 3.3]] [0, 0, 1]] [0, 0, 3.3]]
result = g.op("Mul", self, mask)
result = sym_help._reducesum_helper(g, result, axes_i=[-1], keepdims_i=0)
# Calculate gather indices based on offset and dims
# If offset is greater than zero, set offset to zero as this aids in
# calculation of selection window
offset_op = g.op("Constant", value_t=torch.LongTensor([offset]))
if offset >= 0:
diag_size = g.op(
"Max", g.op("Min", dim1_size, g.op("Sub", dim2_size, offset_op)),
g.op("Constant", value_t=torch.LongTensor([0])))
offset = 0
else:
diag_size = g.op(
"Max", g.op("Min", g.op("Add", dim1_size, offset_op), dim2_size),
g.op("Constant", value_t=torch.LongTensor([0])))
diag_size = g.op("Concat", diag_size, axis_i=0)
# Calculate which diagonal values to select
# For example, in cases with offsets:
# [[0, 1.1, 0]
# [0, 0, 2.2]]
# we need to select the last two columns, so we create a tensor
# with all columns that are to be selected
# So in this example, it is [1, 2]
select_window_ones_fill = ones(g, diag_size, 4, None, None)
select_window = g.op("CumSum", select_window_ones_fill,
g.op("Constant", value_t=torch.LongTensor([0])))
select_window = g.op(
"Add", select_window,
g.op("Constant", value_t=torch.LongTensor([abs(offset) - 1])))
gather_shape = [
size(g, result, dim=g.op("Constant", value_t=torch.LongTensor([axis])))
for axis in list(range(rank))[:-2]
]
gather_shape.append(diag_size)
gather_shape = g.op("Concat", *gather_shape, axis_i=0)
gather_indices = zeros(g, gather_shape, 4, None, None)
# There might be cases where offset value is greater than number of rows/columns
# and might cause the diagonal to overrun and as a result of this, diag_size would be zero.
# For example, if
# offset = 9, dim1_size = 2 (columns), dim2_size = 4 (rows)
# diag_size = max(min(2, (4-9)), 0) = 0, based on calculation above
# Cases with diagonal overrun always result in diag_size = max(0, -ve value) = 0
# In cases without diagonal overrun, we select the appropriate rows/columns along which we
# are calculating diagonal values. In cases with diagonal overrun, we return a tensor which has
# the dimension of the row/column where overrun occurred as 0-dim, as we are essentially
# returning an empty tensor
overrun_cond = g.op(
"Not",
g.op("Equal", diag_size,
g.op("Constant", value_t=torch.tensor(0, dtype=torch.int64))))
if_op = g.op("If", overrun_cond)
if_node = if_op.node()
if_block = _add_block(if_node)
gather_indices_if_block = if_block.op("Add", gather_indices, select_window)
gather_indices_if_block = sym_help._unsqueeze_helper(
if_block, gather_indices_if_block, [rank - 1])
final_non_overrun_ = if_block.op("GatherND",
result,
gather_indices_if_block,
batch_dims_i=rank - 2)
_add_output_to_block(if_block, final_non_overrun_)
else_block = _add_block(if_node)
final_overrun_ = zeros(else_block, gather_shape, 6, None, None)
_add_output_to_block(else_block, final_overrun_)
return if_op
