experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Fixing master regression in slice op (#14329)

Browse Source
This commit is contained in:
Mahmoud Abuzaina 2017-11-09 14:45:51 -08:00 committed by Martin Wicke
parent b1ce37a0fe
commit 96575f01df
1 changed files with 20 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

23
tensorflow/core/kernels/slice_op.cc
View File

@ -252,8 +252,25 @@ class MklSliceOp : public OpKernel {
if (input_dims == 4) {
HandleCase4D(context, begin, size, result);
} else {
functor::Slice<Device, T, input_dims>()(
context->eigen_device<Device>(), result, input, begin, size);
#define HANDLE_DIM(NDIM) \
if (input_dims == NDIM) { \
functor::Slice<Device, T, NDIM>()( \
context->eigen_device<Device>(), result, input, begin, size); \
return; \
}
HANDLE_DIM(1);
HANDLE_DIM(2);
HANDLE_DIM(3);
HANDLE_DIM(4);
HANDLE_DIM(5);
HANDLE_DIM(6);
#undef HANDLE_DIM
// handle cases which dim >= 7
functor::Slice<Device, T, 7>()(
context->eigen_device<Device>(), result, input, begin, size);
}
}
}
@ -375,7 +392,7 @@ class MklSliceOp : public OpKernel {
}
functor::Slice<Device, T, 4>()(
context->eigen_device<Device>(), result, input, begin, size);
context->eigen_device<Device>(), result, context->input(0), begin, size);
}
};
#endif