<|start_filename|>tensorflow/contrib/lite/kernels/space_to_batch_nd.cc<|end_filename|>
/* Copyright 2017 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.
==============================================================================*/
#include <string.h>
#include <vector>
#include "tensorflow/contrib/lite/builtin_op_data.h"
#include "tensorflow/contrib/lite/context.h"
#include "tensorflow/contrib/lite/kernels/internal/optimized/optimized_ops.h"
#include "tensorflow/contrib/lite/kernels/internal/reference/reference_ops.h"
#include "tensorflow/contrib/lite/kernels/internal/tensor.h"
#include "tensorflow/contrib/lite/kernels/kernel_util.h"
#include "tensorflow/contrib/lite/kernels/op_macros.h"

namespace tflite {
namespace ops {
namespace builtin {
namespace space_to_batch_nd {

// This file has two implementations of SpaceToBatchND.
enum KernelType {
  kReference,
  kGenericOptimized,
};

// Inputs specified in the 2nd tensor (block_shape) and 3rd tensor (paddings)
// are ignored. Only use the `block_shape` and `paddings` specified in params.
// TODO(nupurgarg): Support inputs as tensors in SpaceToBatchND.
struct SpaceToBatchNDContext {
  SpaceToBatchNDContext(TfLiteContext* context, TfLiteNode* node) {
    params = reinterpret_cast<TfLiteSpaceToBatchNDParams*>(node->builtin_data);
    input = GetInput(context, node, 0);
    output = GetOutput(context, node, 0);
  }
  TfLiteSpaceToBatchNDParams* params;
  TfLiteTensor* input;
  TfLiteTensor* output;
};

// Currently, only 4D NHWC input/output op_context are supported.
// The 4D array need to have exactly 2 spatial dimensions.
// TODO(nupurgarg): Support arbitrary dimension in SpaceToBatchND.
const int kInputDimensionNum = 4;
const int kOutputDimensionNum = 4;
const int kSpatialDimensionNum = 2;
const int kPaddingDimensionNum = 4;

TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
  TF_LITE_ENSURE(context, NumInputs(node) >= 1 && NumInputs(node) <= 3);
  TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);

  SpaceToBatchNDContext op_context(context, node);
  TF_LITE_ENSURE_EQ(context, NumDimensions(op_context.input),
                    kInputDimensionNum);
  TF_LITE_ENSURE_EQ(context, op_context.params->num_spatial_dimensions,
                    kSpatialDimensionNum);
  TF_LITE_ENSURE_EQ(context, op_context.input->type, op_context.output->type);

  const TfLiteIntArray* input_size = op_context.input->dims;
  const int* block_shape = op_context.params->block_shape;

  TfLiteIntArray* output_size = TfLiteIntArrayCreate(kOutputDimensionNum);

  // Ensures the input height and width (with padding) is a multiple of block
  // shape height and width.
  for (int dim = 0; dim < kSpatialDimensionNum; ++dim) {
    int final_dim_size =
        (input_size->data[dim + 1] + op_context.params->before_paddings[dim] +
         op_context.params->after_paddings[dim]);
    TF_LITE_ENSURE_EQ(context, final_dim_size % block_shape[dim], 0);
    output_size->data[dim + 1] = final_dim_size / block_shape[dim];
  }

  const int output_batch_size =
      input_size->data[0] * block_shape[0] * block_shape[1];
  const int output_channel_size = input_size->data[3];

  output_size->data[0] = output_batch_size;
  output_size->data[3] = output_channel_size;

  return context->ResizeTensor(context, op_context.output, output_size);
}

template <KernelType kernel_type>
TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
  SpaceToBatchNDContext op_context(context, node);

  int block_shape_dims_array[1] = {kSpatialDimensionNum};
  Dims<4> block_shape_dims = GetTensorDims(block_shape_dims_array, 1);

  // Initialize padding array in the format accepted by the kernel code.
  // TODO(nupurgarg): Make kernel code accept padding array format that is
  // consistent with Pad operation (i.e. before_paddings and after_paddings).
  TfLiteIntArray* padding_data = TfLiteIntArrayCreate(kPaddingDimensionNum);
  padding_data->data[0] = op_context.params->before_paddings[0];
  padding_data->data[1] = op_context.params->after_paddings[0];
  padding_data->data[2] = op_context.params->before_paddings[1];
  padding_data->data[3] = op_context.params->after_paddings[1];
  int padding_dims_array[1] = {kPaddingDimensionNum};
  Dims<4> padding_dims = GetTensorDims(padding_dims_array, 1);

#define TF_LITE_SPACE_TO_BATCH_ND(type, scalar)                          \
  type::SpaceToBatchND(GetTensorData<scalar>(op_context.input),          \
                       GetTensorDims(op_context.input),                  \
                       op_context.params->block_shape, block_shape_dims, \
                       padding_data->data, padding_dims,                 \
                       GetTensorData<scalar>(op_context.output),         \
                       GetTensorDims(op_context.output))
  switch (op_context.input->type) {  // Already know in/out types are same.
    case kTfLiteFloat32:
      if (kernel_type == kReference) {
        TF_LITE_SPACE_TO_BATCH_ND(reference_ops, float);
      } else {
        TF_LITE_SPACE_TO_BATCH_ND(optimized_ops, float);
      }
      break;
    case kTfLiteUInt8:
      if (kernel_type == kReference) {
        TF_LITE_SPACE_TO_BATCH_ND(reference_ops, uint8_t);
      } else {
        TF_LITE_SPACE_TO_BATCH_ND(optimized_ops, uint8_t);
      }
      break;
    case kTfLiteInt32:
      if (kernel_type == kReference) {
        TF_LITE_SPACE_TO_BATCH_ND(reference_ops, int32_t);
      } else {
        TF_LITE_SPACE_TO_BATCH_ND(optimized_ops, int32_t);
      }
      break;
    case kTfLiteInt64:
      if (kernel_type == kReference) {
        TF_LITE_SPACE_TO_BATCH_ND(reference_ops, int64_t);
      } else {
        TF_LITE_SPACE_TO_BATCH_ND(optimized_ops, int64_t);
      }
      break;
    default:
      context->ReportError(context,
                           "Type is currently not supported by SpaceToBatch.");
      return kTfLiteError;
  }
#undef TF_LITE_SPACE_TO_BATCH_ND

  TfLiteIntArrayFree(padding_data);
  return kTfLiteOk;
}

}  // namespace space_to_batch_nd

TfLiteRegistration* Register_SPACE_TO_BATCH_ND_REF() {
  static TfLiteRegistration r = {
      nullptr, nullptr, space_to_batch_nd::Prepare,
      space_to_batch_nd::Eval<space_to_batch_nd::kReference>};
  return &r;
}

