chore: import upstream snapshot with attribution

tensorflow/c/kernels/BUILD

@@ -0,0 +1,225 @@
+load("//tensorflow:tensorflow.bzl", "tf_cc_test")
+load("//tensorflow:tensorflow.default.bzl", "filegroup", "tf_gen_op_libs", "tf_kernel_library")
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = ["//visibility:public"],
+    licenses = ["notice"],
+)
+
+tf_kernel_library(
+    name = "bitcast_op",
+    prefix = "bitcast_op",
+    deps = [
+        "//tensorflow/c:kernels",
+        "//tensorflow/c:ops",
+        "//tensorflow/c:tf_datatype",
+        "//tensorflow/c:tf_status",
+        "//tensorflow/c:tf_tensor",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:lib",
+        "@com_google_absl//absl/log:check",
+    ],
+)
+
+tf_kernel_library(
+    name = "summary_op",
+    prefix = "summary_op",
+    deps = [
+        ":tensor_shape_utils",
+        "//tensorflow/c:kernels",
+        "//tensorflow/c:tf_datatype_hdrs",
+        "//tensorflow/c:tf_status",
+        "//tensorflow/c:tf_tensor",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:protos_all_cc",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/strings",
+        "@eigen_archive//:eigen3",
+    ],
+)
+
+tf_kernel_library(
+    name = "histogram_summary_op",
+    prefix = "histogram_summary_op",
+    deps = [
+        "//tensorflow/c:kernels",
+        "//tensorflow/c:tf_status",
+        "//tensorflow/c:tf_tensor",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:protos_all_cc",
+        "@com_google_absl//absl/log:check",
+        "@eigen_archive//:eigen3",
+    ],
+)
+
+tf_kernel_library(
+    name = "merge_summary_op",
+    prefix = "merge_summary_op",
+    deps = [
+        "//tensorflow/c:kernels",
+        "//tensorflow/c:tf_status",
+        "//tensorflow/c:tf_tensor",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:protos_all_cc",
+        "@com_google_absl//absl/log:check",
+    ],
+)
+
+tf_gen_op_libs(
+    op_lib_names = ["bitcast"],
+    deps = [
+        "//tensorflow/c:ops",
+        "//tensorflow/c:tf_datatype",
+        "//tensorflow/c:tf_status",
+        "//tensorflow/c:tf_tensor",
+        "//tensorflow/core:lib",
+    ],
+)
+
+tf_gen_op_libs(
+    op_lib_names = ["summary"],
+    deps = [
+        "//tensorflow/c:ops",
+        "//tensorflow/c:tf_status",
+        "//tensorflow/core:lib",
+    ],
+)
+
+tf_gen_op_libs(
+    op_lib_names = ["histogram_summary"],
+    deps = [
+        "//tensorflow/c:ops",
+        "//tensorflow/c:tf_status",
+        "//tensorflow/core:lib",
+    ],
+)
+
+tf_gen_op_libs(
+    op_lib_names = ["merge_summary"],
+    deps = [
+        "//tensorflow/c:ops",
+        "//tensorflow/c:tf_status",
+        "//tensorflow/core:lib",
+    ],
+)
+
+tf_cc_test(
+    name = "bitcast_op_test",
+    srcs = ["bitcast_op_test.cc"],
+    deps = [
+        ":bitcast_op",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+        "//tensorflow/core:testlib",
+        "@com_google_absl//absl/container:inlined_vector",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
+        "@xla//xla/tsl/protobuf:error_codes_proto_impl_cc",
+    ],
+)
+
+tf_cc_test(
+    name = "summary_op_test",
+    srcs = ["summary_op_test.cc"],
+    deps = [
+        ":summary_op",
+        "//tensorflow/c:kernels",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+        "//tensorflow/core:testlib",
+    ],
+)
+
+tf_cc_test(
+    name = "summary_op_benchmark_test",
+    size = "small",
+    srcs = ["summary_op_benchmark_test.cc"],
+    deps = [
+        ":summary_op",
+        "//tensorflow/c:kernels",
+        "//tensorflow/core:core_cpu",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+        "//tensorflow/core:testlib",
+    ],
+)
+
+cc_library(
+    name = "tensor_shape_utils",
+    srcs = ["tensor_shape_utils.cc"],
+    hdrs = ["tensor_shape_utils.h"],
+    visibility = ["//visibility:private"],
+    deps = [
+        "//tensorflow/c:tf_tensor",
+        "//tensorflow/core:lib",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/strings",
+    ],
+)
+
+tf_cc_test(
+    name = "tensor_shape_utils_test",
+    srcs = ["tensor_shape_utils_test.cc"],
+    deps = [
+        ":tensor_shape_utils",
+        "//tensorflow/c:tf_tensor_internal",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+        "@com_google_absl//absl/status",
+    ],
+)
+
+# Be aware that this target is included by
+# //third_party/tensorflow/core/kernels:portable_core_ops
+#
+filegroup(
+    name = "portable_all_op_kernels",
+    srcs = [
+        ":portable_all_op_kernel_hdrs",
+        ":portable_all_op_kernel_srcs",
+    ],
+)
+
+filegroup(
+    name = "portable_all_op_kernel_srcs",
+    srcs = [
+        "bitcast_op.cc",
+        "histogram_summary_op.cc",
+        "merge_summary_op.cc",
+        "summary_op.cc",
+        "tensor_shape_utils.cc",
+    ],
+)
+
+filegroup(
+    name = "portable_all_op_kernel_hdrs",
+    srcs = [
+        "tensor_shape_utils.h",
+    ],
+)
+
+filegroup(
+    name = "portable_all_ops",
+    srcs = [
+        "ops/bitcast.cc",
+        "ops/histogram_summary.cc",
+        "ops/merge_summary.cc",
+        "ops/summary.cc",
+    ],
+)

tensorflow/c/kernels/bitcast_op.cc

