Commit 8c5d3efa by Siva Committed by Yizhi Liu

[COMPILER][BUG] Fix out of bound access. (#1723)

* [COMPILER][BUG] Fix out of bound access.

* 	* Review comments.
parent f713aa9c
......@@ -67,10 +67,11 @@ inline TShape ReduceShapeImpl(const TShape& ishape,
if (r_axes.ndim() == indim)
return TShape(keepdims ? indim : 1);
CHECK(r_axes.ndim() < indim);
if (keepdims) {
TShape oshape(ishape);
for (unsigned i = 0, j = 0; i < indim; ++i) {
if (i != r_axes[j]) continue;
if (j >= r_axes.ndim() || i != r_axes[j]) continue;
oshape[i] = 1;
++j;
}
......@@ -79,7 +80,7 @@ inline TShape ReduceShapeImpl(const TShape& ishape,
TShape oshape(indim - r_axes.ndim());
for (unsigned i = 0, j = 0, k = 0; i < indim; ++i) {
if (i == r_axes[j]) {
if (j < r_axes.ndim() && i == r_axes[j]) {
++j;
continue;
}
......@@ -95,7 +96,7 @@ inline bool ReduceShape(const nnvm::NodeAttrs& attrs,
CHECK_EQ(out_attrs->size(), 1U);
if ((*in_attrs)[0].ndim() == 0) return false;
const ReduceParam& param = nnvm::get<ReduceParam>(attrs.parsed);
NNVM_ASSIGN_INPUT_SHAPE(
NNVM_ASSIGN_OUTPUT_SHAPE(
attrs, *out_attrs, 0,
ReduceShapeImpl((*in_attrs)[0], param.axis,
param.keepdims, param.exclude));
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment