Proposal and ideas of shape inference with `dim_expr`

Hello, ONNX community. Today, I want to share an idea for better shape inference.

Recently, I’m so motivated to improve shape inference where dimensions are not always int64_t values, but also contain dim_params. Ideally, my goal is to infer all the cases except for those we don’t know the actual values until we execute the model (like, NonZero operator).

To partially achieve this goal, I’ve improved the shape inference with data propagation for some ops (see https://github.com/onnx/onnx/pulls/xuzijian629).

However, there are many models where we cannot infer the shapes even if we are going to improve current inference with dim_values and dim_params.

One common case is Reshape. Reshaping tensors with dim_params is a very common practice. In #3785 (I’m looking forward to feedbacks to this PR, too!), I supported shape inference like

x: ("a", "b", 2, 3)
x.reshape("a", -1, 6)

and inferred the -1 dim by “b”.

However, we cannot currently support the case like

x: ("a", "b", 2, 3)
x.reshape(-1, 6).reshape("a", -1, 6)

because we’ll get an unknown dimension (which should actually be “a * b”) after the first reshape.

Proposal

I would like to propose shape inference with dim_expr to support the above case. In addition to dim_value and dim_params, I introduce a new member that holds how the dimension was computed.

Basically, the idea and the implementation can be summarized as follows.

• A Dimension class object has either one of dim_value, dim_params or dim_expr.
• Internally, dim_expr is something like std::pair<DimOpType, std::vector<Dimension*>>, which forms a tree by recursion.
• We implement operator overloads to Dimension class. In the overload, we compute the dim_value or dim_param if possible. Otherwise, we store the results in dim_expr.

I think a similar thing is done in https://dl.acm.org/doi/abs/10.1145/3394450.3397465

Remarks

• We know we don’t always want exact expression like (unk__1 * (unk__2 + unk__3) + unk__4) * unk__5 ....
• It will be nice if we replace dim_expr with unk__ at the end, if the expression contains some unknown dimension.
• For example, we will have dimensions like n + 5 or 3 * n at the end, but won’t have n + unk__1 or 3 + unk__2.

Pros

• Shape inference will be much stronger.
• By the operator overloads, we won’t have to write complicated if/else in **/defs.cc.

Cons

• To implement this, ONNX ir version must be updated.
• We may need to modify soooo many existing inference functions.

Background

(This section is my personal feelings about the idea) I think this feature will be essential anyway in near future. Neural networks have become more and more realistic way to solve various problems, and there are growing demands for optimizing those models. Models written by non-experts could be slow or not efficient on target processors. At such situations, model compiler will play an important role. In model optimization, tensor shapes are the most important information. They give information how long each operator takes or how much memory it uses. Since ONNX would be already the de facto standard in model compilers, I think this feature will be welcomed.

[xuzijian629]

Sorry, I don’t understand what you want to say…

With current shape inference, we don’t know unk__0 == unk__23 in https://github.com/onnx/onnx/issues/4064#issuecomment-1074712575, but with a smarter shape inference (like using dim_expr), we can infer that equivalence, and do more optimizatons.

[xuzijian629]

I close this issue because it’s not that important for me now. Please reopen this if needed. Thanks.