TfLiteRegistration* Register_SPACE_TO_BATCH_ND_GENERIC_OPT() {
  static TfLiteRegistration r = {
      nullptr, nullptr, space_to_batch_nd::Prepare,
      space_to_batch_nd::Eval<space_to_batch_nd::kGenericOptimized>};
  return &r;
}

TfLiteRegistration* Register_SPACE_TO_BATCH_ND() {
  // return Register_SPACE_TO_BATCH_ND_REF();
  return Register_SPACE_TO_BATCH_ND_GENERIC_OPT();
}

}  // namespace builtin
}  // namespace ops
}  // namespace tflite


<|start_filename|>tensorflow/contrib/lite/kernels/mul_test.cc<|end_filename|>
/* Copyright 2017 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.
==============================================================================*/
#include <gtest/gtest.h>
#include "tensorflow/contrib/lite/interpreter.h"
#include "tensorflow/contrib/lite/kernels/register.h"
#include "tensorflow/contrib/lite/kernels/test_util.h"
#include "tensorflow/contrib/lite/model.h"

namespace tflite {
namespace {

using ::testing::ElementsAreArray;

class BaseMulOpModel : public SingleOpModel {
 public:
  BaseMulOpModel(TensorData input, TensorData output,
                 ActivationFunctionType activation_type) {
    input1_ = AddInput(input);
    input2_ = AddInput(input);
    output_ = AddOutput(output);
    SetBuiltinOp(BuiltinOperator_MUL, BuiltinOptions_MulOptions,
                 CreateMulOptions(builder_, activation_type).Union());
    BuildInterpreter({GetShape(input1_), GetShape(input2_)});
  }

  int input1() { return input1_; }
  int input2() { return input2_; }

 protected:
  int input1_;
  int input2_;
  int output_;
};

class FloatMulOpModel : public BaseMulOpModel {
 public:
  using BaseMulOpModel::BaseMulOpModel;

  std::vector<float> GetOutput() { return ExtractVector<float>(output_); }
};

// For quantized Mul, the error shouldn't exceed (2*step + step^2).
// The param min=-1.0 & max=1.0 is used in the following tests.
// The tolerance value is ~0.0157.
const float kQuantizedStep = 2.0 / 255.0;
const float kQuantizedTolerance =
    2.0 * kQuantizedStep + kQuantizedStep * kQuantizedStep;

class QuantizedMulOpModel : public BaseMulOpModel {
 public:
  using BaseMulOpModel::BaseMulOpModel;

  std::vector<float> GetDequantizedOutput() {
    return Dequantize<uint8_t>(ExtractVector<uint8_t>(output_),
                               GetScale(output_), GetZeroPoint(output_));
  }
};

TEST(FloatMulOpTest, NoActivation) {
  FloatMulOpModel m({TensorType_FLOAT32, {1, 2, 2, 1}},
                    {TensorType_FLOAT32, {}}, ActivationFunctionType_NONE);
  m.PopulateTensor<float>(m.input1(), {-2.0, 0.2, 0.7, 0.8});
  m.PopulateTensor<float>(m.input2(), {0.1, 0.2, 0.3, 0.5});
  m.Invoke();
  EXPECT_THAT(m.GetOutput(),
              ElementsAreArray(ArrayFloatNear({-0.2, 0.04, 0.21, 0.4})));
}

TEST(FloatMulOpTest, ActivationRELU_N1_TO_1) {
  FloatMulOpModel m({TensorType_FLOAT32, {1, 2, 2, 1}},
                    {TensorType_FLOAT32, {}},
                    ActivationFunctionType_RELU_N1_TO_1);
  m.PopulateTensor<float>(m.input1(), {-2.0, 0.2, 0.7, 0.8});
  m.PopulateTensor<float>(m.input2(), {0.1, 0.2, 0.3, 5});
  m.Invoke();
  EXPECT_THAT(m.GetOutput(),
              ElementsAreArray(ArrayFloatNear({-0.2, 0.04, 0.21, 1.0})));
}

TEST(FloatMulOpTest, VariousInputShapes) {
  std::vector<std::initializer_list<int>> test_shapes = {
      {6}, {2, 3}, {2, 1, 3}, {1, 3, 1, 2}};
  for (int i = 0; i < test_shapes.size(); ++i) {
    FloatMulOpModel m({TensorType_FLOAT32, test_shapes[i]},
                      {TensorType_FLOAT32, {}}, ActivationFunctionType_NONE);
    m.PopulateTensor<float>(m.input1(), {-2.0, 0.2, 0.7, 0.8, 1.1, 2.0});
    m.PopulateTensor<float>(m.input2(), {0.1, 0.2, 0.3, 0.5, 1.1, 0.1});
    m.Invoke();
    EXPECT_THAT(
        m.GetOutput(),
        ElementsAreArray(ArrayFloatNear({-0.2, 0.04, 0.21, 0.4, 1.21, 0.2})))
        << "With shape number " << i;
  }
}

TEST(QuantizedMulOpTest, NoActivation) {
  QuantizedMulOpModel m({TensorType_UINT8, {1, 2, 2, 1}, -1.0, 1.0},
                        {TensorType_UINT8, {}, -1.0, 1.0},
                        ActivationFunctionType_NONE);
  m.QuantizeAndPopulate<uint8_t>(m.input1(), {-0.8, 0.2, 0.9, 0.7});
  m.QuantizeAndPopulate<uint8_t>(m.input2(), {0.6, 0.4, 0.9, 0.8});
  m.Invoke();
  EXPECT_THAT(m.GetDequantizedOutput(),
              ElementsAreArray(ArrayFloatNear({-0.48, 0.08, 0.81, 0.56},
                                              kQuantizedTolerance)));
}

}  // namespace
}  // namespace tflite

int main(int argc, char** argv) {
  ::tflite::LogToStderr();
  ::testing::InitGoogleTest(&argc, argv);
  return RUN_ALL_TESTS();
}


<|start_filename|>tensorflow/contrib/lite/examples/label_image/bitmap_helpers_impl.h<|end_filename|>
/* Copyright 2017 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.
==============================================================================*/

#ifndef TENSORFLOW_CONTRIB_LITE_EXAMPLES_LABEL_IMAGE_BITMAP_HELPERS_IMPL_H
#define TENSORFLOW_CONTRIB_LITE_EXAMPLES_LABEL_IMAGE_BITMAP_HELPERS_IMPL_H

#include "tensorflow/contrib/lite/examples/label_image/label_image.h"

