def index_copy(g, self, dim, index, source):
dim_value = sym_help._parse_arg(dim, 'i')
if sym_help._operator_export_type == torch.onnx.OperatorExportTypes.ONNX_ATEN_FALLBACK:
return g.op("ATen",
self,
index,
source,
dim_i=dim_value,
operator_s="index_copy")
expanded_index_shape, expanded_index = sym_help._index_fill_reshape_helper(
g, self, dim, index)
return scatter(g, self, dim, expanded_index, source)
def pad(g, input, pad, mode, value):
mode = sym_help._parse_arg(mode, "s")
if mode == "replicate":
return replication_pad(g, input, pad)
elif mode == "reflect":
return reflection_pad(g, input, pad)
elif mode == "constant":
return constant_pad_nd(g, input, pad, value)
elif mode == "circular":
return _pad_circular(g, input, pad)
else:
raise RuntimeError(f"Unrecognized padding mode {mode}")
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 NotImplementedError("Unknown aten::slice signature")
is_start_none = start.node().kind() == "prim::Constant" and start.type(
).kind() == "NoneType"
is_end_none = end.node().kind() == "prim::Constant" and end.type().kind(
) == "NoneType"
is_start_onnx_const = start.node().kind() == "onnx::Constant"
is_end_onnx_const = end.node().kind() == "onnx::Constant"
step = sym_help._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 sym_help._parse_arg(start, "i")]
end = [
9223372036854775807 if is_end_none else sym_help._parse_arg(
end, "i")
]
dim = [sym_help._parse_arg(dim, "i")]
dynamic_slice = False
return sym_help._slice_helper(g,
self,
axes=dim,
starts=start,
ends=end,
steps=[step],
dynamic_slice=dynamic_slice)
def slice(g, self, *args):
if len(args) == 4:
# aten::slice(Tensor self, int dim, int start, int end, int step) -> Tensor
dim, start, end, step = args
elif len(args) == 3:
# aten::slice(t[] l, int start, int end, int step) -> t[]
start, end, step = args
dim = 0
else:
raise NotImplementedError("Unknown aten::slice signature")
step = sym_help._parse_arg(step, 'i')
if (start.node().kind() != 'onnx::Constant' or
(not isinstance(end, int) and end.node().kind() != 'onnx::Constant') or
(not isinstance(dim, int) and dim.node().kind() != 'onnx::Constant')):
dynamic_slice = True
else:
start = [sym_help._parse_arg(start, 'i')]
end = [sym_help._parse_arg(end, 'i')]
dim = [sym_help._parse_arg(dim, 'i')]
dynamic_slice = False
return sym_help._slice_helper(g, self, axes=dim, starts=start, ends=end, steps=[step], dynamic_slice=dynamic_slice)
def slice(g, input, dim, start, end, step):
if input not in sym_help._quantized_ops:
from torch.onnx.symbolic_opset9 import slice
return slice(g, input, dim, start, end, step)
if step != 1:
raise RuntimeError(
"ONNX quantized slice export only works for step 1.")
start = sym_help._parse_arg(start, 'i')
end = sym_help._parse_arg(end, 'i')
dim = sym_help._parse_arg(dim, 'i')
kwargs = {
"start_idx_i": start,
"end_idx_i": end,
"dim_i": dim,
"Y_scale_f": input.node()["Y_scale"],
"Y_zero_point_i": input.node()["Y_zero_point"],
}
output = g.op("_caffe2::Int8Slice", input, **kwargs)
sym_help._quantized_ops.add(output)
return output
def slice(g, input, dim, start, end, step):
if input not in symbolic_helper._quantized_ops:
return opset9.slice(g, input, dim, start, end, step)
if step != 1:
raise RuntimeError(
"ONNX quantized slice export only works for step 1.")
start = symbolic_helper._parse_arg(start, "i")
end = symbolic_helper._parse_arg(end, "i")
dim = symbolic_helper._parse_arg(dim, "i")
kwargs = {
"start_idx_i": start,
"end_idx_i": end,
"dim_i": dim,
"Y_scale_f": symbolic_helper._node_get(input.node(), "Y_scale"),
"Y_zero_point_i": symbolic_helper._node_get(input.node(),
"Y_zero_point"),
}
output = g.op("_caffe2::Int8Slice", input, **kwargs)
symbolic_helper._quantized_ops.add(output)
return output
def cat(g, tensor_list, dim, scale=None, zero_point=None):
tensors = symbolic_helper._unpack_list(tensor_list)
input = tensors[0]
if input not in symbolic_helper._quantized_ops:
return opset9.cat(g, tensor_list, dim)
dim = symbolic_helper._parse_arg(dim, "i")
kwargs = {
"Y_scale_f": tensors[0].node()["Y_scale"],
"Y_zero_point_i": tensors[0].node()["Y_zero_point"],
}
output = g.op("_caffe2::Int8Concat", *tensors, axis_i=dim, **kwargs)
symbolic_helper._quantized_ops.add(output)
return output
def index_put(g, self, indices_list_value, values, accumulate=False):
indices_list = sym_help._unpack_list(indices_list_value)
if sym_help._operator_export_type == torch.onnx.OperatorExportTypes.ONNX_ATEN_FALLBACK:
args = [self] + indices_list + [values, accumulate]
return g.op("ATen", *args, operator_s='index_put')
from torch.onnx.symbolic_opset9 import add, expand
accumulate = sym_help._parse_arg(accumulate, 'b')
index = indices_list[0]
if len(indices_list) > 1:
for ind in indices_list[1:]:
index = add(g, index, ind)
broadcast_index_shape = g.op("Shape", index)
indices_list = [
g.op("Unsqueeze",
expand(g, ind, broadcast_index_shape, None),
axes_i=[-1]) for ind in indices_list
]
index = g.op("Concat", *indices_list, axis_i=-1)
else:
broadcast_index_shape = g.op("Shape", index)
index = g.op("Unsqueeze", index, axes_i=[-1])
sub_data_shape = sym_help._slice_helper(g,
g.op("Shape", self),
axes=[0],
starts=[len(indices_list)],
ends=[maxsize])
values_shape = g.op("Concat",
broadcast_index_shape,
sub_data_shape,
axis_i=0)
values = g.op("Reshape", values, values_shape)
if accumulate:
dtype = self.type().scalarType()
dtype = sym_help.scalar_type_to_onnx.index(
sym_help.cast_pytorch_to_onnx[dtype])
dtype = sym_help.scalar_type_to_pytorch_type[dtype]
zeros = g.op("ConstantOfShape",
g.op("Shape", self),
value_t=torch.tensor([0], dtype=dtype))
result = g.op("ScatterND", zeros, index, values)
result = add(g, self, result)
else:
result = g.op("ScatterND", self, index, values)
return result
def index_fill(g, self, dim, index, value):
dim_value = sym_help._parse_arg(dim, 'i')
if sym_help._operator_export_type == torch.onnx.OperatorExportTypes.ONNX_ATEN_FALLBACK:
return g.op("ATen",
self,
index,
value,
dim_i=dim_value,
operator_s="index_fill")
expanded_index_shape, expanded_index = sym_help._index_fill_reshape_helper(
g, self, dim, index)
value = sym_help._maybe_get_scalar(value)
value = sym_help._if_scalar_type_as(g, value, self)
expanded_value = expand(g, value, expanded_index_shape, None)
return scatter(g, self, dim, expanded_index, expanded_value)
def upsample_nearest2d(g, input, output_size, align_corners=None):
if input not in sym_help._quantized_ops:
from torch.onnx.symbolic_opset9 import upsample_nearest2d as upsample_nearest2d_impl
return upsample_nearest2d_impl(g, input, output_size, align_corners)
output_size = sym_help._parse_arg(output_size, 'is')
kwargs = {
"output_size_i": output_size,
"Y_scale_f": input.node()["Y_scale"],
"Y_zero_point_i": input.node()["Y_zero_point"],
}
output = g.op("_caffe2::Int8ResizeNearest", input, **kwargs)
sym_help._quantized_ops.add(output)
return output
def index_fill(g, self, dim, index, value):
dim_value = sym_help._parse_arg(dim, "i")
if sym_help.is_caffe2_aten_fallback():
return g.at("index_fill",
self,
index,
value,
overload_name="int_Scalar",
dim_i=dim_value)
expanded_index_shape, expanded_index = sym_help._index_fill_reshape_helper(
g, self, dim, index)
value = sym_help._maybe_get_scalar(value)
value = sym_help._if_scalar_type_as(g, value, self)
expanded_value = expand(g, value, expanded_index_shape, None)
return scatter(g, self, dim, expanded_index, expanded_value)
def upsample_nearest2d(g,
input,
output_size,
align_corners=None,
scales_h=None,
scales_w=None):
if input not in symbolic_helper._quantized_ops:
return opset9.upsample_nearest2d(g, input, output_size, align_corners)
output_size = symbolic_helper._parse_arg(output_size, "is")
kwargs = {
"output_size_i": output_size,
"Y_scale_f": input.node()["Y_scale"],
"Y_zero_point_i": input.node()["Y_zero_point"],
}
input = nchw2nhwc(g, input)
output = g.op("_caffe2::Int8ResizeNearest", input, **kwargs)
output = nhwc2nchw(g, output)
symbolic_helper._quantized_ops.add(output)
return output
def index_put(g, self, indices_list_value, values, accumulate=False):
if sym_help._is_packed_list(indices_list_value):
indices_list = sym_help._unpack_list(indices_list_value)
else:
indices_list = [indices_list_value]
if sym_help._operator_export_type == torch.onnx.OperatorExportTypes.ONNX_ATEN_FALLBACK:
args = [self] + indices_list + [values, accumulate]
return g.op("ATen", *args, operator_s='index_put')
from torch.onnx.symbolic_opset9 import add, expand
accumulate = sym_help._parse_arg(accumulate, 'b')
if len(indices_list) == 0:
return values
index = indices_list[0]
if len(indices_list) > 1:
for ind in indices_list[1:]:
index = add(g, index, ind)
broadcast_index_shape = g.op("Shape", index)
indices_list = [
sym_help._unsqueeze_helper(
g, expand(g, ind, broadcast_index_shape, None), [-1])
for ind in indices_list
]
index = g.op("Concat", *indices_list, axis_i=-1)
else:
# Replace index_put node with masked_scatter or masked_fill
# when inputs to the index_put node contains boolean inputs
#
# index_put -> masked_fill
# * input index contains single tensor of Bool type (e.g.: %24 <- %23).
# * input value contains single element (e.g.: %18).
#
# Torch IR
# %mask : Float(2, 2, 2, strides=[4, 2, 1], requires_grad=0, device=cpu) = aten::clone(%0, %6)
# %16 : Bool(2, 2, 2, strides=[4, 2, 1], requires_grad=0, device=cpu) =
# aten::to(%8, %26, %27, %11, %12, %28, %29, %15)
# %18 : Float(requires_grad=0, device=cpu) = prim::Constant[value={1}]()
# %23 : Bool(8, strides=[1], device=cpu) = aten::view(%16, %22)
# %24 : Tensor?[] = prim::ListConstruct(%23)
# %25 : Float(2, 2, 2, strides=[4, 2, 1], requires_grad=0, device=cpu) =
# aten::index_put(%mask, %24, %18, %30)
# return (%25)
#
#
# index_put -> masked_scatter
# * input index contains single tensor of Bool type (e.g.: %32 <- %31).
# * input value contains multiple elements (e.g.: %28).
#
# Torch IR
# %mask : Float(2, 2, 2, strides=[4, 2, 1], requires_grad=0, device=cpu) = aten::clone(%0, %6)
# %28 : Float(8, strides=[1], requires_grad=0, device=cpu)
# = prim::Constant[value= 1 1 1 1 1 1 1 1 [ CPUFloatType{8} ]]()
# %15 : Bool(2, 2, 2, strides=[4, 2, 1], requires_grad=0, device=cpu)
# = aten::ne(%mask, %some_const)
# %23 : Bool(2, 2, 2, strides=[4, 2, 1], requires_grad=0, device=cpu)
# = aten::to(%15, %34, %35, %18, %19, %36, %37, %22)
# %38 : Long(requires_grad=0, device=cpu) = prim::Constant[value={0}]()
# %30 : int[] = prim::Constant[value=[-1]]()
# %31 : Bool(8, strides=[1], device=cpu) = aten::view(%23, %30)
# %32 : Tensor?[] = prim::ListConstruct(%31)
# %33 : Float(2, 2, 2, strides=[4, 2, 1], requires_grad=0, device=cpu)
# = aten::index_put(%mask, %32, %28, %38)
# return (%33)
bool_inp = index
if bool_inp.type() is not None and bool_inp.type().scalarType(
) == 'Bool':
rank = sym_help._get_tensor_rank(values)
if rank is not None and rank == 0:
from torch.onnx.symbolic_opset9 import masked_fill
return masked_fill(g, self, bool_inp, values)
return masked_scatter(g, self, bool_inp, values)
broadcast_index_shape = g.op("Shape", index)
index = sym_help._unsqueeze_helper(g, index, [-1])
sub_data_shape = sym_help._slice_helper(g,
g.op("Shape", self),
axes=[0],
starts=[len(indices_list)],
ends=[maxsize])
values_shape = g.op("Concat",
broadcast_index_shape,
sub_data_shape,
axis_i=0)
# Check if values is a singular value and expand accordingly
rank = sym_help._get_tensor_rank(values)
if rank is not None and rank == 0:
values = expand(g, values, values_shape, None)
values = g.op("Reshape", values, values_shape)
dtype = self.type().scalarType()
if dtype is not None and dtype != values.type().scalarType():
values = g.op("Cast",
values,
to_i=sym_help.cast_pytorch_to_onnx[dtype])
dtype = sym_help.scalar_type_to_onnx.index(
sym_help.cast_pytorch_to_onnx[dtype])
dtype = sym_help.scalar_type_to_pytorch_type[dtype]
if accumulate:
zeros = g.op("ConstantOfShape",
g.op("Shape", self),
value_t=torch.tensor([0], dtype=dtype))
result = g.op("ScatterND", zeros, index, values)
result = add(g, self, result)
else:
result = g.op("ScatterND", self, index, values)
return result
def index_put(g, self, indices_list_value, values, accumulate=False):
indices_list = sym_help._unpack_list(indices_list_value)
if sym_help._operator_export_type == torch.onnx.OperatorExportTypes.ONNX_ATEN_FALLBACK:
args = [self] + indices_list + [values, accumulate]
return g.op("ATen", *args, operator_s='index_put')
from torch.onnx.symbolic_opset9 import add, expand
accumulate = sym_help._parse_arg(accumulate, 'b')
index = indices_list[0]
if len(indices_list) > 1:
for ind in indices_list[1:]:
index = add(g, index, ind)
broadcast_index_shape = g.op("Shape", index)
indices_list = [
sym_help._unsqueeze_helper(
g, expand(g, ind, broadcast_index_shape, None), [-1])
for ind in indices_list
]
index = g.op("Concat", *indices_list, axis_i=-1)
else:
# Replace index_put node with masked_scatter or masked_fill
# when inputs to the index_put node contains boolean inputs
#
# index_put -> masked_fill
#
# before graph(%0 : Float(2, 2, 2, strides=[4, 2, 1], requires_grad=1, device=cpu),
# %some_const : Float(requires_grad=0, device=cpu)):
# %6 : None = prim::Constant()
# %mask : Float(2, 2, 2, strides=[4, 2, 1], requires_grad=0, device=cpu) = aten::clone(%0, %6)
# %8 : Bool(2, 2, 2, strides=[4, 2, 1], requires_grad=0, device=cpu) = aten::ne(%mask, %some_const)
# %26 : Long(requires_grad=0, device=cpu) = prim::Constant[value={11}]()
# %27 : Long(requires_grad=0, device=cpu) = prim::Constant[value={0}]()
# %11 : Device = prim::Constant[value="cpu"]()
# %12 : None = prim::Constant()
# %28 : Long(requires_grad=0, device=cpu) = prim::Constant[value={0}]()
# %29 : Long(requires_grad=0, device=cpu) = prim::Constant[value={0}]()
# %15 : None = prim::Constant()
# %16 : Bool(2, 2, 2, strides=[4, 2, 1], requires_grad=0, device=cpu) =
# aten::to(%8, %26, %27, %11, %12, %28, %29, %15)
# %18 : Float(requires_grad=0, device=cpu) = prim::Constant[value={1}]()
# %30 : Long(requires_grad=0, device=cpu) = prim::Constant[value={0}]()
# %22 : int[] = prim::Constant[value=[-1]]()
# %23 : Tensor = aten::view(%16, %22)
# %24 : Tensor?[] = prim::ListConstruct(%23)
# %25 : Float(2, 2, 2, strides=[4, 2, 1], requires_grad=0, device=cpu) =
# aten::index_put(%mask, %24, %18, %30)
# return (%25)
#
# after graph(%0 : Float(2, 2, 2, strides=[4, 2, 1], requires_grad=0, device=cpu),
# %some_const : Float(requires_grad=0, device=cpu)):
# %3 : Tensor = onnx::Equal(%0, %some_const)
# %4 : Bool(2, 2, 2, strides=[4, 2, 1], requires_grad=0, device=cpu) = onnx::Not(%3)
# %12 : Bool(2, 2, 2, strides=[4, 2, 1], requires_grad=0, device=cpu) = onnx::Cast[to=9](%4)
# %19 : Tensor = onnx::Cast[to=9](%12)
# %20 : Tensor = onnx::Constant[value={1}]()
# %21 : Float(2, 2, 2, strides=[4, 2, 1], requires_grad=0, device=cpu)
# = onnx::Where(%19, %20, %0)
# return (%21)
#
# index_put -> masked_scatter
#
# before graph(%0 : Float(2, 2, 2, strides=[4, 2, 1], requires_grad=1, device=cpu),
# %some_const : Float(requires_grad=0, device=cpu)):
# %6 : None = prim::Constant()
# %mask : Float(2, 2, 2, strides=[4, 2, 1], requires_grad=0, device=cpu) = aten::clone(%0, %6)
# %28 : Float(8, strides=[1], requires_grad=0, device=cpu)
# = prim::Constant[value= 1 1 1 1 1 1 1 1 [ CPUFloatType{8} ]]()
# %15 : Bool(2, 2, 2, strides=[4, 2, 1], requires_grad=0, device=cpu)
# = aten::ne(%mask, %some_const)
# %34 : Long(requires_grad=0, device=cpu) = prim::Constant[value={11}]()
# %35 : Long(requires_grad=0, device=cpu) = prim::Constant[value={0}]()
# %18 : Device = prim::Constant[value="cpu"]()
# %19 : None = prim::Constant()
# %36 : Long(requires_grad=0, device=cpu) = prim::Constant[value={0}]()
# %37 : Long(requires_grad=0, device=cpu) = prim::Constant[value={0}]()
# %22 : None = prim::Constant()
# %23 : Bool(2, 2, 2, strides=[4, 2, 1], requires_grad=0, device=cpu)
# = aten::to(%15, %34, %35, %18, %19, %36, %37, %22)
# %38 : Long(requires_grad=0, device=cpu) = prim::Constant[value={0}]()
# %30 : int[] = prim::Constant[value=[-1]]()
# %31 : Tensor = aten::view(%23, %30)
# %32 : Tensor?[] = prim::ListConstruct(%31)
# %33 : Float(2, 2, 2, strides=[4, 2, 1], requires_grad=0, device=cpu)
# = aten::index_put(%mask, %32, %28, %38)
# return (%33)
#
# after graph(%0 : Float(2, 2, 2, strides=[4, 2, 1], requires_grad=0, device=cpu),
# %some_const : Float(requires_grad=0, device=cpu)):
# %3 : Float(8, strides=[1], requires_grad=0, device=cpu)
# = onnx::Constant[value= 1 1 1 1 1 1 1 1 [ CPUFloatType{8} ]]()
# %4 : Tensor = onnx::Equal(%0, %some_const)
# %5 : Bool(2, 2, 2, strides=[4, 2, 1], requires_grad=0, device=cpu) = onnx::Not(%4)
# %13 : Bool(2, 2, 2, strides=[4, 2, 1], requires_grad=0, device=cpu) = onnx::Cast[to=9](%5)
# %19 : Tensor = onnx::Shape(%0)
# %20 : Tensor = onnx::Expand(%13, %19)
# %21 : Tensor = onnx::NonZero(%20)
# %22 : Tensor = onnx::Transpose[perm=[1, 0]](%21)
# %23 : Tensor = onnx::Constant[value={-1}]()
# %24 : Tensor = onnx::Reshape(%3, %23)
# %25 : Tensor = onnx::Shape(%22)
# %27 : Tensor = onnx::Constant[value={0}]()
# %28 : Tensor = onnx::Gather[axis=0](%25, %27)
# %29 : Tensor = onnx::Constant[value={0}]()
# %30 : Tensor = onnx::Unsqueeze[axes=[0]](%29)
# %31 : Tensor = onnx::Unsqueeze[axes=[0]](%28)
# %32 : Tensor = onnx::Constant[value={0}]()
# %33 : Tensor = onnx::Unsqueeze[axes=[0]](%32)
# %34 : Tensor = onnx::Slice(%24, %30, %31, %33)
# %35 : Float(2, 2, 2, strides=[4, 2, 1], requires_grad=0, device=cpu)
# = onnx::ScatterND(%0, %22, %34)
# return (%35)
bool_inp = list(index.node().inputs())[0]
if bool_inp.type() is not None and bool_inp.type().scalarType(
) == 'Bool':
rank = sym_help._get_tensor_rank(values)
if rank is not None and rank == 0:
from torch.onnx.symbolic_opset9 import masked_fill
return masked_fill(g, self, bool_inp, values)
return masked_scatter(g, self, bool_inp, values)
broadcast_index_shape = g.op("Shape", index)
index = sym_help._unsqueeze_helper(g, index, [-1])
sub_data_shape = sym_help._slice_helper(g,
g.op("Shape", self),
axes=[0],
starts=[len(indices_list)],
ends=[maxsize])
values_shape = g.op("Concat",
broadcast_index_shape,
sub_data_shape,
axis_i=0)
values = g.op("Reshape", values, values_shape)
if accumulate:
dtype = self.type().scalarType()
dtype = sym_help.scalar_type_to_onnx.index(
sym_help.cast_pytorch_to_onnx[dtype])
dtype = sym_help.scalar_type_to_pytorch_type[dtype]
zeros = g.op("ConstantOfShape",
g.op("Shape", self),
value_t=torch.tensor([0], dtype=dtype))
result = g.op("ScatterND", zeros, index, values)
result = add(g, self, result)
else:
result = g.op("ScatterND", self, index, values)
return result
def index_put(g, self, indices_list_value, values, accumulate=False):
if symbolic_helper._is_packed_list(indices_list_value):
indices_list = symbolic_helper._unpack_list(indices_list_value)
else:
indices_list = [indices_list_value]
if symbolic_helper.is_caffe2_aten_fallback():
args = [self] + indices_list + [values, accumulate]
return g.at("index_put", *args)
accumulate = symbolic_helper._parse_arg(accumulate, "b")
if len(indices_list) == 0:
return values
if len(indices_list) > 1:
for idx_ in range(len(indices_list)):
if indices_list[idx_].type().scalarType() == "Bool": # type: ignore[attr-defined]
# TODO(justinchuby): Remove type ignore after #81112 is checked in.
indices_list[idx_] = g.op("NonZero", indices_list[idx_])
index = indices_list[0]
for ind in indices_list[1:]:
index = opset9.add(g, index, ind)
broadcast_index_shape = g.op("Shape", index)
indices_list = [
symbolic_helper._unsqueeze_helper(
g, opset9.expand(g, ind, broadcast_index_shape, None), [-1]
)
for ind in indices_list
]
index = g.op("Concat", *indices_list, axis_i=-1)
else:
# Replace index_put node with masked_scatter or masked_fill
# when inputs to the index_put node contains a single boolean input.
#
# index_put -> masked_fill
# * input index contains single tensor of Bool type (e.g.: %24 <- %23).
# * input value contains single element (e.g.: %18).
#
# Torch IR
# %mask : Float(2, 2, 2, strides=[4, 2, 1], requires_grad=0, device=cpu) = aten::clone(%0, %6)
# %16 : Bool(2, 2, 2, strides=[4, 2, 1], requires_grad=0, device=cpu) =
# aten::to(%8, %26, %27, %11, %12, %28, %29, %15)
# %18 : Float(requires_grad=0, device=cpu) = prim::Constant[value={1}]()
# %23 : Bool(8, strides=[1], device=cpu) = aten::view(%16, %22)
# %24 : Tensor?[] = prim::ListConstruct(%23)
# %25 : Float(2, 2, 2, strides=[4, 2, 1], requires_grad=0, device=cpu) =
# aten::index_put(%mask, %24, %18, %30)
# return (%25)
#
#
# index_put -> masked_scatter
# * input index contains single tensor of Bool type (e.g.: %32 <- %31).
# * input value contains multiple elements (e.g.: %28).
#
# Torch IR
# %mask : Float(2, 2, 2, strides=[4, 2, 1], requires_grad=0, device=cpu) = aten::clone(%0, %6)
# %28 : Float(8, strides=[1], requires_grad=0, device=cpu)
# = prim::Constant[value= 1 1 1 1 1 1 1 1 [ CPUFloatType{8} ]]()
# %15 : Bool(2, 2, 2, strides=[4, 2, 1], requires_grad=0, device=cpu)
# = aten::ne(%mask, %some_const)
# %23 : Bool(2, 2, 2, strides=[4, 2, 1], requires_grad=0, device=cpu)
# = aten::to(%15, %34, %35, %18, %19, %36, %37, %22)
# %38 : Long(requires_grad=0, device=cpu) = prim::Constant[value={0}]()
# %30 : int[] = prim::Constant[value=[-1]]()
# %31 : Bool(8, strides=[1], device=cpu) = aten::view(%23, %30)
# %32 : Tensor?[] = prim::ListConstruct(%31)
# %33 : Float(2, 2, 2, strides=[4, 2, 1], requires_grad=0, device=cpu)
# = aten::index_put(%mask, %32, %28, %38)
# return (%33)
index = indices_list[0]
bool_inp = index
if bool_inp.type() is not None and bool_inp.type().scalarType() == "Bool": # type: ignore[attr-defined]
# TODO(justinchuby): Remove type ignore after #81112 is checked in.
rank = symbolic_helper._get_tensor_rank(values)
if rank is not None and rank == 0:
return opset9.masked_fill(g, self, bool_inp, values)
return masked_scatter(g, self, bool_inp, values)
broadcast_index_shape = g.op("Shape", index)
index = symbolic_helper._unsqueeze_helper(g, index, [-1])
sub_data_shape = symbolic_helper._slice_helper(
g, g.op("Shape", self), axes=[0], starts=[len(indices_list)], ends=[sys.maxsize]
)
values_shape = g.op("Concat", broadcast_index_shape, sub_data_shape, axis_i=0)
# Check if values is a singular value and expand accordingly
rank = symbolic_helper._get_tensor_rank(values)
if rank is not None and rank == 0:
values = opset9.expand(g, values, values_shape, None)
values = symbolic_helper._reshape_helper(g, values, values_shape)
dtype = self.type().scalarType()
if dtype is not None and dtype != values.type().scalarType():
values = g.op("Cast", values, to_i=symbolic_helper.cast_pytorch_to_onnx[dtype])
dtype = symbolic_helper.scalar_type_to_onnx.index(
symbolic_helper.cast_pytorch_to_onnx[dtype]
)
dtype = symbolic_helper.scalar_type_to_pytorch_type[dtype]
if accumulate:
zeros = g.op(
"ConstantOfShape",
g.op("Shape", self),
value_t=torch.tensor([0], dtype=dtype),
)
result = g.op("ScatterND", zeros, index, values)
result = add(g, self, result)
else:
result = g.op("ScatterND", self, index, values)
return result