pytorch
ba792335 - Export traced aten::unbind (#27247)

Commit
5 years ago
Export traced aten::unbind (#27247) Summary: This PR enables exporting aten::unbind created by the tracer. The traced version IR will always have this pattern ```aten::unbind -> prim::ListUnpack```. Another PR supporting scripted aten::unbind will be submitted separately later. ``` // Unbind is being converted to ONNX as Split + Squeeze. // Example IR // graph(%0 : Float(3, 4, 5)): // %7 : Long() = prim::Constant[value={0}]() // %3 : Tensor[] = aten::unbind(%0, %7) // %4 : Float(4, 5), %5 : Float(4, 5), %6 : Float(4, 5) = prim::ListUnpack(%3) // return (%4, %5, %6) // // Translates to ONNX: // graph(%0 : Float(3, 4, 5)): // %1 : Tensor, %2 : Tensor, %3 : Tensor = onnx::Split[axis=0](%0) // %4 : Float(4, 5) = onnx::Squeeze[axes=[0]](%3) // %5 : Float(4, 5) = onnx::Squeeze[axes=[0]](%2) // %6 : Float(4, 5) = onnx::Squeeze[axes=[0]](%1) // return (%6, %5, %4) ``` Pull Request resolved: https://github.com/pytorch/pytorch/pull/27247 Reviewed By: hl475 Differential Revision: D17791095 Pulled By: houseroad fbshipit-source-id: 83b724275124dd1dedb272583a2fefbdf7035d4c
Author
Parents
Loading