/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

// Tests for ShapeInference.

#include "tensorflow/compiler/jit/shape_inference.h"

#include <cstdint>
#include <initializer_list>
#include <map>
#include <memory>
#include <string>
#include <vector>

#include "absl/log/check.h"
#include "tensorflow/cc/framework/ops.h"
#include "tensorflow/cc/framework/scope.h"
#include "tensorflow/cc/ops/array_ops.h"
#include "tensorflow/cc/ops/const_op.h"
#include "tensorflow/cc/ops/control_flow_ops.h"
#include "tensorflow/cc/ops/control_flow_ops_internal.h"
#include "tensorflow/cc/ops/math_ops.h"
#include "tensorflow/cc/ops/resource_variable_ops.h"
#include "tensorflow/compiler/jit/test_util.h"
#include "xla/tsl/lib/core/status_test_util.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/platform/test.h"

namespace tensorflow {
namespace {

TEST(ShapeInferenceTest, Basics) {
  Scope root = Scope::NewRootScope().ExitOnError();
  auto a = ops::Placeholder(root.WithOpName("A"), DT_FLOAT,
                            ops::Placeholder::Shape({2, 3}));
  auto b = ops::Placeholder(root.WithOpName("B"), DT_FLOAT,
                            ops::Placeholder::Shape({3}));
  auto c = ops::Placeholder(root.WithOpName("C"), DT_FLOAT);
  auto d = ops::Add(root.WithOpName("D"), a, b);
  auto e = ops::Add(root.WithOpName("E"), d, c);
  auto f = ops::Neg(root.WithOpName("F"), e);
  auto g = ops::AddN(root.WithOpName("G"), std::initializer_list<Output>{e, f});

  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
  CHECK_OK(root.ToGraph(graph.get()));

  GraphShapeInfo shape_info;
  TF_ASSERT_OK(InferShapes(graph.get(), /*arg_shapes=*/{},
                           /*fnlib_def=*/nullptr, &shape_info));

  std::map<std::string, std::vector<PartialTensorShape>> expected = {
      {"A", {PartialTensorShape({2, 3})}}, {"B", {PartialTensorShape({3})}},
      {"C", {PartialTensorShape()}},       {"D", {PartialTensorShape({2, 3})}},
      {"E", {PartialTensorShape()}},       {"F", {PartialTensorShape()}},
      {"G", {PartialTensorShape()}},
  };
  TF_EXPECT_OK(ShapeAnnotationsMatch(*graph, shape_info, expected));
}

// Test that shape inference uses user-given `arg_shapes` to propagate shapes.
TEST(ShapeInferenceTest, UseArgShapesForVariableBatchSize) {
  Scope root = Scope::NewRootScope().ExitOnError();
  auto a = ops::Placeholder(root.WithOpName("A"), DT_FLOAT,
                            ops::Placeholder::Shape({-1, 3}));
  auto b = ops::Placeholder(root.WithOpName("B"), DT_FLOAT,
                            ops::Placeholder::Shape({-1, 3}));
  auto c = ops::Add(root.WithOpName("C"), a, b);
  auto d = ops::Neg(root.WithOpName("D"), c);

  a.node()->AddAttr("_index", 0);
  b.node()->AddAttr("_index", 1);

  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
  CHECK_OK(root.ToGraph(graph.get()));

  std::map<int, InferredShape> arg_shapes;
  arg_shapes[0].shape = TensorShape({2, 3});
  arg_shapes[1].shape = TensorShape({2, 3});

  GraphShapeInfo shape_info;
  TF_ASSERT_OK(InferShapes(graph.get(), arg_shapes,
                           /*fnlib_def=*/nullptr, &shape_info));

  std::map<std::string, std::vector<PartialTensorShape>> expected = {
      {"A", {PartialTensorShape({2, 3})}},
      {"B", {PartialTensorShape({2, 3})}},
      {"C", {PartialTensorShape({2, 3})}},
      {"D", {PartialTensorShape({2, 3})}},
  };
  TF_EXPECT_OK(ShapeAnnotationsMatch(*graph, shape_info, expected));
}

// Test that when user-given `arg_shapes` is incomplete (to cover all
// Placeholders), shape inference still succeeds.
TEST(ShapeInferenceTest, UseArgShapesForVariableBatchSizeIncompleteUserArgs) {
  Scope root = Scope::NewRootScope().ExitOnError();
  auto a = ops::Placeholder(root.WithOpName("A"), DT_FLOAT,
                            ops::Placeholder::Shape({-1, 3}));
  auto b = ops::Placeholder(root.WithOpName("B"), DT_FLOAT,
                            ops::Placeholder::Shape({-1, 3}));
  auto c = ops::Add(root.WithOpName("C"), a, b);
  auto d = ops::Neg(root.WithOpName("D"), c);

  a.node()->AddAttr("_index", 0);
  b.node()->AddAttr("_index", 0);

  std::unique_ptr<Graph> graph(new Graph(OpRegistry::Global()));
  CHECK_OK(root.ToGraph(graph.get()));

  std::map<int, InferredShape> arg_shapes;
  arg_shapes[0].shape = TensorShape({2, 3});

  GraphShapeInfo shape_info;
  TF_ASSERT_OK(InferShapes(graph.get(), arg_shapes,
                           /*fnlib_def=*/nullptr, &shape_info));

  std::map<std::string, std::vector<PartialTensorShape>> expected = {
      {"A", {PartialTensorShape({2, 3})}},
      {"B", {PartialTensorShape({2, 3})}},
      {"C", {PartialTensorShape({2, 3})}},
      {"D", {PartialTensorShape({2, 3})}},
  };
  TF_EXPECT_OK(ShapeAnnotationsMatch(*graph, shape_info, expected));
}

TEST(ShapeInferenceTest, WhileLoop) {
  // Graph:
  // x = array_ops.placeholder(dtypes.int32)
  // y = control_flow_ops.while_loop(lambda i: i < 10, lambda i: i + 1, [x])
  Graph graph(OpRegistry::Global());
  {
    Scope scope = Scope::NewRootScope().ExitOnError();

    auto dummy = ops::Placeholder(scope.WithOpName("Dummy"), DT_INT32,
                                  ops::Placeholder::Shape({}));

    auto source = ops::Placeholder(scope.WithOpName("source"), DT_INT32,
                                   ops::Placeholder::Shape({}));
    auto enter =
        ops::internal::Enter(scope.WithOpName("while/Enter"), source, "aloop");
    // Add an unused Enter node. These should be ignored.
    auto enter2 =
        ops::internal::Enter(scope.WithOpName("while/Enter2"), source, "aloop");
    auto merge = ops::Merge(scope.WithOpName("while/Merge"),
                            std::initializer_list<Input>{enter, dummy});
    auto ten = ops::Const<int32_t>(
        scope.WithOpName("while/Less/y").WithControlDependencies(merge.output),
        10);
    auto less = ops::Less(scope.WithOpName("while/Less"), merge.output, ten);
    auto loop_cond = ops::LoopCond(scope.WithOpName("while/LoopCond"), less);
    auto switch_node =
        ops::Switch(scope.WithOpName("while/Switch"), merge.output, loop_cond);
    auto exit = ops::internal::Exit(scope.WithOpName("while/Exit"),
                                    switch_node.output_false);
    auto identity = ops::Identity(scope.WithOpName("while/Identity"),
                                  switch_node.output_true);
    auto identity_shape =
        ops::Const<int32_t>(scope.WithOpName("while/Identity/shape"), {});
    auto identity_reshaped = ops::Reshape(
        scope.WithOpName("while/Identity/reshaped"), identity, identity_shape);

    auto one = ops::Const<int32_t>(
        scope.WithOpName("while/add/y").WithControlDependencies(identity), 1);
    auto add = ops::Add(scope.WithOpName("while/add"), identity_reshaped, one);
    auto next_iteration =
        ops::NextIteration(scope.WithOpName("while/NextIteration"), add);

    auto sink = ops::Identity(scope.WithOpName("sink"), exit);

    // Remove the dummy node and add the loop backedge.
    scope.graph()->RemoveNode(dummy.node());
    scope.graph()->AddEdge(next_iteration.node(), 0, merge.output.node(), 1);

    TF_EXPECT_OK(scope.ToGraph(&graph));
  }

  GraphShapeInfo shape_info;
  TF_ASSERT_OK(InferShapes(&graph, /*arg_shapes=*/{}, /*fnlib_def=*/nullptr,
                           &shape_info));
  std::map<std::string, std::vector<PartialTensorShape>> expected = {
      {"while/Identity", {PartialTensorShape()}},
      {"while/add", {PartialTensorShape({})}},
  };
  TF_EXPECT_OK(ShapeAnnotationsMatch(graph, shape_info, expected));
}

TEST(ShapeInferenceTest, WhileLoopWithResource) {
  // Graph:
  // x = resource_variable_ops.var_handle_op(dtype=dtypes.float32, shape=[2, 3])
  // y = control_flow_ops.while_loop(lambda _: true, lambda x: x, [x])
  Graph graph(OpRegistry::Global());
  {
    Scope scope = Scope::NewRootScope().ExitOnError();

    auto x =
        ops::VarHandleOp(scope.WithOpName("x"), DT_FLOAT, TensorShape({2, 3}));
    auto enter =
        ops::internal::Enter(scope.WithOpName("while/Enter"), x, "aloop");
    auto dummy = ops::Placeholder(scope.WithOpName("dummy"), DT_RESOURCE);
    auto merge = ops::Merge(scope.WithOpName("while/Merge"),
                            std::initializer_list<Input>{enter, dummy});
    auto false_value = ops::Const<bool>(scope.WithOpName("false"), false);
    auto loop_cond =
        ops::LoopCond(scope.WithOpName("while/LoopCond"), false_value);
    auto switch_node =
        ops::Switch(scope.WithOpName("while/Switch"), merge.output, loop_cond);
    auto exit = ops::internal::Exit(scope.WithOpName("while/Exit"),
                                    switch_node.output_false);
    auto identity = ops::Identity(scope.WithOpName("while/Identity"),
                                  switch_node.output_true);
    auto next_iteration =
        ops::NextIteration(scope.WithOpName("while/NextIteration"), identity);
    auto sink = ops::Identity(scope.WithOpName("sink"), exit);

    // Remove the dummy node and add the loop backedge.
    scope.graph()->RemoveNode(dummy.node());
    scope.graph()->AddEdge(next_iteration.node(), 0, merge.output.node(), 1);

    TF_EXPECT_OK(scope.ToGraph(&graph));
  }

  // Check that we can infer shape for "sink" node (Merge node output).
  GraphShapeInfo shape_info;
  TF_ASSERT_OK(InferShapes(&graph, /*arg_shapes=*/{}, /*fnlib_def=*/nullptr,
                           &shape_info));
  auto iter = shape_info.find("sink");
  EXPECT_NE(iter, shape_info.end());
  EXPECT_EQ(iter->second.size(), 1);
  EXPECT_EQ(iter->second.at(0).handle_type, DT_FLOAT);
  TensorShape resource_shape;
  EXPECT_TRUE(iter->second.at(0).handle_shape.AsTensorShape(&resource_shape));
  EXPECT_EQ(resource_shape, TensorShape({2, 3}));
}

}  // namespace
}  // namespace tensorflow