@@ -0,0 +1,176 @@
+/* Copyright 2019 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 <algorithm>
+#include <cstddef>
+#include <cstdint>
+#include <sstream>
+
+#include "absl/log/check.h"
+#include "tensorflow/c/kernels.h"
+#include "tensorflow/c/ops.h"
+#include "tensorflow/c/tf_tensor.h"
+#include "tensorflow/core/framework/common_shape_fns.h"
+#include "tensorflow/core/framework/op.h"
+#include "tensorflow/core/framework/registration/registration.h"
+#include "tensorflow/core/framework/shape_inference.h"
+#include "tensorflow/core/platform/macros.h"
+
+// BitcastOp implements a bitcast kernel, creating an output tensor that shares
+// the same data buffer as the input but with a different shape and/or data
+// type. Its inputs are:
+//
+//   * the input tensor
+//   * an attribute named "T" containing the TF_DataType of the input tensor
+//   * an attribute named "type" containing the TF_DataType of the output tensor
+//
+// Given an input tensor of shape [...], if the input DataType "T" is larger
+// than the output DataType "type", then the shape changes from [...]
+// to [..., sizeof(T)/sizeof(type)].
+//
+// If "T" is smaller than "type", the operator requires that the rightmost
+// dimension be equal to sizeof(type)/sizeof(T). The shape then goes from
+// [..., sizeof(type)/sizeof(T)] to [...].
+//
+// Bitcast is implemented as a low-level cast, so machines with different endian
+// orderings will give different results.
+typedef struct BitcastOp {
+  TF_DataType input_data_type;
+  TF_DataType output_data_type;
+  size_t in_size;
+  size_t out_size;
+} BitcastOp;
+
+static void* BitcastOp_Create(TF_OpKernelConstruction* ctx) {
+  auto* kernel = new BitcastOp;
+
+  TF_Status* s = TF_NewStatus();
+  TF_OpKernelConstruction_GetAttrType(ctx, "T", &kernel->input_data_type, s);
+
+  if (TF_GetCode(s) == TF_OK) {
+    TF_OpKernelConstruction_GetAttrType(ctx, "type", &kernel->output_data_type,
+                                        s);
+  }
+
+  if (TF_GetCode(s) == TF_OK) {
+    kernel->in_size = TF_DataTypeSize(kernel->input_data_type);
+    kernel->out_size = TF_DataTypeSize(kernel->output_data_type);
+
+    size_t check_size = std::max(kernel->in_size, kernel->out_size) %
+                        std::min(kernel->in_size, kernel->out_size);
+    if (check_size != 0) {
+      std::ostringstream err;
+      err << "cannot convert between datatype " << kernel->input_data_type
+          << " and " << kernel->output_data_type;
+      TF_SetStatus(s, TF_INVALID_ARGUMENT, err.str().c_str());
+    }
+  }
+
+  if (TF_GetCode(s) != TF_OK) {
+    TF_OpKernelConstruction_Failure(ctx, s);
+    delete kernel;
+    kernel = nullptr;
+  }
+
+  TF_DeleteStatus(s);
+  return kernel;
+}
+
+static void BitcastOp_Delete(void* kernel) {
+  delete static_cast<BitcastOp*>(kernel);
+}
+
+static void BitcastOp_Compute(void* kernel, TF_OpKernelContext* ctx) {
+  auto* k = static_cast<BitcastOp*>(kernel);
+  int dim_count = 0;
+
+  TF_Tensor* tensor = nullptr;
+  TF_Status* status = TF_NewStatus();
+  TF_GetInput(ctx, 0, &tensor, status);
+  if (TF_GetCode(status) == TF_OK) {
+    dim_count = TF_NumDims(tensor);
+    if (!(k->in_size >= k->out_size ||
+          (dim_count > 0 &&
+           TF_Dim(tensor, dim_count - 1) == k->out_size / k->in_size))) {
+      std::ostringstream err;
+      err << "Cannot bitcast from " << k->input_data_type << " to "
+          << k->output_data_type;
+      TF_SetStatus(status, TF_INVALID_ARGUMENT, err.str().c_str());
+    }
+  }
+
+  if (TF_GetCode(status) == TF_OK) {
+    auto* dims = new int64_t[dim_count + 1];
+    int new_dim_count = dim_count;
+    for (int dim = 0; dim < dim_count; ++dim) {
+      dims[dim] = TF_Dim(tensor, dim);
+    }
+    if (k->out_size < k->in_size) {
+      dims[new_dim_count++] = static_cast<int64_t>(k->in_size / k->out_size);
+    } else if (k->out_size > k->in_size) {
+      --new_dim_count;
+    }
+
+    TF_Tensor* output = TF_AllocateTensor(k->output_data_type, dims, 0,
+                                          TF_DataTypeSize(k->output_data_type));
+    TF_TensorBitcastFrom(tensor, k->output_data_type, output, dims,
+                         new_dim_count, status);
+    if (TF_GetCode(status) == TF_OK) {
+      TF_SetOutput(ctx, 0, output, status);
+    }
+    delete[] dims;
+    TF_DeleteTensor(output);
+  }
+
+  if (TF_GetCode(status) != TF_OK) {
+    TF_OpKernelContext_Failure(ctx, status);
+  }
+  TF_DeleteStatus(status);
+  TF_DeleteTensor(tensor);
+}
+
+void RegisterBitcastOpKernel() {
+  TF_Status* status = TF_NewStatus();
+  {
+    auto* builder = TF_NewKernelBuilder("Bitcast", tensorflow::DEVICE_CPU,
+                                        &BitcastOp_Create, &BitcastOp_Compute,
+                                        &BitcastOp_Delete);
+    TF_RegisterKernelBuilder("BitcastOp", builder, status);
+    CHECK_EQ(TF_OK, TF_GetCode(status))
+        << "Error while registering bitcast kernel";
+  }
+
+#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
+  {
+    auto* builder = TF_NewKernelBuilder("Bitcast", tensorflow::DEVICE_GPU,
+                                        &BitcastOp_Create, &BitcastOp_Compute,
+                                        &BitcastOp_Delete);
+    TF_RegisterKernelBuilder("BitcastOp", builder, status);
+    CHECK_EQ(TF_OK, TF_GetCode(status))
+        << "Error while registering CUDA bitcast kernel";
+  }
+#endif
+
+  TF_DeleteStatus(status);
+}
+
+// A dummy static variable initialized by a lambda whose side-effect is to
+// register the bitcast kernel.
+TF_ATTRIBUTE_UNUSED static bool IsBitcastOpKernelRegistered = []() {
+  if (SHOULD_REGISTER_OP_KERNEL("BitcastOp")) {
+    RegisterBitcastOpKernel();
+  }
+  return true;
+}();

tensorflow/c/kernels/bitcast_op_test.cc

@@ -0,0 +1,170 @@
+/* Copyright 2019 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 <cstdint>
+#include <initializer_list>
+#include <memory>
+#include <vector>
+
+#include "absl/container/inlined_vector.h"
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "xla/tsl/protobuf/error_codes.pb.h"
+#include "tensorflow/core/framework/attr_value.pb.h"
+#include "tensorflow/core/framework/attr_value_util.h"
+#include "tensorflow/core/framework/fake_input.h"
+#include "tensorflow/core/framework/node_def.pb.h"
+#include "tensorflow/core/framework/node_def_builder.h"
+#include "tensorflow/core/framework/op_def.pb.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/shape_inference.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace tensorflow {
+namespace {
+
+class DummyDevice : public DeviceBase {
+ public:
+  explicit DummyDevice(Env* env) : DeviceBase(env) {}
+  Allocator* GetAllocator(AllocatorAttributes /*attr*/) override {
+    return cpu_allocator();
+  }
+};
+
+void TestBitcastOp(Tensor* input_tensor, DataType out_type,
+                   TensorShape expected_shape, error::Code expected_code) {
+  absl::Status status;
+  NodeDef def;
+  def.set_op("Bitcast");
+  def.set_device(DEVICE_CPU);
+
+  AttrValue typeAttr;
+  SetAttrValue(input_tensor->dtype(), &typeAttr);
+
+  AttrValue outTypeAttr;
+  SetAttrValue(out_type, &outTypeAttr);
+
+  (*def.mutable_attr())["T"] = typeAttr;
+  (*def.mutable_attr())["type"] = outTypeAttr;
+
+  def.add_input(
+      absl::StrCat("input1: ", DataTypeString(input_tensor->dtype())));
+
+  std::unique_ptr<OpKernel> kernel =
+      CreateOpKernel(DeviceType(DEVICE_CPU), nullptr, nullptr, def, 1, &status);
+  ASSERT_TRUE(status.ok()) << status.ToString();
+
+  OpKernelContext::Params params;
+  DummyDevice dummy_device(nullptr);
+  params.device = &dummy_device;
+  params.op_kernel = kernel.get();
+  absl::InlinedVector<TensorValue, 4UL> inputs;
+  inputs.emplace_back(input_tensor);
+  params.inputs = inputs;
+
+  OpKernelContext ctx(&params);
+  kernel->Compute(&ctx);
+  ASSERT_EQ(expected_code, ctx.status().code());
+  if (expected_code == error::OK) {
+    ASSERT_EQ(expected_shape, ctx.mutable_output(0)->shape())
+        << ctx.mutable_output(0)->shape().DebugString();
+  }
+}
+
+TEST(BitcastOpTest, TestUpcast) {
+  Tensor int8_input(DT_UINT8, {8});
+  for (int i = 0; i < 8; i++) {
+    int8_input.vec<uint8_t>()(i) = static_cast<uint8_t>(1);
+  }
+  TestBitcastOp(&int8_input, DT_UINT64, TensorShape(), error::OK);
+}
+
+TEST(BitcastOpTest, TestDowncast) {
+  Tensor int64_input(static_cast<uint64_t>(1));
+  TestBitcastOp(&int64_input, DT_UINT8, TensorShape({8}), error::OK);
+}
+
+TEST(BitcastOpTest, TestCastToSameSize) {
+  Tensor int32_input(DT_UINT32, {4, 6});
+  TestBitcastOp(&int32_input, DT_UINT8, TensorShape({4, 6, 4}), error::OK);
+}
+
+TEST(BitcastOpTest, TestImpossibleCast) {
+  Tensor int8_input(DT_UINT8, {1});
+  TestBitcastOp(&int8_input, DT_UINT32, TensorShape(), error::INVALID_ARGUMENT);
+}
+
+PartialTensorShape S(std::initializer_list<int64_t> dims) {
+  return PartialTensorShape(dims);
+}
+
+TEST(BitcastOpTest, TestShapeInference_LargerShape) {
+  const OpRegistrationData* reg;
+  TF_CHECK_OK(OpRegistry::Global()->LookUp("Bitcast", &reg));
+  OpDef op_def = reg->op_def;
+  NodeDef def;
+  TF_CHECK_OK(NodeDefBuilder("dummy", &op_def)
+                  .Attr("type", DT_INT8)
+                  .Attr("T", DT_INT64)
+                  .Input(FakeInput(DT_INT64))
+                  .Finalize(&def));
+  shape_inference::InferenceContext c(0, def, op_def, {S({3, 4})}, {}, {}, {});
+  std::vector<shape_inference::ShapeHandle> input_shapes;
+  TF_CHECK_OK(c.input("input", &input_shapes));
+  ASSERT_EQ("[3,4]", c.DebugString(input_shapes[0]));
+  TF_CHECK_OK(reg->shape_inference_fn(&c));
+  ASSERT_EQ("[3,4,8]", c.DebugString(c.output(0)));
+}
+
+TEST(BitcastOpTest, TestShapeInference_SmallerShape) {
+  const OpRegistrationData* reg;
+  TF_CHECK_OK(OpRegistry::Global()->LookUp("Bitcast", &reg));
+  OpDef op_def = reg->op_def;
+  NodeDef def;
+  TF_CHECK_OK(NodeDefBuilder("dummy", &op_def)
+                  .Attr("type", DT_INT64)
+                  .Attr("T", DT_INT8)
+                  .Input(FakeInput(DT_INT8))
+                  .Finalize(&def));
+  shape_inference::InferenceContext c(0, def, op_def, {S({3, 4, 8})}, {}, {},
+                                      {});
+  std::vector<shape_inference::ShapeHandle> input_shapes;
+  TF_CHECK_OK(c.input("input", &input_shapes));
+  ASSERT_EQ("[3,4,8]", c.DebugString(input_shapes[0]));
+  TF_CHECK_OK(reg->shape_inference_fn(&c));
+  ASSERT_EQ("[3,4]", c.DebugString(c.output(0)));
+}
+
+TEST(BitcastOpTest, TestShapeInference_SameShape) {
+  const OpRegistrationData* reg;
+  TF_CHECK_OK(OpRegistry::Global()->LookUp("Bitcast", &reg));
+  OpDef op_def = reg->op_def;
+  NodeDef def;
+  TF_CHECK_OK(NodeDefBuilder("dummy", &op_def)
+                  .Attr("type", DT_INT32)
+                  .Attr("T", DT_FLOAT)
+                  .Input(FakeInput(DT_FLOAT))
+                  .Finalize(&def));
+  shape_inference::InferenceContext c(0, def, op_def, {S({3, 4})}, {}, {}, {});
+  std::vector<shape_inference::ShapeHandle> input_shapes;
+  TF_CHECK_OK(c.input("input", &input_shapes));
+  ASSERT_EQ("[3,4]", c.DebugString(input_shapes[0]));
+  TF_CHECK_OK(reg->shape_inference_fn(&c));
+  ASSERT_EQ("[3,4]", c.DebugString(c.output(0)));
+}
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/c/kernels/histogram_summary_op.cc

