[Anderspapitto patch] fix the shape inference for broadcasting (#1368)

houseroad · web-flow · commit 8aec4e23a0bf · 2018-08-31T15:35:18.000-07:00
* fix case in shape inference where dimension is unknown on both sides * fix case in shape inference where dimension is unknown on both sides (#1367)
diff --git a/onnx/defs/shape_inference.h b/onnx/defs/shape_inference.h
@@ -388,7 +388,9 @@ inline void multidirectionalBroadcastShapeInference(
       resultShape.add_dim()->set_dim_value(dim_value);
     } else if (num_symbolic_dims == 1) {
       *resultShape.add_dim() = symbolic_dim;
-	}
+    } else {
+      resultShape.add_dim();
+    }
   }
 }
 
diff --git a/onnx/defs/tensor/defs.cc b/onnx/defs/tensor/defs.cc
@@ -300,7 +300,10 @@ ONNX_OPERATOR_SET_SCHEMA(
             fail_shape_inference("Required attribute axis is missing");
           }
           int axis = static_cast<int>(axisAttr->i());
-          if (axis < 0 || rank <= axis) {
+          if (rank <= axis) {
+            fail_shape_inference("rank must be greater than axis");
+          }
+          if (axis < 0) {
             return; // TODO: check if negative axis must be supported
           }
 

Original file line number	Diff line number	Diff line change
`@@ -388,7 +388,9 @@ inline void multidirectionalBroadcastShapeInference(`
`388`	`388`	`resultShape.add_dim()->set_dim_value(dim_value);`
`389`	`389`	`} else if (num_symbolic_dims == 1) {`
`390`	`390`	`*resultShape.add_dim() = symbolic_dim;`
`391`		`- }`
	`391`	`+ } else {`
	`392`	`+ resultShape.add_dim();`
	`393`	`+ }`
`392`	`394`	`}`
`393`	`395`	`}`
`394`	`396`
Original file line number	Diff line number	Diff line change
`@@ -300,7 +300,10 @@ ONNX_OPERATOR_SET_SCHEMA(`
`300`	`300`	`fail_shape_inference("Required attribute axis is missing");`
`301`	`301`	`}`
`302`	`302`	`int axis = static_cast<int>(axisAttr->i());`
`303`		`- if (axis < 0 \|\| rank <= axis) {`
	`303`	`+ if (rank <= axis) {`
	`304`	`+ fail_shape_inference("rank must be greater than axis");`
	`305`	`+ }`
	`306`	`+ if (axis < 0) {`
`304`	`307`	`return; // TODO: check if negative axis must be supported`
`305`	`308`	`}`
`306`	`309`