Skip to content

Fix ConcatFromSequence lowering for opsets 11 and 12 #3161

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
DiadkoShmek wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Fix ConcatFromSequence lowering for opsets 11 and 12 #3161

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 15 additions & 5 deletions onnxscript/optimizer/_constant_folding.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -189,6 +189,7 @@ class Replacement:

class OptimizerState:
def __init__(self):
self.opset_imports: dict[str, int] = {}
self._sym_value_map: dict[ir.Value, SymbolicValue] = {}
self._initializer_inputs: list[set[ir.Value]] = []

Expand Down Expand Up @@ -796,13 +797,21 @@ def concat_from_sequence(node: ir.Node, op, state: OptimizerState) -> ReturnValu
logger.debug("ConcatFromSequence => Concat: %s", [x.name for x in inputs])
return op.Concat(*inputs, axis=axis)
if new_axis == 1:
# Unsqueeze the inputs with concat axis if new_axis is 1
axis_value = op.Constant(value_int=axis)
opset_version = state.opset_imports.get(node.domain)
if opset_version is None:
return None
# Unsqueeze takes axes as an attribute before opset 13.
axis_value = op.Constant(value_int=axis) if opset_version >= 13 else None
unsqueezed_inputs = []
for node_input in inputs:
unsqueezed_input = op.Unsqueeze(
node_input, axis_value, _outputs=[f"{node_input.name}_unsqueeze"]
)
if opset_version < 13:
unsqueezed_input = op.Unsqueeze(
node_input, axes=[axis], _outputs=[f"{node_input.name}_unsqueeze"]
)
else:
unsqueezed_input = op.Unsqueeze(
node_input, axis_value, _outputs=[f"{node_input.name}_unsqueeze"]
)
unsqueezed_inputs.append(unsqueezed_input)
# Send unsqueezed outputs to Concat
logger.debug(
Expand Down Expand Up @@ -1450,6 +1459,7 @@ def visit_function(self, function: ir.Function) -> None:
def call(self, model: ir.Model) -> FoldConstantsResult:
self._reset()
self._opset_imports = model.opset_imports
self._state.opset_imports = self._opset_imports
self.visit_graph(model.graph)
for function in model.functions.values():
# TODO(rama): Should we specialize functions?
Expand Down
88 changes: 88 additions & 0 deletions onnxscript/optimizer/_constant_folding_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -608,6 +608,94 @@ def test_static_split_to_sequence_with_list_split_no_keepdims_and_squence_at_is_
ReferenceEvaluator(onnx.parser.parse_model(model)).run(None, {"x": data})[0],
)

def test_concat_from_sequence_preserves_opset_axes_signature(self):
a = np.arange(6, dtype=np.float32).reshape(2, 3)
b = a + 10
for opset in (11, 12, 13, 18):
for new_axis in (0, 1):
for axis in (1, -1):
with self.subTest(opset=opset, new_axis=new_axis, axis=axis):
expected = (
np.stack([a, b], axis=axis)
if new_axis
else np.concatenate([a, b], axis=axis)
)
model = onnx.helper.make_model(
onnx.helper.make_graph(
[
onnx.helper.make_node(
"SequenceConstruct", ["a", "b"], ["seq"]
),
onnx.helper.make_node(
"ConcatFromSequence",
["seq"],
["y"],
axis=axis,
new_axis=new_axis,
),
],
"concat_sequence",
[
onnx.helper.make_tensor_value_info(
name, onnx.TensorProto.FLOAT, [2, 3]
)
for name in ("a", "b")
],
[
onnx.helper.make_tensor_value_info(
"y", onnx.TensorProto.FLOAT, expected.shape
)
],
),
opset_imports=[onnx.helper.make_opsetid("", opset)],
ir_version=10,
)
onnx.checker.check_model(model, full_check=True)
optimized = optimizer.optimize(model)
onnx.checker.check_model(optimized, full_check=True)
actual = ReferenceEvaluator(optimized).run(None, {"a": a, "b": b})[0]
np.testing.assert_array_equal(actual, expected)
self.assertEqual(
{o.domain: o.version for o in optimized.opset_import}, {"": opset}
)

def test_concat_from_sequence_inherits_opset_in_if_branches(self):
a = np.arange(6, dtype=np.float32).reshape(2, 3)
b = a + 10
for opset in (11, 12, 13, 18):
with self.subTest(opset=opset):
model = onnx.parser.parse_model(
f"""
<ir_version: 10, opset_import: ["" : {opset}]>
main (float[2,3] a, float[2,3] b, bool cond) => (float[2,2,3] y) {{
y = If (cond) <
then_branch = then_graph () => (then_y) {{
seq = SequenceConstruct (a, b)
then_y = ConcatFromSequence <axis=1, new_axis=1> (seq)
}},
else_branch = else_graph () => (else_y) {{
seq = SequenceConstruct (b, a)
else_y = ConcatFromSequence <axis=1, new_axis=1> (seq)
}}
>
}}
"""
)
onnx.checker.check_model(model, full_check=True)
optimized = optimizer.optimize(model)
onnx.checker.check_model(optimized, full_check=True)
for attr in optimized.graph.node[0].attribute:
self.assertEqual(
[node.op_type for node in attr.g.node],
["Unsqueeze", "Unsqueeze", "Concat"],
)
for cond in (True, False):
actual = ReferenceEvaluator(optimized).run(
None, {"a": a, "b": b, "cond": np.array(cond)}
)[0]
expected = np.stack([a, b] if cond else [b, a], axis=1)
np.testing.assert_array_equal(actual, expected)

def test_split_to_sequence_and_concat_from_sequence_with_new_axis_0(
self,
):
Expand Down