@@ -0,0 +1,168 @@
+/* Copyright 2020 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 <cstdint>
+#include <memory>
+#include <sstream>
+#include <string>
+
+#include "absl/log/check.h"
+#include "unsupported/Eigen/CXX11/Tensor"  // from @eigen_archive
+#include "tensorflow/c/kernels.h"
+#include "tensorflow/c/tf_status.h"
+#include "tensorflow/c/tf_tensor.h"
+#include "tensorflow/core/framework/registration/registration.h"
+#include "tensorflow/core/framework/summary.pb.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/lib/histogram/histogram.h"
+#include "tensorflow/core/platform/bfloat16.h"
+#include "tensorflow/core/platform/logging.h"
+#include "tensorflow/core/platform/macros.h"
+#include "tensorflow/core/platform/protobuf.h"
+#include "tensorflow/core/platform/tstring.h"
+#include "tensorflow/core/platform/types.h"
+
+namespace {
+
+// Operators used to create a std::unique_ptr for TF_Tensor and TF_Status.
+struct TFTensorDeleter {
+  void operator()(TF_Tensor* tf_tensor) const { TF_DeleteTensor(tf_tensor); }
+};
+
+struct TFStatusDeleter {
+  void operator()(TF_Status* tf_status) const { TF_DeleteStatus(tf_status); }
+};
+
+// Struct that wraps TF_Tensor and TF_Status to delete once out of scope.
+using Safe_TF_TensorPtr = std::unique_ptr<TF_Tensor, TFTensorDeleter>;
+using Safe_TF_StatusPtr = std::unique_ptr<TF_Status, TFStatusDeleter>;
+
+// Used to pass the operation node name from kernel construction to
+// kernel computation.
+struct HistogramSummaryOp {
+  std::string op_node_name;
+};
+
+void* HistogramSummaryOp_Create(TF_OpKernelConstruction* ctx) {
+  HistogramSummaryOp* kernel = new HistogramSummaryOp;
+  TF_StringView string_view_name = TF_OpKernelConstruction_GetName(ctx);
+  kernel->op_node_name =
+      std::string(string_view_name.data, string_view_name.len);
+  return kernel;
+}
+
+void HistogramSummaryOp_Delete(void* kernel) {
+  delete static_cast<HistogramSummaryOp*>(kernel);
+}
+
+template <typename T>
+void HistogramSummaryOp_Compute(void* kernel, TF_OpKernelContext* ctx) {
+  HistogramSummaryOp* k = static_cast<HistogramSummaryOp*>(kernel);
+  TF_Tensor* tags = nullptr;
+  TF_Tensor* values = nullptr;
+  Safe_TF_StatusPtr status(TF_NewStatus());
+  TF_GetInput(ctx, 0, &tags, status.get());
+  Safe_TF_TensorPtr safe_tags_ptr(tags);
+  if (TF_GetCode(status.get()) != TF_OK) {
+    TF_OpKernelContext_Failure(ctx, status.get());
+    return;
+  }
+  TF_GetInput(ctx, 1, &values, status.get());
+  Safe_TF_TensorPtr safe_values_ptr(values);
+  if (TF_GetCode(status.get()) != TF_OK) {
+    TF_OpKernelContext_Failure(ctx, status.get());
+    return;
+  }
+  if (TF_NumDims(safe_tags_ptr.get()) != 0) {
+    TF_SetStatus(status.get(), TF_INVALID_ARGUMENT, "tags must be scalar");
+    TF_OpKernelContext_Failure(ctx, status.get());
+    return;
+  }
+  // Cast values to array to access tensor elements by index
+  auto values_array = static_cast<T*>(TF_TensorData(safe_values_ptr.get()));
+  tensorflow::histogram::Histogram histo;
+  for (int64_t i = 0; i < TF_TensorElementCount(safe_values_ptr.get()); ++i) {
+    const double double_val = static_cast<double>(values_array[i]);
+    if (Eigen::numext::isnan(double_val)) {
+      std::ostringstream err;
+      err << "Nan in summary histogram for: " << k->op_node_name;
+      TF_SetStatus(status.get(), TF_INVALID_ARGUMENT, err.str().c_str());
+      TF_OpKernelContext_Failure(ctx, status.get());
+      return;
+    } else if (Eigen::numext::isinf(double_val)) {
+      std::ostringstream err;
+      err << "Infinity in Histogram for: " << k->op_node_name;
+      TF_SetStatus(status.get(), TF_INVALID_ARGUMENT, err.str().c_str());
+      TF_OpKernelContext_Failure(ctx, status.get());
+      return;
+    }
+    histo.Add(double_val);
+  }
+  tensorflow::Summary s;
+  tensorflow::Summary::Value* v = s.add_value();
+  const tensorflow::tstring& tag =
+      *(static_cast<tensorflow::tstring*>(TF_TensorData(safe_tags_ptr.get())));
+  v->set_tag(tag.data(), tag.size());
+  histo.EncodeToProto(v->mutable_histo(), false /* Drop zero buckets */);
+
+  Safe_TF_TensorPtr summary_tensor(TF_AllocateOutput(
+      /*context=*/ctx, /*index=*/0, /*dtype=*/TF_ExpectedOutputDataType(ctx, 0),
+      /*dims=*/nullptr, /*num_dims=*/0,
+      /*len=*/sizeof(tensorflow::tstring), status.get()));
+
+  if (TF_GetCode(status.get()) != TF_OK) {
+    TF_OpKernelContext_Failure(ctx, status.get());
+    return;
+  }
+  tensorflow::tstring* output_tstring = reinterpret_cast<tensorflow::tstring*>(
+      TF_TensorData(summary_tensor.get()));
+  CHECK(SerializeToTString(s, output_tstring));
+}
+
+template <typename T>
+void RegisterHistogramSummaryOpKernel() {
+  TF_Status* status = TF_NewStatus();
+  {
+    auto* builder = TF_NewKernelBuilder(
+        "HistogramSummary", tensorflow::DEVICE_CPU, &HistogramSummaryOp_Create,
+        &HistogramSummaryOp_Compute<T>, &HistogramSummaryOp_Delete);
+    TF_KernelBuilder_TypeConstraint(
+        builder, "T",
+        static_cast<TF_DataType>(tensorflow::DataTypeToEnum<T>::v()), status);
+    CHECK_EQ(TF_OK, TF_GetCode(status)) << "Error while adding type constraint";
+    TF_RegisterKernelBuilder("HistogramSummary", builder, status);
+    CHECK_EQ(TF_OK, TF_GetCode(status))
+        << "Error while registering Histogram Summmary kernel";
+  }
+  TF_DeleteStatus(status);
+}
+
+// A dummy static variable initialized by a lambda whose side-effect is to
+// register the Histogram Summary kernel.
+TF_ATTRIBUTE_UNUSED static bool IsHistogramSummaryOpKernelRegistered = []() {
+  if (SHOULD_REGISTER_OP_KERNEL("HistogramSummary")) {
+    RegisterHistogramSummaryOpKernel<int64_t>();
+    RegisterHistogramSummaryOpKernel<uint64_t>();
+    RegisterHistogramSummaryOpKernel<int32_t>();
+    RegisterHistogramSummaryOpKernel<uint32_t>();
+    RegisterHistogramSummaryOpKernel<uint16_t>();
+    RegisterHistogramSummaryOpKernel<int16_t>();
+    RegisterHistogramSummaryOpKernel<int8_t>();
+    RegisterHistogramSummaryOpKernel<uint8_t>();
+    RegisterHistogramSummaryOpKernel<Eigen::half>();
+    RegisterHistogramSummaryOpKernel<tensorflow::bfloat16>();
+    RegisterHistogramSummaryOpKernel<float>();
+    RegisterHistogramSummaryOpKernel<double>();
+  }
+  return true;
+}();
+}  // namespace

