diff options
Diffstat (limited to 'tensorflow/core/ops/array_ops_test.cc')
| -rw-r--r-- | tensorflow/core/ops/array_ops_test.cc | 15 |
1 files changed, 15 insertions, 0 deletions
diff --git a/tensorflow/core/ops/array_ops_test.cc b/tensorflow/core/ops/array_ops_test.cc index b1463338fb..c15409a246 100644 --- a/tensorflow/core/ops/array_ops_test.cc +++ b/tensorflow/core/ops/array_ops_test.cc @@ -27,6 +27,21 @@ limitations under the License. namespace tensorflow { +TEST(ArrayOpsTest, UnravelIndex_ShapeFn) { + ShapeInferenceTestOp op("UnravelIndex"); + + INFER_OK(op, "?;?", "?"); + + INFER_OK(op, "[];[?]", "[d1_0]"); + + INFER_OK(op, "[4,5];[?]", "[d1_0,20]"); + INFER_OK(op, "[2,3,4];[?]", "[d1_0,24]"); + INFER_OK(op, "?;[?]", "?"); + INFER_OK(op, "[?];[?]", "[d1_0,?]"); + + INFER_ERROR("Shape must be rank 1 but is rank 2", op, "?;[1,1]"); +} + TEST(ArrayOpsTest, Pack_ShapeFn) { ShapeInferenceTestOp op("Pack"); auto set_axis = [&op](int axis) { |