I am wondering if we enable the flag, how many tests will pass?

@spandantiwari do you want to give it a try?

@houseroad, maybe as the next PR. :-)

i mean locally first, i am wondering how stable it is :-)

new line character

Fixed.

I think Block::param_node() is what you want here. Confusingly, it is called param_node, which overloads the word "parameter" in this pass 😢

That's useful. Thanks for pointing to it.

as mentioned below, blocks have direct accessors for their param nodes: b.param_node()

Fixed.

you could AT_ASSERT this instead of returning; it should be an invariant that every block has a param node.

Done.

this does not appear to be used anywhere. In general, we discourage converting a Symbol (basically an integer) to its display string unless for printing purposes.

Removed.

I think it's clearer to do inputNode->kindOf(attr::value) == AttributeKind::t

Yes. Fixed.

I think instead of adding more outputs to the block's source node (which will add inputs to the graph), you should just insert a new onnx::constant node. Unless you want to expose them as parameters later instead of embedding them in the graph?

As we discussed, the use of initializers is preferred from ONNX point of view (to distinguish between learnable parameters and real constants.), hence keeping it.

yeah this is fine. Or you can run an EliminateDeadCode() pass after this and it should clean up the nodes for you

We are using names/indices a lot here when it's more idiomatic to use Value* directly. (see [1])

Fixed.

If you have the data structure described in [2] this check would just be:

for input in node:
  if isParam(input):
    toRemove.push_back(input)

Agreed. Fixed.

Similarly, if you had the data structure described in [2], this would be

for v in toRemove:
  if !v->hasUses():
    name = data_structure[v].first
    paramsDict.erase(name)

Done.

These changes look good! My one remaining nit is that we shouldn't operate on an entity called sourceNode, since it more clearly expresses your intent to operate on block inputs (as a reminder, block->inputs() is equivalent to param_node()->outputs()).

Good point. Updated.

@houseroad, can you take a look at these cases? I'm not familiar with the onnx->aten semantics

Yes, that will be great. Just fyi - I have tested this part on multiple large production models, including BERT, to make sure that these computations are correct, and the numbers match exactly.

Updated based on @houseroad 's feedback on this section.

I think it would be clearer for you to return a c10::optionalat::Tensor from this function, and return c10::nullopt in the case that node is not a supported op for constant folding.

Done.

same as above; returning c10::nullopt instead of an empty tensor will express your intent more clearly.

Done.

can we just call _jit_flatten unconditionally?

I tried all the tests in test_pytorch_onnx_caffe2.py and nothing seems to break when we call _jit_flatten unconditionally. Conditional removed.

Document the argument do_constant_folding like other arugments?

Done.

Can we change run_model_test's interface as well? both run_actual_test and run_debug_test will need to add do_constant_folding argument.

Yes, sure. Updated.

why do we need this?

We don't. Removed.

what's this for?

Vestige from older version. Removed.

usually, the optimization pass should not directly throw exception. And axes are optional... we should handle such cases. Probably just skip the node.

OK. Removed the exception, just skipping node instead.

check axis exist before running?

Done.

We can also check whether the opset is supported by our constant folding. Op spec may change, and the optimization may not be applicable any more.

Yes, good point. Now checking for opset version (opset 9).

can we extract this as a helper function?

OK. Done.

We can put everything into anonymous namespace except ConstantFoldONNX. So they don't need to be static i think.

I think both are viable and have the same effect. But your suggestion is cleaner. Updated with anonymous namespace.