tensorflow/c/kernels/merge_summary_op.cc

@@ -0,0 +1,125 @@
+/* Copyright 2020 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 <memory>
+#include <sstream>
+#include <string>
+#include <unordered_set>
+
+#include "absl/log/check.h"
+#include "tensorflow/c/kernels.h"
+#include "tensorflow/c/tf_status.h"
+#include "tensorflow/c/tf_tensor.h"
+#include "tensorflow/core/framework/registration/registration.h"
+#include "tensorflow/core/framework/summary.pb.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/platform/logging.h"
+#include "tensorflow/core/platform/macros.h"
+#include "tensorflow/core/platform/protobuf.h"
+#include "tensorflow/core/platform/tstring.h"
+
+namespace {
+
+// Operators used to create a std::unique_ptr for TF_Tensor and TF_Status
+struct TFTensorDeleter {
+  void operator()(TF_Tensor* tf_tensor) const { TF_DeleteTensor(tf_tensor); }
+};
+
+struct TFStatusDeleter {
+  void operator()(TF_Status* tf_status) const { TF_DeleteStatus(tf_status); }
+};
+
+// Struct that wraps TF_Tensor and TF_Status to delete once out of scope
+using Safe_TF_TensorPtr = std::unique_ptr<TF_Tensor, TFTensorDeleter>;
+using Safe_TF_StatusPtr = std::unique_ptr<TF_Status, TFStatusDeleter>;
+
+// dummy functions used for kernel registration
+void* MergeSummaryOp_Create(TF_OpKernelConstruction* ctx) { return nullptr; }
+
+void MergeSummaryOp_Delete(void* kernel) {}
+
+void MergeSummaryOp_Compute(void* kernel, TF_OpKernelContext* ctx) {
+  tensorflow::Summary s;
+  std::unordered_set<std::string> tags;
+  Safe_TF_StatusPtr status(TF_NewStatus());
+  for (int input_num = 0; input_num < TF_NumInputs(ctx); ++input_num) {
+    TF_Tensor* input = nullptr;
+    TF_GetInput(ctx, input_num, &input, status.get());
+    Safe_TF_TensorPtr safe_input_ptr(input);
+    if (TF_GetCode(status.get()) != TF_OK) {
+      TF_OpKernelContext_Failure(ctx, status.get());
+      return;
+    }
+    auto tags_array =
+        static_cast<tensorflow::tstring*>(TF_TensorData(safe_input_ptr.get()));
+    for (int i = 0; i < TF_TensorElementCount(safe_input_ptr.get()); ++i) {
+      const tensorflow::tstring& s_in = tags_array[i];
+      tensorflow::Summary summary_in;
+      if (!tensorflow::ParseProtoUnlimited(&summary_in, s_in)) {
+        TF_SetStatus(status.get(), TF_INVALID_ARGUMENT,
+                     "Could not parse one of the summary inputs");
+        TF_OpKernelContext_Failure(ctx, status.get());
+        return;
+      }
+      for (int v = 0; v < summary_in.value_size(); ++v) {
+        // This tag is unused by the TensorSummary op, so no need to check for
+        // duplicates.
+        const std::string& tag = summary_in.value(v).tag();
+        if ((!tag.empty()) && !tags.insert(tag).second) {
+          std::ostringstream err;
+          err << "Duplicate tag " << tag << " found in summary inputs ";
+          TF_SetStatus(status.get(), TF_INVALID_ARGUMENT, err.str().c_str());
+          TF_OpKernelContext_Failure(ctx, status.get());
+          return;
+        }
+        *s.add_value() = summary_in.value(v);
+      }
+    }
+  }
+  Safe_TF_TensorPtr summary_tensor(TF_AllocateOutput(
+      /*context=*/ctx, /*index=*/0, /*dtype=*/TF_ExpectedOutputDataType(ctx, 0),
+      /*dims=*/nullptr, /*num_dims=*/0,
+      /*len=*/sizeof(tensorflow::tstring), status.get()));
+  if (TF_GetCode(status.get()) != TF_OK) {
+    TF_OpKernelContext_Failure(ctx, status.get());
+    return;
+  }
+  tensorflow::tstring* output_tstring = reinterpret_cast<tensorflow::tstring*>(
+      TF_TensorData(summary_tensor.get()));
+  CHECK(SerializeToTString(s, output_tstring));
+}
+
+void RegisterMergeSummaryOpKernel() {
+  TF_Status* status = TF_NewStatus();
+  {
+    auto* builder = TF_NewKernelBuilder(
+        "MergeSummary", tensorflow::DEVICE_CPU, &MergeSummaryOp_Create,
+        &MergeSummaryOp_Compute, &MergeSummaryOp_Delete);
+    TF_RegisterKernelBuilder("MergeSummary", builder, status);
+    CHECK_EQ(TF_OK, TF_GetCode(status))
+        << "Error while registering Merge Summmary kernel";
+  }
+  TF_DeleteStatus(status);
+}
+
+// A dummy static variable initialized by a lambda whose side-effect is to
+// register the Histogram Summary kernel.
+TF_ATTRIBUTE_UNUSED static bool IsMergeSummaryOpKernelRegistered = []() {
+  if (SHOULD_REGISTER_OP_KERNEL("MergeSummary")) {
+    RegisterMergeSummaryOpKernel();
+  }
+  return true;
+}();
+
+}  // namespace

tensorflow/c/kernels/ops/bitcast.cc

