introduce shape_as_tensor and reshape_from_variable_shape (#5824)

anderspapitto · colesbury · commit f446b82e70ca · 2018-03-19T13:30:27.000-04:00
diff --git a/torch/onnx/operators.py b/torch/onnx/operators.py
@@ -0,0 +1,20 @@
+import torch
+import torch.onnx
+
+
+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())