namespace tflite {
namespace label_image {

template <class T>
void downsize(T* out, uint8_t* in, int image_height, int image_width,
              int image_channels, int wanted_height, int wanted_width,
              int wanted_channels, Settings* s) {
  for (int y = 0; y < wanted_height; ++y) {
    const int in_y = (y * image_height) / wanted_height;
    uint8_t* in_row = in + (in_y * image_width * image_channels);
    T* out_row = out + (y * wanted_width * wanted_channels);
    for (int x = 0; x < wanted_width; ++x) {
      const int in_x = (x * image_width) / wanted_width;
      uint8_t* in_pixel = in_row + (in_x * image_channels);
      T* out_pixel = out_row + (x * wanted_channels);
      for (int c = 0; c < wanted_channels; ++c) {
        if (s->input_floating)
          out_pixel[c] = (in_pixel[c] - s->input_mean) / s->input_std;
        else
          out_pixel[c] = in_pixel[c];
      }
    }
  }
}

}  // namespace label_image
}  // namespace tflite

#endif  // TENSORFLOW_CONTRIB_LITE_EXAMPLES_LABEL_IMAGE_BITMAP_HELPERS_IMPL_H


<|start_filename|>tensorflow/contrib/lite/kernels/space_to_batch_nd_test.cc<|end_filename|>
/* Copyright 2017 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.
==============================================================================*/

#include <gtest/gtest.h>
#include "tensorflow/contrib/lite/interpreter.h"
#include "tensorflow/contrib/lite/kernels/register.h"
#include "tensorflow/contrib/lite/kernels/test_util.h"
#include "tensorflow/contrib/lite/model.h"

namespace tflite {
namespace {

using ::testing::ElementsAreArray;

class SpaceToBatchNDOpModel : public SingleOpModel {
 public:
  SpaceToBatchNDOpModel(std::initializer_list<int> input_shape,
                        std::initializer_list<int> block_shape,
                        std::initializer_list<int> before_paddings,
                        std::initializer_list<int> after_paddings) {
    input_ = AddInput(TensorType_FLOAT32);
    output_ = AddOutput(TensorType_FLOAT32);
    SetBuiltinOp(BuiltinOperator_SPACE_TO_BATCH_ND,
                 BuiltinOptions_SpaceToBatchNDOptions,
                 CreateSpaceToBatchNDOptions(
                     builder_, builder_.CreateVector<int>(block_shape),
                     builder_.CreateVector<int>(before_paddings),
                     builder_.CreateVector<int>(after_paddings))
                     .Union());
    BuildInterpreter({input_shape});
  }

  void SetInput(std::initializer_list<float> data) {
    PopulateTensor<float>(input_, data);
  }

  std::vector<float> GetOutput() { return ExtractVector<float>(output_); }
  std::vector<int> GetOutputShape() { return GetTensorShape(output_); }

 private:
  int input_;
  int output_;
};

TEST(SpaceToBatchNDOpTest, InvalidShapeTest) {
  EXPECT_DEATH(SpaceToBatchNDOpModel({1, 3, 3, 1}, {2, 2}, {0, 0}, {0, 0}),
               "Cannot allocate tensors");
}

TEST(SpaceToBatchNDOpTest, SimpleTest) {
  SpaceToBatchNDOpModel m({1, 4, 4, 1}, {2, 2}, {0, 0}, {0, 0});
  m.SetInput({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16});
  m.Invoke();
  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({4, 2, 2, 1}));
  EXPECT_THAT(m.GetOutput(), ElementsAreArray({1, 3, 9, 11, 2, 4, 10, 12, 5, 7,
                                               13, 15, 6, 8, 14, 16}));
}

TEST(SpaceToBatchNDOpTest, MultipleInputBatches) {
  SpaceToBatchNDOpModel m({2, 2, 4, 1}, {2, 2}, {0, 0}, {0, 0});
  m.SetInput({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16});
  m.Invoke();
  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({8, 1, 2, 1}));
  EXPECT_THAT(m.GetOutput(), ElementsAreArray({1, 3, 9, 11, 2, 4, 10, 12, 5, 7,
                                               13, 15, 6, 8, 14, 16}));
}

TEST(SpaceToBatchNDOpTest, SimplePadding) {
  SpaceToBatchNDOpModel m({1, 5, 2, 1}, {3, 2}, {1, 2}, {0, 0});
  m.SetInput({1, 2, 3, 4, 5, 6, 7, 8, 9, 10});
  m.Invoke();
  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({6, 2, 2, 1}));
  EXPECT_THAT(m.GetOutput(), ElementsAreArray({
                                 0, 0, 0, 5, 0, 0, 0, 6, 0, 1, 0, 7,
                                 0, 2, 0, 8, 0, 3, 0, 9, 0, 4, 0, 10,
                             }));
}

TEST(SpaceToBatchNDOpTest, ComplexPadding) {
  SpaceToBatchNDOpModel m({1, 4, 2, 1}, {3, 2}, {1, 2}, {1, 4});
  m.SetInput({1, 2, 3, 4, 5, 6, 7, 8});
  m.Invoke();
  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({6, 2, 4, 1}));
  EXPECT_THAT(m.GetOutput(), ElementsAreArray({
                                 0, 0, 0, 0, 0, 5, 0, 0, 0, 0, 0, 0, 0, 6, 0, 0,
                                 0, 1, 0, 0, 0, 7, 0, 0, 0, 2, 0, 0, 0, 8, 0, 0,
                                 0, 3, 0, 0, 0, 0, 0, 0, 0, 4, 0, 0, 0, 0, 0, 0,
                             }));
}

}  // namespace
}  // namespace tflite

int main(int argc, char** argv) {
  ::tflite::LogToStderr();
  ::testing::InitGoogleTest(&argc, argv);
  return RUN_ALL_TESTS();
}


<|start_filename|>tensorflow/core/grappler/optimizers/memory_optimizer_test.cc<|end_filename|>
/* Copyright 2017 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.
==============================================================================*/

#include "tensorflow/core/grappler/optimizers/memory_optimizer.h"

#include <vector>

#include "tensorflow/cc/ops/standard_ops.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/grappler/clusters/virtual_cluster.h"
#include "tensorflow/core/grappler/grappler_item.h"
#include "tensorflow/core/grappler/utils.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/test.h"