@@ -0,0 +1,134 @@
+/* Copyright 2019 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 <sstream>
+#include <string>
+
+#include "tensorflow/c/ops.h"
+#include "tensorflow/c/tf_datatype.h"
+#include "tensorflow/c/tf_status.h"
+#include "tensorflow/core/framework/registration/registration.h"
+#include "tensorflow/core/platform/logging.h"
+#include "tensorflow/core/platform/macros.h"
+
+static void ComputeNewShape(TF_ShapeInferenceContext* ctx,
+                            TF_ShapeHandle* shape, TF_DataType input_type,
+                            TF_DataType output_type, TF_Status* status) {
+  size_t input_type_size = TF_DataTypeSize(input_type);
+  size_t output_type_size = TF_DataTypeSize(output_type);
+
+  if (input_type_size == 0 || output_type_size == 0) {
+    std::ostringstream err;
+    err << "Cannot bitcast type " << input_type << " to " << output_type
+        << " because one of the type sizes is zero";
+    TF_SetStatus(status, TF_INVALID_ARGUMENT, err.str().c_str());
+    return;
+  }
+
+  TF_SetStatus(status, TF_OK, "");
+  if (input_type_size < output_type_size) {
+    TF_ShapeInferenceContextWithRankAtLeast(ctx, shape, 1, shape, status);
+
+    if (TF_GetCode(status) == TF_OK) {
+      TF_DimensionHandle* last_dim = TF_NewDimensionHandle();
+      size_t divisor_val = output_type_size / input_type_size;
+      TF_ShapeInferenceContextDim(ctx, shape, -1, last_dim);
+      if (!TF_DimensionHandleValueKnown(last_dim) ||
+          TF_DimensionHandleValue(last_dim) == divisor_val) {
+        TF_ShapeInferenceContextSubshape(ctx, shape, 0, -1, shape, status);
+      } else {
+        std::ostringstream err;
+        err << "Cannot bitcast from " << input_type << " to " << output_type
+            << " due to shape. " << TF_DimensionHandleValue(last_dim)
+            << " does not match " << divisor_val;
+        TF_SetStatus(status, TF_INVALID_ARGUMENT, err.str().c_str());
+      }
+      TF_DeleteDimensionHandle(last_dim);
+    }
+  } else if (input_type_size > output_type_size) {
+    // Input type size is larger than output type size.
+    size_t divisor_val = input_type_size / output_type_size;
+    TF_ShapeHandle* extension =
+        TF_ShapeInferenceContextVectorFromSize(ctx, divisor_val);
+    TF_ShapeInferenceContextConcatenateShapes(ctx, shape, extension, shape,
+                                              status);
+    TF_DeleteShapeHandle(extension);
+  }
+}
+
+static void bitcast_shape_inference_fn(TF_ShapeInferenceContext* ctx,
+                                       TF_Status* status) {
+  TF_ShapeHandle* result = TF_NewShapeHandle();
+  TF_ShapeInferenceContextGetInput(ctx, 0, result, status);
+  if (TF_GetCode(status) == TF_OK &&
+      !TF_ShapeInferenceContextRankKnown(ctx, result)) {
+    TF_ShapeInferenceContextSetUnknownShape(ctx, status);
+    TF_DeleteShapeHandle(result);
+    return;
+  }
+
+  // Find the size of the input and output data types.
+  TF_DataType input_type;
+  TF_DataType output_type;
+
+  if (TF_GetCode(status) == TF_OK) {
+    TF_ShapeInferenceContext_GetAttrType(ctx, "T", &input_type, status);
+  }
+
+  if (TF_GetCode(status) == TF_OK) {
+    TF_ShapeInferenceContext_GetAttrType(ctx, "type", &output_type, status);
+  }
+
+  if (TF_GetCode(status) == TF_OK) {
+    ComputeNewShape(ctx, result, input_type, output_type, status);
+  }
+
+  if (TF_GetCode(status) == TF_OK) {
+    TF_ShapeInferenceContextSetOutput(ctx, 0, result, status);
+  }
+  TF_DeleteShapeHandle(result);
+}
+
+void RegisterBitcastOp() {
+  TF_Status* status = TF_NewStatus();
+
+  TF_OpDefinitionBuilder* op_builder = TF_NewOpDefinitionBuilder("Bitcast");
+  TF_OpDefinitionBuilderAddInput(op_builder, "input: T");
+  TF_OpDefinitionBuilderAddOutput(op_builder, "output: type");
+  TF_OpDefinitionBuilderAddAttr(
+      op_builder,
+      "T: {bfloat16, half, float, double, int64, int32, uint8, uint16, "
+      "uint32, uint64, int8, int16, complex64, complex128, qint8, quint8, "
+      "qint16, quint16, qint32}");
+  TF_OpDefinitionBuilderAddAttr(
+      op_builder,
+      "type: {bfloat16, half, float, double, int64, int32, uint8, uint16, "
+      "uint32, uint64, int8, int16, complex64, complex128, qint8, quint8, "
+      "qint16, quint16, qint32}");
+  TF_OpDefinitionBuilderSetShapeInferenceFunction(op_builder,
+                                                  &bitcast_shape_inference_fn);
+
+  TF_RegisterOpDefinition(op_builder, status);
+  CHECK_EQ(TF_GetCode(status), TF_OK)
+      << "Bitcast op registration failed: " << TF_Message(status);
+  TF_DeleteStatus(status);
+}
+
+TF_ATTRIBUTE_UNUSED static bool IsBitcastOpRegistered = []() {
+  if ((&TF_NewStatus != nullptr) && SHOULD_REGISTER_OP("Bitcast")) {
+    RegisterBitcastOp();
+  }
+  return true;
+}();

tensorflow/c/kernels/ops/histogram_summary.cc

@@ -0,0 +1,50 @@
+/* Copyright 2020 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/c/ops.h"
+#include "tensorflow/c/tf_status.h"
+#include "tensorflow/core/framework/registration/registration.h"
+#include "tensorflow/core/platform/logging.h"
+#include "tensorflow/core/platform/macros.h"
+
+static void histogram_summary_shape_inference_fn(TF_ShapeInferenceContext* ctx,
+                                                 TF_Status* status) {
+  TF_SetStatus(status, TF_OK, "");
+  TF_ShapeHandle* result = TF_ShapeInferenceContextScalar(ctx);
+  TF_ShapeInferenceContextSetOutput(ctx, 0, result, status);
+  TF_DeleteShapeHandle(result);
+}
+
+void Register_HistogramSummaryOp() {
+  TF_Status* status = TF_NewStatus();
+
+  TF_OpDefinitionBuilder* op_builder =
+      TF_NewOpDefinitionBuilder("HistogramSummary");
+  TF_OpDefinitionBuilderAddInput(op_builder, "tag: string");
+  TF_OpDefinitionBuilderAddInput(op_builder, "values: T");
+  TF_OpDefinitionBuilderAddOutput(op_builder, "summary: string");
+  TF_OpDefinitionBuilderAddAttr(op_builder, "T: realnumbertype = DT_FLOAT");
+  TF_OpDefinitionBuilderSetShapeInferenceFunction(
+      op_builder, &histogram_summary_shape_inference_fn);
+
+  TF_RegisterOpDefinition(op_builder, status);
+  CHECK_EQ(TF_GetCode(status), TF_OK)
+      << "HistogramSummary op registration failed: " << TF_Message(status);
+  TF_DeleteStatus(status);
+}
+
+TF_ATTRIBUTE_UNUSED static bool HistogramSummaryOpRegistered = []() {
+  if (&TF_NewStatus != nullptr && SHOULD_REGISTER_OP("HistogramSummary")) {
+    Register_HistogramSummaryOp();
+  }
+  return true;
+}();

tensorflow/c/kernels/ops/merge_summary.cc

@@ -0,0 +1,52 @@
+/* Copyright 2020 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/c/ops.h"
+#include "tensorflow/c/tf_status.h"
+#include "tensorflow/core/framework/registration/registration.h"
+#include "tensorflow/core/platform/logging.h"
+#include "tensorflow/core/platform/macros.h"
+
+static void merge_summary_shape_inference_fn(TF_ShapeInferenceContext* ctx,
+                                             TF_Status* status) {
+  TF_SetStatus(status, TF_OK, "");
+  TF_ShapeHandle* result = TF_ShapeInferenceContextScalar(ctx);
+  TF_ShapeInferenceContextSetOutput(ctx, 0, result, status);
+  TF_DeleteShapeHandle(result);
+}
+
+void Register_MergeSummaryOp() {
+  TF_Status* status = TF_NewStatus();
+
+  TF_OpDefinitionBuilder* op_builder =
+      TF_NewOpDefinitionBuilder("MergeSummary");
+  TF_OpDefinitionBuilderAddInput(op_builder, "inputs: N * string");
+  TF_OpDefinitionBuilderAddOutput(op_builder, "summary: string");
+  TF_OpDefinitionBuilderAddAttr(op_builder, "N: int >= 1");
+  TF_OpDefinitionBuilderSetShapeInferenceFunction(
+      op_builder, &merge_summary_shape_inference_fn);
+
+  TF_RegisterOpDefinition(op_builder, status);
+  CHECK_EQ(TF_GetCode(status), TF_OK)
+      << "MergeSummary op registration failed: " << TF_Message(status);
+  TF_DeleteStatus(status);
+}
+
+TF_ATTRIBUTE_UNUSED static bool MergeSummaryOpRegistered = []() {
+  if ((&TF_NewStatus != nullptr) && SHOULD_REGISTER_OP("MergeSummary")) {
+    Register_MergeSummaryOp();
+  }
+  return true;
+}();

tensorflow/c/kernels/ops/summary.cc

@@ -0,0 +1,53 @@
+/* Copyright 2020 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/c/ops.h"
+#include "tensorflow/c/tf_status.h"
+#include "tensorflow/core/framework/registration/registration.h"
+#include "tensorflow/core/platform/logging.h"
+#include "tensorflow/core/platform/macros.h"
+
+static void scalar_summary_shape_inference_fn(TF_ShapeInferenceContext* ctx,
+                                              TF_Status* status) {
+  TF_SetStatus(status, TF_OK, "");
+  TF_ShapeHandle* result = TF_ShapeInferenceContextScalar(ctx);
+  TF_ShapeInferenceContextSetOutput(ctx, 0, result, status);
+  TF_DeleteShapeHandle(result);
+}
+
+void Register_ScalarSummaryOp() {
+  TF_Status* status = TF_NewStatus();
+
+  TF_OpDefinitionBuilder* op_builder =
+      TF_NewOpDefinitionBuilder("ScalarSummary");
+  TF_OpDefinitionBuilderAddInput(op_builder, "tags: string");
+  TF_OpDefinitionBuilderAddInput(op_builder, "values: T");
+  TF_OpDefinitionBuilderAddOutput(op_builder, "summary: string");
+  TF_OpDefinitionBuilderAddAttr(op_builder, "T: realnumbertype");
+  TF_OpDefinitionBuilderSetShapeInferenceFunction(
+      op_builder, &scalar_summary_shape_inference_fn);
+
+  TF_RegisterOpDefinition(op_builder, status);
+  CHECK_EQ(TF_GetCode(status), TF_OK)
+      << "ScalarSummary op registration failed: " << TF_Message(status);
+  TF_DeleteStatus(status);
+}
+
+TF_ATTRIBUTE_UNUSED static bool SummaryScalarOpRegistered = []() {
+  if (&TF_NewStatus != nullptr && SHOULD_REGISTER_OP("ScalarSummary")) {
+    Register_ScalarSummaryOp();
+  }
+  return true;
+}();

tensorflow/c/kernels/summary_op.cc

@@ -0,0 +1,174 @@
+
+/* Copyright 2020 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 <cstdint>
+#include <sstream>
+#include <string>
+
+#include "absl/log/check.h"
+#include "absl/strings/str_cat.h"
+#include "Eigen/Core"  // from @eigen_archive
+#include "unsupported/Eigen/CXX11/Tensor"  // from @eigen_archive
+#include "tensorflow/c/kernels.h"
+#include "tensorflow/c/kernels/tensor_shape_utils.h"
+#include "tensorflow/c/tf_datatype.h"
+#include "tensorflow/c/tf_status.h"
+#include "tensorflow/c/tf_tensor.h"
+#include "xla/tsl/platform/macros.h"
+#include "tensorflow/core/framework/registration/registration.h"
+#include "tensorflow/core/framework/summary.pb.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/platform/bfloat16.h"
+#include "tensorflow/core/platform/tstring.h"
+#include "tsl/platform/protobuf.h"
+
+namespace {
+
+// Struct that stores the status and TF_Tensor inputs to the opkernel.
+// Used to delete tensor and status in its destructor upon kernel return.
+struct Params {
+  TF_Tensor* tags;
+  TF_Tensor* values;
+  TF_Status* status;
+  explicit Params(TF_OpKernelContext* ctx)
+      : tags(nullptr), values(nullptr), status(nullptr) {
+    status = TF_NewStatus();
+    TF_GetInput(ctx, 0, &tags, status);
+    if (TF_GetCode(status) == TF_OK) {
+      TF_GetInput(ctx, 1, &values, status);
+    }
+  }
+  ~Params() {
+    TF_DeleteStatus(status);
+    TF_DeleteTensor(tags);
+    TF_DeleteTensor(values);
+  }
+};
+
+// dummy functions used for kernel registration
+void* ScalarSummaryOp_Create(TF_OpKernelConstruction* ctx) { return nullptr; }
+
+void ScalarSummaryOp_Delete(void* kernel) {}
+
+// Helper functions for compute method
+bool IsSameSize(TF_Tensor* tensor1, TF_Tensor* tensor2);
+// Returns a string representation of a single tag or empty string if there
+// are multiple tags
+std::string SingleTag(TF_Tensor* tags);
+
+template <typename T>
+void ScalarSummaryOp_Compute(void* kernel, TF_OpKernelContext* ctx) {
+  Params params(ctx);
+  if (TF_GetCode(params.status) != TF_OK) {
+    TF_OpKernelContext_Failure(ctx, params.status);
+    return;
+  }
+  if (!IsSameSize(params.tags, params.values)) {
+    std::ostringstream err;
+    err << "tags and values are not the same shape: "
+        << tensorflow::ShapeDebugString(params.tags)
+        << " != " << tensorflow::ShapeDebugString(params.values)
+        << SingleTag(params.tags);
+    TF_SetStatus(params.status, TF_INVALID_ARGUMENT, err.str().c_str());
+    TF_OpKernelContext_Failure(ctx, params.status);
+    return;
+  }
+  // Convert tags and values tensor to array to access elements by index
+  tensorflow::Summary s;
+  auto tags_array =
+      static_cast<tensorflow::tstring*>(TF_TensorData(params.tags));
+  auto values_array = static_cast<T*>(TF_TensorData(params.values));
+  // Copy tags and values into summary protobuf
+  for (int64_t i = 0; i < TF_TensorElementCount(params.tags); ++i) {
+    tensorflow::Summary::Value* v = s.add_value();
+    const tensorflow::tstring& Ttags_i = tags_array[i];
+    v->set_tag(Ttags_i.data(), Ttags_i.size());
+    v->set_simple_value(static_cast<float>(values_array[i]));
+  }
+  TF_Tensor* summary_tensor =
+      TF_AllocateOutput(ctx, 0, TF_ExpectedOutputDataType(ctx, 0), nullptr, 0,
+                        sizeof(tensorflow::tstring), params.status);
+  if (TF_GetCode(params.status) != TF_OK) {
+    TF_DeleteTensor(summary_tensor);
+    TF_OpKernelContext_Failure(ctx, params.status);
+    return;
+  }
+  tensorflow::tstring* output_tstring =
+      reinterpret_cast<tensorflow::tstring*>(TF_TensorData(summary_tensor));
+  CHECK(SerializeToTString(s, output_tstring));
+  TF_DeleteTensor(summary_tensor);
+}
+
+bool IsSameSize(TF_Tensor* tensor1, TF_Tensor* tensor2) {
+  if (TF_NumDims(tensor1) != TF_NumDims(tensor2)) {
+    return false;
+  }
+  for (int d = 0; d < TF_NumDims(tensor1); d++) {
+    if (TF_Dim(tensor1, d) != TF_Dim(tensor2, d)) {
+      return false;
+    }
+  }
+  return true;
+}
+
+std::string SingleTag(TF_Tensor* tags) {
+  if (TF_TensorElementCount(tags) == 1) {
+    const char* single_tag =
+        static_cast<tensorflow::tstring*>(TF_TensorData(tags))->c_str();
+    return absl::StrCat(" (tag '", single_tag, "')");
+  } else {
+    return "";
+  }
+}
+
+template <typename T>
+void RegisterScalarSummaryOpKernel() {
+  TF_Status* status = TF_NewStatus();
+  {
+    auto* builder = TF_NewKernelBuilder(
+        "ScalarSummary", tensorflow::DEVICE_CPU, &ScalarSummaryOp_Create,
+        &ScalarSummaryOp_Compute<T>, &ScalarSummaryOp_Delete);
+    TF_KernelBuilder_TypeConstraint(
+        builder, "T",
+        static_cast<TF_DataType>(tensorflow::DataTypeToEnum<T>::v()), status);
+    CHECK_EQ(TF_OK, TF_GetCode(status)) << "Error while adding type constraint";
+    TF_RegisterKernelBuilder("ScalarSummary", builder, status);
+    CHECK_EQ(TF_OK, TF_GetCode(status))
+        << "Error while registering Scalar Summmary kernel";
+  }
+  TF_DeleteStatus(status);
+}
+
+// A dummy static variable initialized by a lambda whose side-effect is to
+// register the ScalarSummary kernel.
+TF_ATTRIBUTE_UNUSED bool IsScalarSummaryOpKernelRegistered = []() {
+  if (SHOULD_REGISTER_OP_KERNEL("ScalarSummary")) {
+    RegisterScalarSummaryOpKernel<int64_t>();
+    RegisterScalarSummaryOpKernel<uint64_t>();
+    RegisterScalarSummaryOpKernel<int32_t>();
+    RegisterScalarSummaryOpKernel<uint32_t>();
+    RegisterScalarSummaryOpKernel<uint16_t>();
+    RegisterScalarSummaryOpKernel<int16_t>();
+    RegisterScalarSummaryOpKernel<int8_t>();
+    RegisterScalarSummaryOpKernel<uint8_t>();
+    RegisterScalarSummaryOpKernel<Eigen::half>();
+    RegisterScalarSummaryOpKernel<tensorflow::bfloat16>();
+    RegisterScalarSummaryOpKernel<float>();
+    RegisterScalarSummaryOpKernel<double>();
+  }
+  return true;
+}();
+}  // namespace

tensorflow/c/kernels/summary_op_benchmark_test.cc

@@ -0,0 +1,70 @@
+/* Copyright 2020 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>
+
+#include "tensorflow/core/common_runtime/kernel_benchmark_testlib.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/graph/node_builder.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/platform/test_benchmark.h"
+
+namespace tensorflow {
+namespace {
+
+Graph* BM_ScalarSummaryOp(TensorShape shape, std::string tag, float value) {
+  Graph* g = new Graph(OpRegistry::Global());
+  Tensor tags(DT_STRING, shape);
+  Tensor values(DT_FLOAT, shape);
+  for (int i = 0; i < tags.NumElements(); ++i) {
+    tags.flat<tstring>()(i) = tag;
+    values.flat<float>()(i) = value;
+  }
+  Node* ret;
+  TF_CHECK_OK(NodeBuilder(g->NewName("dummy"), "ScalarSummary")
+                  .Input(test::graph::Constant(g, tags))
+                  .Input(test::graph::Constant(g, values))
+                  .Attr("T", DT_FLOAT)
+                  .Finalize(g, &ret));
+  return g;
+}
+
+// Macro used to parse initializer list for tensorshape
+#define DIMARGS(...) \
+  { __VA_ARGS__ }
+// // Random parameters for testing
+constexpr char longTagParam[] = "LONGTAG____________________________";
+constexpr float largeValueParam = 2352352.2623433;
+
+#define BM_ScalarSummaryDev(device, dims, name, tag, value)                 \
+  void BM_ScalarSummary##name##device(::testing::benchmark::State& state) { \
+    TensorShape tensorshape(DIMARGS dims);                                  \
+    auto g = BM_ScalarSummaryOp(tensorshape, #tag, value);                  \
+    test::Benchmark("cpu", g, /*old_benchmark_api=*/false).Run(state);      \
+  }                                                                         \
+  BENCHMARK(BM_ScalarSummary##name##device);
+
+BM_ScalarSummaryDev(Cpu, (5, 10, 100), Base, Tag, 5.2);
+// Benchmark for large shapes
+BM_ScalarSummaryDev(Cpu, (500, 100, 100), LargeShape, Tag, 5.2);
+// Benchmark for large tag tstring
+BM_ScalarSummaryDev(Cpu, (5, 10, 100), LongTag, longTagParam, 5.2);
+// Benchmark for large values
+BM_ScalarSummaryDev(Cpu, (500, 100, 100), LargeValue, Tag, largeValueParam);
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/c/kernels/summary_op_test.cc

@@ -0,0 +1,188 @@
+/* Copyright 2020 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/c/kernels.h"
+#include "tensorflow/core/framework/allocator.h"
+#include "tensorflow/core/framework/attr_value.pb.h"
+#include "tensorflow/core/framework/attr_value_util.h"
+#include "tensorflow/core/framework/device_base.h"
+#include "tensorflow/core/framework/node_def.pb.h"
+#include "tensorflow/core/framework/op.h"
+#include "tensorflow/core/framework/op_def_builder.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/summary.pb.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/lib/gtl/inlined_vector.h"
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/protobuf.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/platform/strcat.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/protobuf/error_codes.pb.h"
+
+namespace tensorflow {
+namespace {
+
+class DummyDevice : public DeviceBase {
+ public:
+  explicit DummyDevice(Env* env) : DeviceBase(env) {}
+  Allocator* GetAllocator(AllocatorAttributes /*attr*/) override {
+    return cpu_allocator();
+  }
+};
+
+// Helper for comparing output and expected output
+void ExpectSummaryMatches(const Summary& actual,
+                          const std::string& expected_str) {
+  Summary expected;
+  ASSERT_TRUE(protobuf::TextFormat::ParseFromString(expected_str, &expected));
+  EXPECT_EQ(expected.DebugString(), actual.DebugString());
+}
+
+void TestScalarSummaryOp(Tensor* tags, Tensor* values,
+                         std::string expected_output,
+                         error::Code expected_code) {
+  // Initialize node used to fetch OpKernel
+  absl::Status status;
+  NodeDef def;
+  def.set_op("ScalarSummary");
+
+  def.set_device(DEVICE_CPU);
+
+  AttrValue valuesTypeAttr;
+  SetAttrValue(values->dtype(), &valuesTypeAttr);
+  (*def.mutable_attr())["T"] = valuesTypeAttr;
+
+  def.add_input(absl::StrCat("input1: ", DataTypeString(tags->dtype())));
+  def.add_input(absl::StrCat("input2: ", DataTypeString(values->dtype())));
+
+  std::unique_ptr<OpKernel> kernel =
+      CreateOpKernel(DeviceType(DEVICE_CPU), nullptr, nullptr, def, 1, &status);
+  ASSERT_TRUE(status.ok()) << status.ToString();
+  OpKernelContext::Params params;
+  DummyDevice dummy_device(nullptr);
+  params.device = &dummy_device;
+  params.op_kernel = kernel.get();
+  AllocatorAttributes alloc_attrs;
+  params.output_attr_array = &alloc_attrs;
+  absl::InlinedVector<TensorValue, 4UL> inputs;
+  inputs.emplace_back(tags);
+  inputs.emplace_back(values);
+  params.inputs = inputs;
+  OpKernelContext ctx(&params, 1);
+  kernel->Compute(&ctx);
+  ASSERT_EQ(expected_code, ctx.status().code());
+  if (expected_code == error::OK) {
+    Summary summary;
+    ASSERT_TRUE(ParseProtoUnlimited(
+        &summary, ctx.mutable_output(0)->scalar<tstring>()()));
+    ExpectSummaryMatches(summary, expected_output);
+  } else {
+    EXPECT_TRUE(absl::StrContains(ctx.status().ToString(), expected_output))
+        << ctx.status();
+  }
+}
+
+TEST(ScalarSummaryOpTest, SimpleFloat) {
+  int vectorSize = 3;
+  Tensor tags(DT_STRING, {vectorSize});
+  Tensor values(DT_FLOAT, {vectorSize});
+  tags.vec<tstring>()(0) = "tag1";
+  tags.vec<tstring>()(1) = "tag2";
+  tags.vec<tstring>()(2) = "tag3";
+  values.vec<float>()(0) = 1.0f;
+  values.vec<float>()(1) = -0.73f;
+  values.vec<float>()(2) = 10000.0f;
+  TestScalarSummaryOp(&tags, &values, R"(
+                      value { tag: 'tag1' simple_value: 1.0 }
+                      value { tag: 'tag2' simple_value: -0.73}
+                      value { tag: 'tag3' simple_value: 10000.0})",
+                      error::OK);
+}
+
+TEST(ScalarSummaryOpTest, SimpleDouble) {
+  int vectorSize = 3;
+  Tensor tags(DT_STRING, {vectorSize});
+  Tensor values(DT_DOUBLE, {vectorSize});
+  tags.vec<tstring>()(0) = "tag1";
+  tags.vec<tstring>()(1) = "tag2";
+  tags.vec<tstring>()(2) = "tag3";
+  values.vec<double>()(0) = 1.0;
+  values.vec<double>()(1) = -0.73;
+  values.vec<double>()(2) = 10000.0;
+  TestScalarSummaryOp(&tags, &values, R"(
+                      value { tag: 'tag1' simple_value: 1.0 }
+                      value { tag: 'tag2' simple_value: -0.73}
+                      value { tag: 'tag3' simple_value: 10000.0})",
+                      error::OK);
+}
+
+TEST(ScalarSummaryOpTest, SimpleHalf) {
+  int vectorSize = 3;
+  Tensor tags(DT_STRING, {vectorSize});
+  Tensor values(DT_HALF, {vectorSize});
+  tags.vec<tstring>()(0) = "tag1";
+  tags.vec<tstring>()(1) = "tag2";
+  tags.vec<tstring>()(2) = "tag3";
+  values.vec<Eigen::half>()(0) = Eigen::half(1.0);
+  values.vec<Eigen::half>()(1) = Eigen::half(-2.0);
+  values.vec<Eigen::half>()(2) = Eigen::half(10000.0);
+  TestScalarSummaryOp(&tags, &values, R"(
+                      value { tag: 'tag1' simple_value: 1.0 }
+                      value { tag: 'tag2' simple_value: -2.0}
+                      value { tag: 'tag3' simple_value: 10000.0})",
+                      error::OK);
+}
+
+TEST(ScalarSummaryOpTest, Error_WrongDimsTags) {
+  Tensor tags(DT_STRING, {2, 1});
+  Tensor values(DT_FLOAT, {2});
+  tags.matrix<tstring>()(0, 0) = "tag1";
+  tags.matrix<tstring>()(1, 0) = "tag2";
+  values.vec<float>()(0) = 1.0f;
+  values.vec<float>()(1) = -2.0f;
+  TestScalarSummaryOp(&tags, &values, "tags and values are not the same shape",
+                      error::INVALID_ARGUMENT);
+}
+
+TEST(ScalarSummaryOpTest, Error_WrongValuesTags) {
+  Tensor tags(DT_STRING, {2});
+  Tensor values(DT_FLOAT, {2, 1});
+  tags.vec<tstring>()(0) = "tag1";
+  tags.vec<tstring>()(1) = "tag2";
+  values.matrix<float>()(0, 0) = 1.0f;
+  values.matrix<float>()(1, 0) = -2.0f;
+  TestScalarSummaryOp(&tags, &values, "tags and values are not the same shape",
+                      error::INVALID_ARGUMENT);
+}
+
+TEST(ScalarSummaryOpTest, Error_WrongWithSingleTag) {
+  Tensor tags(DT_STRING, {1});
+  Tensor values(DT_FLOAT, {2, 1});
+  tags.vec<tstring>()(0) = "tag1";
+  values.matrix<float>()(0, 0) = 1.0f;
+  values.matrix<float>()(1, 0) = -2.0f;
+  TestScalarSummaryOp(&tags, &values, "tags and values are not the same shape",
+                      error::INVALID_ARGUMENT);
+}
+
+TEST(ScalarSummaryOpTest, IsRegistered) {
+  const OpRegistrationData* reg;
+  TF_CHECK_OK(OpRegistry::Global()->LookUp("ScalarSummary", &reg));
+}
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/c/kernels/tensor_shape_utils.cc

