Support matching variadic inputs

[JoanYu]

I have a ONNX model graph: multi-input -> concat -> transpose, for a better performance for some target device, we need to swap transpose and concat, now the new graph will be multi x (input-> transpose) -> concat (the axis will change), so I wrote a code based on onnxscript.rewriter.pattern

def pattern(op, shape, *seq):
    x = op.Concat(*seq, axis=1)
    x = op.Transpose(x, perm=[0,1,3,2])
    return op.Reshape(x, shape)

def rewrite(op, shape, *seq: Sequence[ValuePattern]):
    print(seq)
    x = op.Concat(seq, axis=3)
    return op.Transpose(x, perm=[0,3,1,2])

but seq will be considered as one input while converted to _pattern_ir._target_pattern

the number of inputs will be dynamic, it seemed to be the RewriteRule class not support the feature now.

BRs, Joan.

[justinchuby]

Thanks for the issue - cc @gramalingam

[JoanYu]

Hi @justinchuby , @gramalingam
here is the demo model for this rewriter rule:

import onnx
from onnx import helper, numpy_helper, shape_inference
import numpy as np

inputs = []
for i in range(12):
    inputs.append(
        helper.make_tensor_value_info(
            'input_'+str(i), 1, [1,1,12,64]
        )
    )
concat0_node = helper.make_node(
    'Concat', ['input_'+str(i) for i in range(12)],
    ['Concat_output_0'], 'Concat',
)
concat0_node.attribute.append(
    helper.make_attribute(
        'axis', 1
    )
)

concat1_node = helper.make_node(
    'Concat', ['input_'+str(i) for i in range(12)],
    ['Concat_output_0'], 'Concat',
)
concat1_node.attribute.append(
    helper.make_attribute(
        'axis', 3
    )
)

transpose0_node = helper.make_node(
    'Transpose', ['Concat_output_0'], ['Transpose_output_0'], 'Transpose'
)
transpose0_node.attribute.append(
    helper.make_attribute(
        'perm', [0,1,3,2]
    )
)
transpose1_node = helper.make_node(
    'Transpose', ['Concat_output_0'], ['output'], 'Transpose'
)
transpose1_node.attribute.append(
    helper.make_attribute(
        'perm', [0,3,1,2]
    )
)
shape = numpy_helper.from_array(np.array([1,768,1,12]).astype(np.int64), 'shape')
reshape_node = helper.make_node(
    'Reshape', ['Transpose_output_0', 'shape'], ['output'], 'Reshape'
)
outputs = [
    helper.make_tensor_value_info(
            'output', 1, [1,768,1,12]
        )
]

graph0 = helper.make_graph(
    [concat0_node, transpose0_node, reshape_node], 'grpah0',
    inputs, outputs, [shape]
)

graph1 = helper.make_graph(
    [concat1_node, transpose1_node], 'grpah1',
    inputs, outputs, [shape]
)


model0 = helper.make_model(
    graph0
)

model1 = helper.make_model(
    graph1
)
onnx.save(shape_inference.infer_shapes(model0), 'model0.onnx')
onnx.save(shape_inference.infer_shapes(model1), 'model1.onnx')

the generated two onnx models are the mathematical equivalence. when I try to modify the model1 into model2, I wrote the rewriter rule (not includes check()). Actually, the number of inputs are not always 12. If we have the different number of inputs, such as ,4, 8, 16, ... Does the tools have a common script that can cover veriadic inputs?

Thanks,

BRs, Joan

[justinchuby]

There is a _allow_other_inputs=True you can use on the pattern