namespace tensorflow {
namespace grappler {
namespace {

class RecomputeSubgraphTest : public ::testing::Test {};

TEST_F(RecomputeSubgraphTest, SimpleSubgraph) {
  tensorflow::Scope s = tensorflow::Scope::NewRootScope();

  Output a = ops::Variable(s.WithOpName("a"), {2, 3, 4}, DT_FLOAT);
  Output b = ops::Identity(s.WithOpName("b"), a);  // Recomputed
  Output c = ops::Identity(s.WithOpName("c"), b);
  Output d = ops::AddN(s.WithOpName("gradients/d"), {c});
  Output e = ops::AddN(s.WithOpName("gradients/e"), {d, b});
  Output f = ops::AddN(s.WithOpName("gradients/f"), {e, a});

  GrapplerItem item;
  TF_CHECK_OK(s.ToGraphDef(&item.graph));
  EXPECT_EQ(6, item.graph.node_size());
  NodeMap pre_transform_node_map(&item.graph);
  (*pre_transform_node_map.GetNode("b")->mutable_attr())["_recompute_hint"]
      .set_i(0);

  MemoryOptimizer optimizer(RewriterConfig::MANUAL);
  GraphDef output;
  Status status = optimizer.Optimize(nullptr, item, &output);

  TF_EXPECT_OK(status);
  NodeMap post_transform_node_map(&output);
  EXPECT_EQ(8, output.node_size());
  NodeDef* transformed_e = post_transform_node_map.GetNode(e.name());
  EXPECT_EQ(2, transformed_e->input_size());
  EXPECT_EQ("gradients/d", transformed_e->input(0));
  EXPECT_EQ("Recomputed/b", transformed_e->input(1));
  NodeDef* recomputed_b = post_transform_node_map.GetNode("Recomputed/b");
  EXPECT_EQ(2, recomputed_b->input_size());
  EXPECT_EQ("a", recomputed_b->input(0));
  EXPECT_EQ("^RecomputeTrigger/b", recomputed_b->input(1));
  NodeDef* recompute_trigger =
      post_transform_node_map.GetNode("RecomputeTrigger/b");
  EXPECT_EQ(1, recompute_trigger->input_size());
  EXPECT_EQ("^gradients/d", recompute_trigger->input(0));
}

TEST_F(RecomputeSubgraphTest, NoFeedsRecomputed) {
  tensorflow::Scope s = tensorflow::Scope::NewRootScope();

  Output a = ops::Variable(s.WithOpName("a"), {2, 3, 4}, DT_FLOAT);
  Output b = ops::Identity(s.WithOpName("b"), a);  // Would be recomputed, but
                                                   // for being fed
  Output c = ops::Identity(s.WithOpName("c"), b);
  Output d = ops::AddN(s.WithOpName("gradients/d"), {c});
  Output e = ops::AddN(s.WithOpName("gradients/e"), {d, b});
  Output f = ops::AddN(s.WithOpName("gradients/f"), {e, a});

  GrapplerItem item;
  TF_CHECK_OK(s.ToGraphDef(&item.graph));
  item.feed.emplace_back("b", Tensor());
  EXPECT_EQ(6, item.graph.node_size());
  NodeMap pre_transform_node_map(&item.graph);
  (*pre_transform_node_map.GetNode("b")->mutable_attr())["_recompute_hint"]
      .set_i(0);

  MemoryOptimizer optimizer(RewriterConfig::MANUAL);
  GraphDef output;
  Status status = optimizer.Optimize(nullptr, item, &output);

  TF_EXPECT_OK(status);
  EXPECT_EQ(6, output.node_size());
}

TEST_F(RecomputeSubgraphTest, TwoInputSubgraphs) {
  tensorflow::Scope s = tensorflow::Scope::NewRootScope();

  Output a = ops::Variable(s.WithOpName("a"), {2, 3, 4}, DT_FLOAT);
  Output b = ops::Variable(s.WithOpName("b"), {2, 3, 4}, DT_FLOAT);
  Output d = ops::AddN(
      s.WithOpName("some_name_scope/gradients/two_subgraph_inputs"), {a, b});

  GrapplerItem item;
  TF_CHECK_OK(s.ToGraphDef(&item.graph));
  EXPECT_EQ(3, item.graph.node_size());
  NodeMap pre_transform_node_map(&item.graph);
  (*pre_transform_node_map.GetNode("a")->mutable_attr())["_recompute_hint"]
      .set_i(0);
  (*pre_transform_node_map.GetNode("b")->mutable_attr())["_recompute_hint"]
      .set_i(0);

  MemoryOptimizer optimizer(RewriterConfig::MANUAL,
                            "some_name_scope/gradients");
  GraphDef output;
  Status status = optimizer.Optimize(nullptr, item, &output);

  TF_EXPECT_OK(status);
  NodeMap post_transform_node_map(&output);
  // Mostly checking that this case does not crash.
  EXPECT_EQ(7, output.node_size());
  EXPECT_NE(post_transform_node_map.GetNode("Recomputed/a"), nullptr);
  EXPECT_NE(post_transform_node_map.GetNode("Recomputed/b"), nullptr);
  EXPECT_NE(post_transform_node_map.GetNode("RecomputeTrigger/a"), nullptr);
  EXPECT_NE(post_transform_node_map.GetNode("RecomputeTrigger/b"), nullptr);
}

TEST_F(RecomputeSubgraphTest, MultiNode) {
  tensorflow::Scope s = tensorflow::Scope::NewRootScope();

  Output a = ops::Variable(s.WithOpName("Conv"), {2, 3, 4}, DT_FLOAT);
  Output b = ops::Identity(s.WithOpName("BN"), a);    // Recomputed
  Output c = ops::Identity(s.WithOpName("ReLU"), b);  // Recomputed
  Output d = ops::Identity(s.WithOpName("Conv1"), c);

  // The "gradients/" prefix means the heuristic will pick these up as
  // candidates to have their inputs recomputed.
  Output trigger = ops::AddN(s.WithOpName("gradients/BN1Grad"), {d});
  Output e = ops::AddN(s.WithOpName("gradients/Conv1Grad"), {trigger, c});
  Output f = ops::AddN(s.WithOpName("gradients/ReLUGrad"), {e, c});
  Output g = ops::AddN(s.WithOpName("gradients/BNGrad"), {f, a});
  Output h = ops::AddN(s.WithOpName("gradients/ConvGrad"), {g});

  GrapplerItem item;
  TF_CHECK_OK(s.ToGraphDef(&item.graph));
  EXPECT_EQ(9, item.graph.node_size());
  NodeMap pre_transform_node_map(&item.graph);
  // Set op types so that the heuristic will pick these nodes up to be
  // recomputed
  pre_transform_node_map.GetNode("BN")->set_op("FusedBatchNorm");
  pre_transform_node_map.GetNode("ReLU")->set_op("Relu");

  MemoryOptimizer optimizer(RewriterConfig::RECOMPUTATION_HEURISTICS);
  GraphDef first_pass_output;
  Status first_pass_status =
      optimizer.Optimize(nullptr, item, &first_pass_output);
  TF_EXPECT_OK(first_pass_status);

  NodeMap post_transform_node_map(&first_pass_output);
  EXPECT_EQ(13, first_pass_output.node_size());
  NodeDef* transformed_e = post_transform_node_map.GetNode(e.name());
  EXPECT_EQ(2, transformed_e->input_size());
  EXPECT_EQ("gradients/BN1Grad", transformed_e->input(0));
  EXPECT_EQ("Recomputed/ReLU", transformed_e->input(1));
  NodeDef* transformed_f = post_transform_node_map.GetNode(f.name());
  EXPECT_EQ(2, transformed_f->input_size());
  EXPECT_EQ("gradients/Conv1Grad", transformed_f->input(0));
  EXPECT_EQ("Recomputed/ReLU", transformed_f->input(1));
  NodeDef* transformed_g = post_transform_node_map.GetNode(g.name());
  EXPECT_EQ(2, transformed_g->input_size());
  EXPECT_EQ("gradients/ReLUGrad", transformed_g->input(0));
  EXPECT_EQ("Conv", transformed_g->input(1));

  NodeDef* recomputed_b = post_transform_node_map.GetNode("Recomputed/BN");
  EXPECT_EQ(2, recomputed_b->input_size());
  EXPECT_EQ("Conv", recomputed_b->input(0));
  EXPECT_EQ("^RecomputeTrigger/BN", recomputed_b->input(1));
  NodeDef* recompute_trigger_b =
      post_transform_node_map.GetNode("RecomputeTrigger/BN");
  EXPECT_EQ(1, recompute_trigger_b->input_size());
  EXPECT_EQ("^RecomputeTrigger/ReLU", recompute_trigger_b->input(0));

  NodeDef* recomputed_c = post_transform_node_map.GetNode("Recomputed/ReLU");
  EXPECT_EQ(2, recomputed_c->input_size());
  EXPECT_EQ("Recomputed/BN", recomputed_c->input(0));
  EXPECT_EQ("^RecomputeTrigger/ReLU", recomputed_c->input(1));
  NodeDef* recompute_trigger_c =
      post_transform_node_map.GetNode("RecomputeTrigger/ReLU");
  EXPECT_EQ(1, recompute_trigger_c->input_size());
  EXPECT_EQ("^gradients/BN1Grad", recompute_trigger_c->input(0));
}

class MemoryOptimizerTest : public ::testing::Test {
 public:
  static std::unique_ptr<VirtualCluster> CreateVirtualCluster() {
    DeviceProperties cpu_device;
    cpu_device.set_type("CPU");
    cpu_device.set_frequency(1000);
    cpu_device.set_num_cores(4);
    cpu_device.set_bandwidth(32);
    DeviceProperties gpu_device;
    gpu_device.set_type("GPU");
    gpu_device.set_frequency(1000);
    gpu_device.set_num_cores(24);
    gpu_device.set_bandwidth(128);
    gpu_device.set_memory_size(1024 * 1024);
    gpu_device.mutable_environment()->insert({"architecture", "6"});
    std::unordered_map<string, DeviceProperties> devices;
    devices["/job:localhost/replica:0/task:0/cpu:0"] = cpu_device;
    devices["/job:localhost/replica:0/task:0/gpu:0"] = gpu_device;
    return std::unique_ptr<VirtualCluster>(new VirtualCluster(devices));
  }
};

TEST_F(MemoryOptimizerTest, SimpleSwapping) {
  // Build a simple graph with an op that's marked for swapping.
  tensorflow::Scope s = tensorflow::Scope::NewRootScope();

  Output a = ops::Variable(s.WithOpName("a"), {10, 10}, DT_FLOAT);
  Output b = ops::AddN(s.WithOpName("b"), {a});
  Output c = ops::AddN(s.WithOpName("c"), {b});
  Output d = ops::AddN(s.WithOpName("d"), {c});
  Output e = ops::AddN(s.WithOpName("e"), {b, d});

  GrapplerItem item;
  TF_CHECK_OK(s.ToGraphDef(&item.graph));

  EXPECT_EQ(5, item.graph.node_size());
  EXPECT_EQ(NodeName(e.name()), item.graph.node(4).name());
  AttrValue& val =
      (*item.graph.mutable_node(4)->mutable_attr())["_swap_to_host"];
  val.mutable_list()->add_i(0);

  std::unique_ptr<VirtualCluster> cluster(CreateVirtualCluster());

  MemoryOptimizer optimizer(RewriterConfig::MANUAL);
  GraphDef output;
  Status status = optimizer.Optimize(cluster.get(), item, &output);
  TF_EXPECT_OK(status);

  EXPECT_EQ(7, output.node_size());
  const NodeDef& new_e = output.node(4);
  EXPECT_EQ(NodeName(e.name()), new_e.name());

  EXPECT_EQ(2, new_e.input_size());
  EXPECT_EQ(NodeName(d.name()), new_e.input(1));
  EXPECT_EQ("swap_in_e_0", new_e.input(0));

  const NodeDef& swap_out = output.node(5);
  EXPECT_EQ("swap_out_e_0", swap_out.name());

  const NodeDef& swap_in = output.node(6);
  EXPECT_EQ("swap_in_e_0", swap_in.name());

  EXPECT_EQ(NodeName(b.name()), swap_out.input(0));
  EXPECT_EQ(NodeName(swap_out.name()), swap_in.input(0));
  EXPECT_EQ("^c", swap_in.input(1));

  const NodeDef& new_c = output.node(2);
  EXPECT_EQ(NodeName(c.name()), new_c.name());
  EXPECT_EQ("^swap_out_e_0", new_c.input(1));

  // Run the optimizer a second time to ensure it's idempotent.
  item.graph.Swap(&output);
  status = optimizer.Optimize(cluster.get(), item, &output);
  TF_EXPECT_OK(status);
}

TEST_F(MemoryOptimizerTest, SwappingHeuristics) {
  tensorflow::Scope s = tensorflow::Scope::NewRootScope();
  Output v = ops::Variable(s.WithOpName("v").WithDevice("/gpu:0"),
                           {128, 128, 8}, DT_FLOAT);
  Output a = ops::Identity(s.WithOpName("a").WithDevice("/gpu:0"), v);
  Output b = ops::Square(s.WithOpName("b").WithDevice("/gpu:0"), v);
  Output c = ops::Sqrt(s.WithOpName("c").WithDevice("/gpu:0"), a);
  Output d = ops::Identity(s.WithOpName("d").WithDevice("/gpu:0"), b);
  Output axis = ops::Const(s.WithOpName("axis"), 0);
  Output e =
      ops::Concat(s.WithOpName("e").WithDevice("/gpu:0"), {a, b, c, d}, axis);
  Output f = ops::Square(s.WithOpName("f").WithDevice("/gpu:0"), a);
  Output g = ops::Sqrt(s.WithOpName("g").WithDevice("/gpu:0"), b);
  Output h = ops::Exp(s.WithOpName("h").WithDevice("/gpu:0"), c);
  Output i = ops::Log(s.WithOpName("i").WithDevice("/gpu:0"), d);

  GrapplerItem item;
  TF_CHECK_OK(s.ToGraphDef(&item.graph));
  item.fetch = {"e", "f", "g", "h", "i"};

  std::unique_ptr<VirtualCluster> cluster(CreateVirtualCluster());

  MemoryOptimizer optimizer(RewriterConfig::SWAPPING_HEURISTICS);
  GraphDef output;
  Status status = optimizer.Optimize(cluster.get(), item, &output);
  TF_EXPECT_OK(status);

  for (const auto& node : output.node()) {
    if (node.name() == "e") {
      EXPECT_EQ(5, node.input_size());
      EXPECT_EQ("a", node.input(0));
      EXPECT_EQ("swap_in_e_1", node.input(1));
      EXPECT_EQ("swap_in_e_2", node.input(2));
      EXPECT_EQ("swap_in_e_3", node.input(3));
      EXPECT_EQ("axis", node.input(4));
    }
  }
}

TEST_F(MemoryOptimizerTest, UnswappableInputs) {
  tensorflow::Scope s = tensorflow::Scope::NewRootScope();
  Output v = ops::Variable(s.WithOpName("v").WithDevice("/gpu:0"),
                           {128, 128, 8}, DT_FLOAT);
  Output a = ops::Square(s.WithOpName("a").WithDevice("/gpu:0"), v);
  Output b = ops::Identity(s.WithOpName("b").WithDevice("/gpu:0"), {a});
  Output c = ops::Identity(s.WithOpName("c").WithDevice("/gpu:0"), {a});
  Output index = ops::Const(s.WithOpName("index"), {0});
  Output indices = ops::Tile(s.WithOpName("indices"), index, {128});
  Output d =
      ops::ScatterAdd(s.WithOpName("d").WithDevice("/gpu:0"), v, indices, c);
  Output axis = ops::Const(s.WithOpName("axis"), 0);
  Output e =
      ops::Concat(s.WithOpName("e").WithDevice("/gpu:0"), {b, c, d}, axis);

  GrapplerItem item;
  TF_CHECK_OK(s.ToGraphDef(&item.graph));
  item.fetch = {"e"};

  std::unique_ptr<VirtualCluster> cluster(CreateVirtualCluster());

  MemoryOptimizer optimizer(RewriterConfig::SWAPPING_HEURISTICS);
  GraphDef output;
  Status status = optimizer.Optimize(cluster.get(), item, &output);
  TF_EXPECT_OK(status);

  for (const auto& node : output.node()) {
    if (node.name() == "e") {
      // The d node isn't swappable.
      EXPECT_EQ(5, node.input_size());
      EXPECT_EQ("d", node.input(2));
      EXPECT_EQ("^swap_out_d_2", node.input(4));
    }
  }
}

TEST_F(MemoryOptimizerTest, AccumulationRewrites) {
  tensorflow::Scope s = tensorflow::Scope::NewRootScope();
  Output a = ops::Variable(s.WithOpName("a").WithDevice("/gpu:0"),
                           {128, 128, 8}, DT_FLOAT);
  Output b = ops::Variable(s.WithOpName("b").WithDevice("/gpu:0"),
                           {128, 128, 8}, DT_FLOAT);
  Output c = ops::Variable(s.WithOpName("c").WithDevice("/gpu:0"),
                           {128, 128, 8}, DT_FLOAT);
  Output d = ops::AddN(s.WithOpName("d").WithDevice("/gpu:0"), {a, b, c});

  GrapplerItem item;
  TF_CHECK_OK(s.ToGraphDef(&item.graph));
  item.fetch = {"d"};

  std::unique_ptr<VirtualCluster> cluster(CreateVirtualCluster());
  MemoryOptimizer optimizer(RewriterConfig::SCHEDULING_HEURISTICS);
  GraphDef output;
  Status status = optimizer.Optimize(cluster.get(), item, &output);
  TF_EXPECT_OK(status);

  int count = 0;
  for (const auto& node : output.node()) {
    if (node.name() == "d") {
      EXPECT_EQ("DestroyTemporaryVariable", node.op());
      count++;
    } else if (node.name() == "d/tmp_var_initializer") {
      EXPECT_EQ("Assign", node.op());
      count++;
    } else if (node.name() == "d/tmp_var") {
      EXPECT_EQ("TemporaryVariable", node.op());
      count++;
    }
  }
  EXPECT_EQ(3, count);
}

}  // namespace
}  // namespace grappler
}  // namespace tensorflow


