[ONNX] Fix controlflow shape inference with contrib op (#60707) (#62762)
authorBowenBao <bowbao@microsoft.com>
Fri, 20 Aug 2021 19:44:29 +0000 (12:44 -0700)
committerFacebook GitHub Bot <facebook-github-bot@users.noreply.github.com>
Fri, 20 Aug 2021 19:45:53 +0000 (12:45 -0700)
commita65d1ae7cc8c104d041f554d460da0a84c0f2d4e
tree1ff2705f8f193b2e5f4ced8be57abe3b3e95e936
parent125e2d02e575612eb427104e7c67f1c28f090db8
[ONNX] Fix controlflow shape inference with contrib op (#60707) (#62762)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/62762

`ONNXShapeTypeInference` for node `n` is skipped if `n` is non ONNX namespace, or if `n` contains any non ONNX namespace nodes. This prevents controlflow nodes containing contrib ops from running `SpecialPostProcess`, which sets up correct node output shape/type information in rare cases.

This PR depends on opset 14 export https://github.com/pytorch/pytorch/pull/59486

Test Plan: Imported from OSS

Reviewed By: SplitInfinity

Differential Revision: D30375180

Pulled By: msaroufim

fbshipit-source-id: 5deacec39f091deb4d75ddd9e660e12fca7f16c5

Co-authored-by: BowenBao <bowbao@microsoft.com>
test/onnx/test_custom_ops.py
torch/csrc/jit/passes/onnx/fixup_onnx_controlflow.cpp
torch/csrc/jit/passes/onnx/fixup_onnx_controlflow.h
torch/csrc/jit/passes/onnx/shape_type_inference.cpp