Method TestRelaxHandles

tensorflow/core/framework/shape_inference_test.cc:1717–1807

}

void ShapeInferenceTest::TestRelaxHandles(bool input_not_output) {
  NodeDef def;
  InferenceContext c(kVersion, &def, MakeOpDef(2, 2), {S({}), S({})}, {}, {},
                     {});
  auto make_shape = [&c](std::initializer_list<int64> dim_sizes) {
    ShapeHandle s;
    TF_CHECK_OK(c.MakeShapeFromPartialTensorShape(S(dim_sizes), &s));
    return s;
  };
  auto get_shapes_and_types_from_context = [&](int idx) {
    if (input_not_output) {
      return c.input_handle_shapes_and_types(idx);
    } else {
      return c.output_handle_shapes_and_types(idx);
    }
  };
  auto relax_shapes_and_types_to_context =
      [&](int idx, const std::vector<ShapeAndType>& shapes_and_types) {
        if (input_not_output) {
          return c.RelaxInputHandleShapesAndMergeTypes(idx, shapes_and_types);
        } else {
          return c.RelaxOutputHandleShapesAndMergeTypes(idx, shapes_and_types);
        }
      };

  EXPECT_TRUE(get_shapes_and_types_from_context(0) == nullptr);
  EXPECT_TRUE(get_shapes_and_types_from_context(1) == nullptr);

  // First relax will take the input completely.
  std::vector<ShapeAndType> t{{make_shape({1, 2, 3}), DT_FLOAT},
                              {c.UnknownShape(), DT_INVALID},
                              {make_shape({4, 3, 2, 1}), DT_INT32}};
  ASSERT_TRUE(relax_shapes_and_types_to_context(0, t));
  ASSERT_TRUE(get_shapes_and_types_from_context(0) != nullptr);
  std::vector<ShapeAndType> v = *get_shapes_and_types_from_context(0);
  ASSERT_EQ(3, v.size());
  for (int i = 0; i < v.size(); ++i) {
    EXPECT_TRUE(SameHandle(t[i].shape, v[i].shape)) << i;
    EXPECT_EQ(t[i].dtype, v[i].dtype);
  }

  // Relax that fails because wrong number of values passed.
  // Fails, and no changes made.
  ASSERT_FALSE(relax_shapes_and_types_to_context(
      0, std::vector<ShapeAndType>{{make_shape({1, 2, 3}), DT_FLOAT}}));
  v = *get_shapes_and_types_from_context(0);
  ASSERT_EQ(3, v.size());
  for (int i = 0; i < v.size(); ++i) {
    EXPECT_TRUE(SameHandle(t[i].shape, v[i].shape)) << i;
    EXPECT_EQ(t[i].dtype, v[i].dtype);
  }

  // Only difference is in a mismatched shape. This should replace
  // the mismatched dimension with an UnknownDim.
  auto t2 = t;
  t2[2].shape = make_shape({4, 3, 4, 1});
  ASSERT_TRUE(relax_shapes_and_types_to_context(0, t2));
  v = *get_shapes_and_types_from_context(0);