<|start_filename|>tensorflow/contrib/lite/kernels/mean_test.cc<|end_filename|>
/* 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.
==============================================================================*/
#include <gtest/gtest.h>
#include "tensorflow/contrib/lite/interpreter.h"
#include "tensorflow/contrib/lite/kernels/register.h"
#include "tensorflow/contrib/lite/kernels/test_util.h"
#include "tensorflow/contrib/lite/model.h"

namespace tflite {
namespace {

using ::testing::ElementsAreArray;

class BaseMeanOpModel : public SingleOpModel {
 public:
  BaseMeanOpModel(const TensorData& input, const TensorData& output,
                  std::initializer_list<int> axis, bool keep_dims) {
    input_ = AddInput(input);
    output_ = AddOutput(output);
    SetBuiltinOp(
        BuiltinOperator_MEAN, BuiltinOptions_MeanOptions,
        CreateMeanOptions(builder_, builder_.CreateVector<int>(axis), keep_dims)
            .Union());
    BuildInterpreter({GetShape(input_)});
  }

  int input() { return input_; }

 protected:
  int input_;
  int output_;
};

class FloatMeanOpModel : public BaseMeanOpModel {
 public:
  using BaseMeanOpModel::BaseMeanOpModel;

  void SetInput(std::initializer_list<float> data) {
    PopulateTensor(input_, data);
  }

