Softmax symbolic should account for negative dim (#5846)

James Reed · ezyang · commit ba64724aeea8 · 2018-03-19T15:43:41.000-04:00
diff --git a/torch/onnx/symbolic.py b/torch/onnx/symbolic.py
@@ -356,7 +356,9 @@ def softmax(g, input, dim=None):
     #           [0.167, 0.167, 0.167]]
     # So only when dim and axis both equal to ndim - 1 (the last dimension),
     # their semantics are equivalent.
-    if len(input.type().sizes()) != dim + 1:
+    if dim < 0:
+        check_dim = len(input.type().sizes()) + dim
+    if len(input.type().sizes()) != check_dim + 1:
         return _unimplemented("dim", "ONNX and PyTorch use different strategies to split the input.")
     return g.op('Softmax', input, axis_i=dim)
 
