| r"""This file provides a location for operators that help exporting |
| models via onnx. E.g. shape_as_tensor and reshape_from_tensor_shape |
| are to make all dynamic sizes operations traceble. |
| |
| """ |
| |
| import torch |
| import torch.onnx |
| import torch.onnx.utils |
| |
| |
| def _shape_as_tensor(g, input): |
| return g.op('Shape', input) |
| |
| |
| @torch.onnx.symbolic_override(_shape_as_tensor) |
| def shape_as_tensor(x): |
| return torch.LongTensor(tuple(x.shape)) |
| |
| |
| def _reshape_from_tensor_shape(g, input, shape): |
| return g.op('Reshape', input, shape) |
| |
| |
| @torch.onnx.symbolic_override(_reshape_from_tensor_shape) |
| def reshape_from_tensor_shape(x, shape): |
| return x.reshape(shape.tolist()) |