  std::vector<float> GetOutput() { return ExtractVector<float>(output_); }
  std::vector<int> GetOutputShape() { return GetTensorShape(output_); }
};

TEST(FloatMeanOpTest, NotKeepDims) {
  std::initializer_list<float> data = {
      1.0,  2.0,  3.0,  4.0,  5.0,  6.0,  7.0,  8.0,  9.0,  10.0, 11.0, 12.0,
      13.0, 14.0, 15.0, 16.0, 17.0, 18.0, 19.0, 20.0, 21.0, 22.0, 23.0, 24.0};
  FloatMeanOpModel m({TensorType_FLOAT32, {4, 3, 2}}, {TensorType_FLOAT32, {2}},
                     {1, 0, -3, -3}, false);
  m.SetInput(data);
  m.Invoke();
  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({2}));
  EXPECT_THAT(m.GetOutput(), ElementsAreArray(ArrayFloatNear({12, 13})));
}

TEST(FloatMeanOpTest, KeepDims) {
  std::initializer_list<float> data = {
      1.0,  2.0,  3.0,  4.0,  5.0,  6.0,  7.0,  8.0,  9.0,  10.0, 11.0, 12.0,
      13.0, 14.0, 15.0, 16.0, 17.0, 18.0, 19.0, 20.0, 21.0, 22.0, 23.0, 24.0};
  FloatMeanOpModel m({TensorType_FLOAT32, {4, 3, 2}}, {TensorType_FLOAT32, {3}},
                     {0, 2}, true);
  m.SetInput(data);
  m.Invoke();
  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 3, 1}));
  EXPECT_THAT(m.GetOutput(),
              ElementsAreArray(ArrayFloatNear({10.5, 12.5, 14.5})));
}

}  // namespace
}  // namespace tflite