@@ -0,0 +1,43 @@
+/* Copyright 2020 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/c/kernels/tensor_shape_utils.h"
+
+#include <cstdint>
+#include <string>
+
+#include "absl/log/check.h"
+#include "absl/strings/str_cat.h"
+#include "tensorflow/c/tf_tensor.h"
+#include "tensorflow/core/platform/logging.h"
+#include "tensorflow/core/platform/strcat.h"
+
+namespace tensorflow {
+
+std::string ShapeDebugString(TF_Tensor* tensor) {
+  // A TF_Tensor cannot have an unknown rank.
+  CHECK_GE(TF_NumDims(tensor), 0);
+  std::string s = "[";
+  for (int i = 0; i < TF_NumDims(tensor); ++i) {
+    if (i > 0) absl::StrAppend(&s, ",");
+    int64_t dim = TF_Dim(tensor, i);
+    // A TF_Tensor cannot have an unknown dimension.
+    CHECK_GE(dim, 0);
+    absl::StrAppend(&s, dim);
+  }
+  absl::StrAppend(&s, "]");
+  return s;
+}
+}  // namespace tensorflow

tensorflow/c/kernels/tensor_shape_utils.h

@@ -0,0 +1,37 @@
+/* Copyright 2020 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.
+==============================================================================*/
+
+// This file contains shape utilities to be used by kernels and is not part of
+// the C API. As such, it is subject to change at any time.
+
+#ifndef TENSORFLOW_C_KERNELS_TENSOR_SHAPE_UTILS_H_
+#define TENSORFLOW_C_KERNELS_TENSOR_SHAPE_UTILS_H_
+
+#include <string>
+
+#include "tensorflow/c/tf_tensor.h"
+
+namespace tensorflow {
+
+// The following are utils for the shape of a TF_Tensor type.
+// These functions may later be subsumed by the methods for a
+// TF_TensorShape type.
+
+// Returns a string representation of the TF_Tensor shape.
+std::string ShapeDebugString(TF_Tensor* tensor);
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_C_KERNELS_TENSOR_SHAPE_UTILS_H_

tensorflow/c/kernels/tensor_shape_utils_test.cc

@@ -0,0 +1,53 @@
+/* Copyright 2020 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/c/kernels/tensor_shape_utils.h"
+
+#include "absl/status/status.h"
+#include "tensorflow/c/tf_tensor_internal.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/platform/types.h"
+
+namespace tensorflow {
+
+namespace {
+
+// A wrapper that will automatically delete the allocated TF_Tensor
+// once out of scope.
+struct TF_TensorWrapper {
+  TF_Tensor* tf_tensor;
+  explicit TF_TensorWrapper(TF_Tensor* tensor) { tf_tensor = tensor; }
+  ~TF_TensorWrapper() { TF_DeleteTensor(tf_tensor); }
+};
+
+void TestShapeMatch(TensorShape shape) {
+  Tensor tensor(DT_FLOAT, shape);
+  absl::Status status;
+  TF_Tensor* tf_tensor = TF_TensorFromTensor(tensor, &status);
+  TF_TensorWrapper tensor_wrapper = TF_TensorWrapper(tf_tensor);
+  ASSERT_TRUE(status.ok()) << status.ToString();
+  ASSERT_EQ(tensor.shape().DebugString(), ShapeDebugString(tf_tensor));
+}
+
+TEST(ShapeDebugString, RegularShape) { TestShapeMatch(TensorShape({5, 4, 7})); }
+
+TEST(ShapeDebugString, ScalarShape) { TestShapeMatch(TensorShape({})); }
+
+}  // namespace
+}  // namespace tensorflow