diff options
Diffstat (limited to 'chapters/tensor_ops.adoc')
| -rw-r--r-- | chapters/tensor_ops.adoc | 5 |
1 files changed, 4 insertions, 1 deletions
diff --git a/chapters/tensor_ops.adoc b/chapters/tensor_ops.adoc index 656af85..cb268fa 100644 --- a/chapters/tensor_ops.adoc +++ b/chapters/tensor_ops.adoc @@ -17,7 +17,7 @@ include::{generated}/operators/ARGMAX.adoc[] [source,c++] ---- -ERROR_IF(axis < 0 || axis >= rank(shape1)); +ERROR_IF(axis < 0 || axis >= rank(shape1) || rank(shape1) > 4); if (axis == 0) { left_shape = []; } else { @@ -111,6 +111,7 @@ ERROR_IF(dilation_y < 1 || dilation_x < 1); ERROR_IF(OH != idiv_check(IH - 1 + pad_top + pad_bottom - (KH - 1) * dilation_y, stride_y) + 1); ERROR_IF(OW != idiv_check(IW - 1 + pad_left + pad_right - (KW - 1) * dilation_x, stride_x) + 1); +pad = flatten([0,0], pad, [0,0]); for_each(0 <= n < N, 0 <= oy < OH, 0 <= ox < OW; 0 <= oc < OC) { out_t acc = 0; index_t iy = oy * stride_y - pad_top; @@ -148,6 +149,7 @@ ERROR_IF(OD != idiv_check(ID - 1 + pad_d0 + pad_d1 - (KD - 1) * dilation_d, ERROR_IF(OH != idiv_check(IH - 1 + pad_top + pad_bottom - (KH - 1) * dilation_y, stride_y) + 1); ERROR_IF(OW != idiv_check(IW - 1 + pad_left + pad_right - (KW - 1) * dilation_x, stride_x) + 1); +pad = flatten([0,0], pad, [0,0]); for_each(0 <= n < N, 0 <= od < OD, 0 <= oy < OH, 0 <= ox < OW; 0 <= oc < OC) { out_t acc = 0; index_t id = od * stride_d - pad_d0; @@ -186,6 +188,7 @@ ERROR_IF(dilation_y < 1 || dilation_x < 1); ERROR_IF(OH != idiv_check(IH - 1 + pad_top + pad_bottom - (KH - 1) * dilation_y, stride_y) + 1); ERROR_IF(OW != idiv_check(IW - 1 + pad_left + pad_right - (KW - 1) * dilation_x, stride_x) + 1); +pad = flatten([0,0], pad, [0,0]); for_each(0 <= n < N, 0 <= oy < OH, 0 <= ox < OW; 0 <= c < C, 0 <= m < M) { out_t acc = 0; index_t iy = oy * stride_y - pad_top; |