int main(int argc, char** argv) {
  ::tflite::LogToStderr();
  ::testing::InitGoogleTest(&argc, argv);
  return RUN_ALL_TESTS();
}


<|start_filename|>tensorflow/core/framework/dataset.h<|end_filename|>
/* 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.
==============================================================================*/

#ifndef TENSORFLOW_FRAMEWORK_DATASET_H_
#define TENSORFLOW_FRAMEWORK_DATASET_H_

namespace tensorflow {
namespace dataset {
// Registry for stateful ops that need to be used in dataset functions.
// See below macro for usage details.
class WhitelistedStatefulOpRegistry {
 public:
  Status Add(StringPiece op_name) {
    op_names_.insert(op_name);
    return Status::OK();
  }

  bool Contains(StringPiece op_name) {
    return op_names_.find(op_name) != op_names_.end();
  }

  static WhitelistedStatefulOpRegistry* Global() {
    static WhitelistedStatefulOpRegistry* reg =
        new WhitelistedStatefulOpRegistry;
    return reg;
  }

 private:
  WhitelistedStatefulOpRegistry() {}
  WhitelistedStatefulOpRegistry(WhitelistedStatefulOpRegistry const& copy);
  WhitelistedStatefulOpRegistry operator=(
      WhitelistedStatefulOpRegistry const& copy);
  std::set<StringPiece> op_names_;
};

}  // namespace dataset

// Use this macro to whitelist an op that is marked stateful but needs to be
// used inside a map_fn in an input pipeline. This is only needed if you wish
// to be able to checkpoint the state of the input pipeline. We currently
// do not allow stateful ops to be defined inside of map_fns since it is not
// possible to save their state.
// Note that the state of the whitelisted ops inside functions will not be
// saved during checkpointing, hence this should only be used if the op is
// marked stateful for reasons like to avoid constant folding during graph
// optimiztion but is not stateful.
// If possible, try to remove the stateful flag on the op first.
// Example usage:
//
//   WHITELIST_STATEFUL_OP_FOR_DATASET_FUNCTIONS("LegacyStatefulReader");
//
#define WHITELIST_STATEFUL_OP_FOR_DATASET_FUNCTIONS(name) \
  WHITELIST_STATEFUL_OP_FOR_DATASET_FUNCTIONS_UNIQ_HELPER(__COUNTER__, name)
#define WHITELIST_STATEFUL_OP_FOR_DATASET_FUNCTIONS_UNIQ_HELPER(ctr, name) \
  WHITELIST_STATEFUL_OP_FOR_DATASET_FUNCTIONS_UNIQ(ctr, name)
#define WHITELIST_STATEFUL_OP_FOR_DATASET_FUNCTIONS_UNIQ(ctr, name)        \
  static ::tensorflow::Status whitelist_op##ctr TF_ATTRIBUTE_UNUSED =      \
      ::tensorflow::dataset::WhitelistedStatefulOpRegistry::Global()->Add( \
          name)

}  // namespace tensorflow

#endif  // TENSORFLOW_FRAMEWORK_DATASET_H_


<|start_filename|>tensorflow/core/ops/resource_variable_ops.cc<|end_filename|>
// Copyright 2016 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.
// ============================================================================

#include "tensorflow/core/framework/common_shape_fns.h"
#include "tensorflow/core/framework/node_def_util.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/resource_mgr.h"
#include "tensorflow/core/framework/shape_inference.h"

using ::tensorflow::shape_inference::InferenceContext;
using ::tensorflow::shape_inference::ShapeAndType;
using ::tensorflow::shape_inference::ShapeHandle;

namespace tensorflow {

namespace {

Status ValidateVariableResourceHandle(InferenceContext* c,
                                      ShapeAndType* shape_and_type) {
  auto* handle_data = c->input_handle_shapes_and_types(0);
  if (handle_data == nullptr || handle_data->empty()) {
    shape_and_type->shape = c->UnknownShape();
    shape_and_type->dtype = DT_INVALID;
  } else {
    *shape_and_type = (*handle_data)[0];
    DataType value_dtype;
    TF_RETURN_IF_ERROR(c->GetAttr("dtype", &value_dtype));
    if (shape_and_type->dtype != value_dtype) {
      return errors::InvalidArgument(
          "Trying to read variable with wrong dtype. "
          "Expected ",
          DataTypeString(shape_and_type->dtype), " got ",
          DataTypeString(value_dtype));
    }
  }
  return Status::OK();
}

Status ReadVariableShapeFn(InferenceContext* c) {
  ShapeAndType shape_and_type;
  TF_RETURN_IF_ERROR(ValidateVariableResourceHandle(c, &shape_and_type));
  c->set_output(0, shape_and_type.shape);
  return Status::OK();
}

}  // namespace

REGISTER_OP("VarHandleOp")
    .Attr("container: string = ''")
    .Attr("shared_name: string = ''")
    .Attr("dtype: type")
    .Attr("shape: shape")
    .Output("resource: resource")
    .SetIsStateful()
    .SetShapeFn([](InferenceContext* c) {
      c->set_output(0, c->Scalar());
      DataType t;
      TF_RETURN_IF_ERROR(c->GetAttr("dtype", &t));
      PartialTensorShape p;
      TF_RETURN_IF_ERROR(c->GetAttr("shape", &p));
      ShapeHandle s;
      TF_RETURN_IF_ERROR(c->MakeShapeFromPartialTensorShape(p, &s));
      c->set_output_handle_shapes_and_types(0,
                                            std::vector<ShapeAndType>{{s, t}});

      return Status::OK();
    });

REGISTER_OP("ReadVariableOp")
    .Input("resource: resource")
    .Output("value: dtype")
    .Attr("dtype: type")
    .SetShapeFn(ReadVariableShapeFn);

REGISTER_OP("DestroyResourceOp")
    .Input("resource: resource")
    .Attr("ignore_lookup_error: bool = true")
    .SetIsStateful()
    .SetShapeFn(shape_inference::NoOutputs);

Status CreateAssignShapeFn(InferenceContext* c) {
  ShapeAndType handle_shape_and_type;
  TF_RETURN_IF_ERROR(ValidateVariableResourceHandle(c, &handle_shape_and_type));

  ShapeHandle value_shape = c->input(1);
  ShapeHandle unused;
  TF_RETURN_IF_ERROR(
      c->Merge(handle_shape_and_type.shape, value_shape, &unused));
  return Status::OK();
}

REGISTER_OP("AssignVariableOp")
    .Input("resource: resource")
    .Input("value: dtype")
    .Attr("dtype: type")
    .SetShapeFn(CreateAssignShapeFn);

REGISTER_OP("AssignAddVariableOp")
    .Input("resource: resource")
    .Input("value: dtype")
    .Attr("dtype: type")
    .SetShapeFn(CreateAssignShapeFn);

REGISTER_OP("AssignSubVariableOp")
    .Input("resource: resource")
    .Input("value: dtype")
    .Attr("dtype: type")
    .SetShapeFn(CreateAssignShapeFn);

REGISTER_OP("VarIsInitializedOp")
    .Input("resource: resource")
    .Output("is_initialized: bool")
    .SetShapeFn(tensorflow::shape_inference::ScalarShape);

Status VariableShapeShapeFn(InferenceContext* c) {
  auto* handle_data = c->input_handle_shapes_and_types(0);
  if (handle_data == nullptr || handle_data->empty()) {
    return errors::InvalidArgument("Handle doesn't have shape information.");
  }
  ShapeHandle var_shape = (*handle_data)[0].shape;
  int64 rank = c->RankKnown(var_shape) ? c->Rank(var_shape)
                                       : InferenceContext::kUnknownDim;
  c->set_output(0, c->Vector(rank));
  return Status::OK();
}

REGISTER_OP("VariableShape")
    .Input("input: resource")
    .Output("output: out_type")
    .Attr("out_type: {int32, int64} = DT_INT32")
    .SetShapeFn(VariableShapeShapeFn);

REGISTER_OP("ResourceGather")
    .Input("resource: resource")
    .Input("indices: Tindices")
    .Attr("validate_indices: bool = true")
    .Output("output: dtype")
    .Attr("dtype: type")
    .Attr("Tindices: {int32,int64}")
    .SetShapeFn([](InferenceContext* c) {
      ShapeAndType handle_shape_and_type;
      TF_RETURN_IF_ERROR(
          ValidateVariableResourceHandle(c, &handle_shape_and_type));

      ShapeHandle unused;
      TF_RETURN_IF_ERROR(
          c->WithRankAtLeast(handle_shape_and_type.shape, 1, &unused));
      ShapeHandle params_subshape;
      TF_RETURN_IF_ERROR(
          c->Subshape(handle_shape_and_type.shape, 1, &params_subshape));
      ShapeHandle indices_shape = c->input(1);
      ShapeHandle out;
      TF_RETURN_IF_ERROR(c->Concatenate(indices_shape, params_subshape, &out));
      c->set_output(0, out);
      return Status::OK();
    });

REGISTER_OP("ResourceScatterAdd")
    .Input("resource: resource")
    .Input("indices: Tindices")
    .Input("updates: dtype")
    .Attr("dtype: numbertype")
    .Attr("Tindices: {int32, int64}")
    .SetShapeFn([](InferenceContext* c) {
      ShapeAndType handle_shape_and_type;
      TF_RETURN_IF_ERROR(
          ValidateVariableResourceHandle(c, &handle_shape_and_type));
      ShapeHandle var_shape = handle_shape_and_type.shape;
      ShapeHandle indices_shape = c->input(1);

      ShapeHandle unused_updates_shape;
      ShapeHandle concat;
      ShapeHandle var_subshape;
      TF_RETURN_IF_ERROR(c->Subshape(var_shape, 1, &var_subshape));
      TF_RETURN_IF_ERROR(c->Concatenate(indices_shape, var_subshape, &concat));
      TF_RETURN_IF_ERROR(c->Merge(c->input(2), concat, &unused_updates_shape));
      return Status::OK();
    });

REGISTER_OP("ResourceScatterUpdate")
    .Input("resource: resource")
    .Input("indices: Tindices")
    .Input("updates: dtype")
    .Attr("dtype: numbertype")
    .Attr("Tindices: {int32, int64}")
    .SetShapeFn([](InferenceContext* c) {
      ShapeAndType handle_shape_and_type;
      TF_RETURN_IF_ERROR(
          ValidateVariableResourceHandle(c, &handle_shape_and_type));
      ShapeHandle var_shape = handle_shape_and_type.shape;
      ShapeHandle indices_shape = c->input(1);

      ShapeHandle unused_updates_shape;
      ShapeHandle concat;
      ShapeHandle var_subshape;
      TF_RETURN_IF_ERROR(c->Subshape(var_shape, 1, &var_subshape));
      TF_RETURN_IF_ERROR(c->Concatenate(indices_shape, var_subshape, &concat));
      TF_RETURN_IF_ERROR(c->Merge(c->input(2), concat, &unused_updates_shape));
      return Status::OK();
    });

REGISTER_OP("CriticalSectionOp")
    .Attr("container: string = ''")
    .Attr("shared_name: string = ''")
    .Output("resource: resource")
    .SetIsStateful()
    .SetShapeFn([](InferenceContext* c) {
      c->set_output(0, c->Scalar());
      return Status::OK();
    });

REGISTER_OP("ExecuteInCriticalSection")
    .Input("critical_section: resource")
    .Input("arguments: Targuments")
    .Output("outputs: output_types")
    .Attr("f: func")
    .Attr("Targuments: list(type) >= 0")
    .Attr("output_types: list(type) >= 0")
    .Attr("output_shapes: list(shape) >= 0")
    .SetShapeFn([](InferenceContext* c) {
      std::vector<PartialTensorShape> output_shapes;
      TF_RETURN_IF_ERROR(c->GetAttr("output_shapes", &output_shapes));
      for (int i = 0; i < output_shapes.size(); ++i) {
        ShapeHandle s;
        TF_RETURN_IF_ERROR(
            c->MakeShapeFromPartialTensorShape(output_shapes[i], &s));
        c->set_output(i, s);
      }
      return Status::OK();
    });

}  // namespace tensorflow
