chore: import upstream snapshot with attribution

tensorflow/examples/README.md

@@ -0,0 +1,20 @@
+# TensorFlow C++ Examples
+
+This directory contains examples of the TensorFlow C++ API (and some redirects).
+If that's not what you're looking for here are some links:
+
+*   For TensorFlow python examples see
+    [the tutorials on tensorflow.org](https://tensorflow.org/tutorials)
+*   For community maintained keras examples goto
+    [keras.io/examples](https://keras.io/examples/)
+*   For TensorFlow Lite examples see
+    [the tensorflow/examples repository](https://github.com/tensorflow/examples/tree/master/lite)
+*   For the Udacity course notebooks, refer to
+    [this directory](https://github.com/tensorflow/examples/tree/master/courses)
+
+## About these examples
+
+*   The C++ API is only easily buildable from within the TensorFlow `bazel`
+    build. If you need a stand alone build
+    [see the C API](https://www.tensorflow.org/install/lang_c).
+*   This directory is not actively maintained.

tensorflow/examples/adding_an_op/BUILD

@@ -0,0 +1,162 @@
+# Description:
+# Code examples referenced by adding_an_op
+
+load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
+load("//tensorflow:tensorflow.bzl", "py_test", "tf_cc_binary", "tf_custom_op_library")
+load(
+    "//tensorflow/core/platform:build_config_root.bzl",
+    "tf_cuda_tests_tags",
+    "tf_exec_properties",
+)
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    licenses = ["notice"],
+)
+
+tf_custom_op_library(
+    name = "zero_out_op_kernel_1.so",
+    srcs = ["zero_out_op_kernel_1.cc"],
+)
+
+py_library(
+    name = "zero_out_op_1",
+    srcs = ["zero_out_op_1.py"],
+    data = [":zero_out_op_kernel_1.so"],
+    strict_deps = True,
+    deps = ["//tensorflow:tensorflow_py"],
+)
+
+tf_custom_op_library(
+    name = "zero_out_op_kernel_2.so",
+    srcs = ["zero_out_op_kernel_2.cc"],
+)
+
+py_library(
+    name = "zero_out_op_2",
+    srcs = ["zero_out_op_2.py"],
+    data = [":zero_out_op_kernel_2.so"],
+    strict_deps = True,
+    deps = ["//tensorflow:tensorflow_py"],
+)
+
+tf_custom_op_library(
+    name = "zero_out_op_kernel_3.so",
+    srcs = ["zero_out_op_kernel_3.cc"],
+)
+
+py_library(
+    name = "zero_out_op_3",
+    srcs = ["zero_out_op_3.py"],
+    data = [":zero_out_op_kernel_3.so"],
+    strict_deps = True,
+    deps = ["//tensorflow:tensorflow_py"],
+)
+
+py_library(
+    name = "zero_out_grad_2",
+    srcs = ["zero_out_grad_2.py"],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:sparse_ops",
+    ],
+)
+
+py_test(
+    name = "zero_out_1_test",
+    size = "medium",
+    srcs = ["zero_out_1_test.py"],
+    strict_deps = True,
+    tags = [
+        "notap",
+    ],
+    deps = [
+        ":zero_out_op_1",
+        # copybara:uncomment "//third_party/py/google/protobuf:use_fast_cpp_protos",
+        "//tensorflow:tensorflow_py",
+    ],
+)
+
+py_test(
+    name = "zero_out_2_test",
+    size = "medium",
+    srcs = ["zero_out_2_test.py"],
+    strict_deps = True,
+    tags = [
+        "notap",
+    ],
+    deps = [
+        ":zero_out_grad_2",
+        ":zero_out_op_2",
+        # copybara:uncomment "//third_party/py/google/protobuf:use_fast_cpp_protos",
+        "//tensorflow:tensorflow_py",
+    ],
+)
+
+py_test(
+    name = "zero_out_3_test",
+    size = "medium",
+    srcs = ["zero_out_3_test.py"],
+    strict_deps = True,
+    tags = [
+        "notap",
+    ],
+    deps = [
+        ":zero_out_op_3",
+        # copybara:uncomment "//third_party/py/google/protobuf:use_fast_cpp_protos",
+        "//tensorflow:tensorflow_py",
+    ],
+)
+
+tf_custom_op_library(
+    name = "cuda_op_kernel.so",
+    srcs = ["cuda_op_kernel.cc"],
+    gpu_srcs = ["cuda_op_kernel.cu.cc"],
+)
+
+py_library(
+    name = "cuda_op",
+    srcs = ["cuda_op.py"],
+    data = [":cuda_op_kernel.so"],
+    strict_deps = True,
+    deps = ["//tensorflow:tensorflow_py"],
+)
+
+py_test(
+    name = "cuda_op_test",
+    size = "small",
+    srcs = ["cuda_op_test.py"],
+    exec_properties = tf_exec_properties({"tags": tf_cuda_tests_tags()}),
+    strict_deps = True,
+    tags = tf_cuda_tests_tags() + [
+        "notap",
+    ],
+    deps = [
+        ":cuda_op",
+        # copybara:uncomment "//third_party/py/google/protobuf:use_fast_cpp_protos",
+        "//tensorflow:tensorflow_py",
+    ],
+)
+
+py_test(
+    name = "fact_test",
+    size = "small",
+    srcs = ["fact_test.py"],
+    strict_deps = True,
+    deps = [
+        # copybara:uncomment "//third_party/py/google/protobuf:use_fast_cpp_protos",
+        "//tensorflow:tensorflow_py",
+        "//tensorflow/python/framework:test_lib",
+    ],
+)
+
+tf_cc_binary(
+    name = "attr_examples",
+    srcs = ["attr_examples.cc"],
+    deps = [
+        "//tensorflow/core",
+        "//tensorflow/core:framework",
+    ],
+)

tensorflow/examples/adding_an_op/README.md

@@ -0,0 +1,2 @@
+This is the code as described in https://tensorflow.org/guide/create_op.
+See `../custom_ops_doc` for more examples.

tensorflow/examples/adding_an_op/attr_examples.cc

@@ -0,0 +1,46 @@
+/* Copyright 2015 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 <stdio.h>
+#include "tensorflow/core/framework/op.h"
+
+REGISTER_OP("RestrictedTypeExample").Attr("t: {int32, float, bool}");
+
+REGISTER_OP("NumberType").Attr("t: numbertype");
+
+REGISTER_OP("EnumExample").Attr("e: {'apple', 'orange'}");
+
+REGISTER_OP("MinIntExample").Attr("a: int >= 2");
+
+REGISTER_OP("TypeListExample").Attr("a: list({int32, float}) >= 3");
+
+REGISTER_OP("AttrDefaultExample").Attr("i: int = 0");
+
+REGISTER_OP("AttrDefaultExampleForAllTypes")
+    .Attr("s: string = 'foo'")
+    .Attr("i: int = 0")
+    .Attr("f: float = 1.0")
+    .Attr("b: bool = true")
+    .Attr("ty: type = DT_INT32")
+    .Attr("sh: shape = { dim { size: 1 } dim { size: 2 } }")
+    .Attr("te: tensor = { dtype: DT_INT32 int_val: 5 }")
+    .Attr("l_empty: list(int) = []")
+    .Attr("l_int: list(int) = [2, 3, 5, 7]");
+
+int main(int argc, char* argv[]) {
+  printf("All registered ops:\n%s\n",
+         tensorflow::OpRegistry::Global()->DebugString(false).c_str());
+  return 0;
+}

tensorflow/examples/adding_an_op/cuda_op.py

@@ -0,0 +1,23 @@
+# Copyright 2015 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.
+# ==============================================================================
+"""Cuda op Python library."""
+import os.path
+
+import tensorflow as tf
+
+if tf.test.is_built_with_cuda():
+  _cuda_op_module = tf.load_op_library(os.path.join(
+      tf.compat.v1.resource_loader.get_data_files_path(), 'cuda_op_kernel.so'))
+  add_one = _cuda_op_module.add_one

tensorflow/examples/adding_an_op/cuda_op_kernel.cc

@@ -0,0 +1,55 @@
+/* Copyright 2015 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/op.h"
+#include "tensorflow/core/framework/op_kernel.h"
+
+using namespace tensorflow;  // NOLINT(build/namespaces)
+
+REGISTER_OP("AddOne")
+    .Input("input: int32")
+    .Output("output: int32")
+    .Doc(R"doc(
+Adds 1 to all elements of the tensor.
+
+output: A Tensor.
+  output = input + 1
+)doc");
+
+void AddOneKernelLauncher(const int* in, int N, int* out);
+
+class AddOneOp : public OpKernel {
+ public:
+  explicit AddOneOp(OpKernelConstruction* context) : OpKernel(context) {}
+
+  void Compute(OpKernelContext* context) override {
+    // Grab the input tensor
+    const Tensor& input_tensor = context->input(0);
+    auto input = input_tensor.flat<int32_t>();
+
+    // Create an output tensor
+    Tensor* output_tensor = nullptr;
+    OP_REQUIRES_OK(context, context->allocate_output(0, input_tensor.shape(),
+                                                     &output_tensor));
+    auto output = output_tensor->template flat<int32_t>();
+
+    // Set all but the first element of the output tensor to 0.
+    const int N = input.size();
+    // Call the cuda kernel launcher
+    AddOneKernelLauncher(input.data(), N, output.data());
+  }
+};
+
+REGISTER_KERNEL_BUILDER(Name("AddOne").Device(DEVICE_GPU), AddOneOp);

tensorflow/examples/adding_an_op/cuda_op_kernel.cu.cc

@@ -0,0 +1,34 @@
+/* Copyright 2015 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.
+==============================================================================*/
+
+#if GOOGLE_CUDA
+#define EIGEN_USE_GPU
+#include "unsupported/Eigen/CXX11/Tensor"  // from @eigen_archive
+#include "tensorflow/core/util/gpu_kernel_helper.h"
+#include "tensorflow/core/util/gpu_launch_config.h"
+
+__global__ void AddOneKernel(const int* in, const int N, int* out) {
+  for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < N;
+       i += blockDim.x * gridDim.x) {
+    out[i] = in[i] + 1;
+  }
+}
+
+void AddOneKernelLauncher(const int* in, const int N, int* out) {
+  TF_CHECK_OK(::tensorflow::GpuLaunchKernel(AddOneKernel, 32, 256, 0, nullptr,
+                                            in, N, out));
+}
+
+#endif

tensorflow/examples/adding_an_op/cuda_op_test.py

@@ -0,0 +1,30 @@
+# Copyright 2015 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.
+# ==============================================================================
+"""Test for version 1 of the zero_out op."""
+
+import tensorflow as tf
+from tensorflow.examples.adding_an_op import cuda_op
+
+
+class AddOneTest(tf.test.TestCase):
+
+  def test(self):
+    if tf.test.is_built_with_cuda():
+      result = cuda_op.add_one([5, 4, 3, 2, 1])
+      self.assertAllEqual(result, [6, 5, 4, 3, 2])
+
+
+if __name__ == '__main__':
+  tf.test.main()

tensorflow/examples/adding_an_op/fact_test.py

@@ -0,0 +1,30 @@
+# Copyright 2015 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.
+# ==============================================================================
+
+"""Test that user ops can be used as expected."""
+import tensorflow as tf
+from tensorflow.python.framework import test_util
+
+
+class FactTest(tf.test.TestCase):
+
+  @test_util.run_deprecated_v1
+  def test(self):
+    with self.cached_session():
+      print(tf.compat.v1.user_ops.my_fact().eval())
+
+
+if __name__ == '__main__':
+  tf.test.main()

tensorflow/examples/adding_an_op/zero_out_1_test.py

@@ -0,0 +1,51 @@
+# Copyright 2015 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.
+# ==============================================================================
+
+"""Test for version 1 of the zero_out op."""
+
+import os.path
+
+import tensorflow as tf
+from tensorflow.examples.adding_an_op import zero_out_op_1
+
+
+class ZeroOut1Test(tf.test.TestCase):
+
+  def test(self):
+    result = zero_out_op_1.zero_out([5, 4, 3, 2, 1])
+    self.assertAllEqual(result, [5, 0, 0, 0, 0])
+
+  def test_namespace(self):
+    result = zero_out_op_1.namespace_zero_out([5, 4, 3, 2, 1])
+    self.assertAllEqual(result, [5, 0, 0, 0, 0])
+
+  def test_namespace_call_op_on_op(self):
+    x = zero_out_op_1.namespace_zero_out([5, 4, 3, 2, 1])
+    result = zero_out_op_1.namespace_zero_out(x)
+    self.assertAllEqual(result, [5, 0, 0, 0, 0])
+
+  def test_namespace_nested(self):
+    result = zero_out_op_1.namespace_nested_zero_out([5, 4, 3, 2, 1])
+    self.assertAllEqual(result, [5, 0, 0, 0, 0])
+
+  def test_load_twice(self):
+    zero_out_loaded_again = tf.load_op_library(os.path.join(
+        tf.compat.v1.resource_loader.get_data_files_path(),
+        'zero_out_op_kernel_1.so'))
+    self.assertEqual(zero_out_loaded_again, zero_out_op_1._zero_out_module)
+
+
+if __name__ == '__main__':
+  tf.test.main()

tensorflow/examples/adding_an_op/zero_out_2_test.py

@@ -0,0 +1,49 @@
+# Copyright 2015 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.
+# ==============================================================================
+
+"""Test for version 2 of the zero_out op."""
+
+import tensorflow as tf
+
+
+from tensorflow.examples.adding_an_op import zero_out_grad_2  # pylint: disable=unused-import
+from tensorflow.examples.adding_an_op import zero_out_op_2
+
+
+class ZeroOut2Test(tf.test.TestCase):
+
+  def test(self):
+    result = zero_out_op_2.zero_out([5, 4, 3, 2, 1])
+    self.assertAllEqual(result, [5, 0, 0, 0, 0])
+
+  def test_2d(self):
+    result = zero_out_op_2.zero_out([[6, 5, 4], [3, 2, 1]])
+    self.assertAllEqual(result, [[6, 0, 0], [0, 0, 0]])
+
+  def test_grad(self):
+    x = tf.constant([5, 4, 3, 2, 1], dtype=tf.float32)
+    theoretical, numerical = tf.test.compute_gradient(zero_out_op_2.zero_out,
+                                                      tuple([x]))
+    self.assertAllClose(theoretical, numerical)
+
+  def test_grad_2d(self):
+    x = tf.constant([[6, 5, 4], [3, 2, 1]], dtype=tf.float32)
+    theoretical, numerical = tf.test.compute_gradient(zero_out_op_2.zero_out,
+                                                      tuple([x]))
+    self.assertAllClose(theoretical, numerical)
+
+
+if __name__ == '__main__':
+  tf.test.main()

tensorflow/examples/adding_an_op/zero_out_3_test.py

@@ -0,0 +1,42 @@
+# Copyright 2015 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.
+# ==============================================================================
+
+"""Test for version 3 of the zero_out op."""
+
+import tensorflow as tf
+from tensorflow.examples.adding_an_op import zero_out_op_3
+
+
+class ZeroOut3Test(tf.test.TestCase):
+
+  def test(self):
+    result = zero_out_op_3.zero_out([5, 4, 3, 2, 1])
+    self.assertAllEqual(result, [5, 0, 0, 0, 0])
+
+  def test_attr(self):
+    result = zero_out_op_3.zero_out([5, 4, 3, 2, 1], preserve_index=3)
+    self.assertAllEqual(result, [0, 0, 0, 2, 0])
+
+  def test_negative(self):
+    with self.assertRaisesOpError("Need preserve_index >= 0, got -1"):
+      self.evaluate(zero_out_op_3.zero_out([5, 4, 3, 2, 1], preserve_index=-1))
+
+  def test_large(self):
+    with self.assertRaisesOpError("preserve_index out of range"):
+      self.evaluate(zero_out_op_3.zero_out([5, 4, 3, 2, 1], preserve_index=17))
+
+
+if __name__ == "__main__":
+  tf.test.main()

tensorflow/examples/adding_an_op/zero_out_grad_2.py

@@ -0,0 +1,40 @@
+# Copyright 2015 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.
+# ==============================================================================
+
+"""The gradient of the tutorial zero_out op."""
+
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import sparse_ops
+
+
+@ops.RegisterGradient("ZeroOut")
+def _zero_out_grad(op, grad):
+  """The gradients for `zero_out`.
+
+  Args:
+    op: The `zero_out` `Operation` that we are differentiating, which we can use
+      to find the inputs and outputs of the original op.
+    grad: Gradient with respect to the output of the `zero_out` op.
+
+  Returns:
+    Gradients with respect to the input of `zero_out`.
+  """
+  to_zero = op.inputs[0]
+  shape = array_ops.shape(to_zero)
+  index = array_ops.zeros_like(shape)
+  first_grad = array_ops.reshape(grad, [-1])[0]
+  to_zero_grad = sparse_ops.sparse_to_dense([index], shape, first_grad, 0)
+  return [to_zero_grad]  # List of one Tensor, since we have one input

tensorflow/examples/adding_an_op/zero_out_op_1.py

@@ -0,0 +1,25 @@
+# Copyright 2015 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.
+# ==============================================================================
+"""ZeroOut op Python library."""
+import os.path
+
+import tensorflow as tf
+
+_zero_out_module = tf.load_op_library(
+    os.path.join(tf.compat.v1.resource_loader.get_data_files_path(),
+                 'zero_out_op_kernel_1.so'))
+zero_out = _zero_out_module.zero_out
+namespace_zero_out = _zero_out_module.namespace_zero_out
+namespace_nested_zero_out = _zero_out_module.namespace_nested_zero_out

tensorflow/examples/adding_an_op/zero_out_op_2.py

@@ -0,0 +1,25 @@
+# Copyright 2015 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.
+# ==============================================================================
+"""ZeroOut ops Python library."""
+import os.path
+
+import tensorflow as tf
+
+_zero_out_module = tf.load_op_library(
+    os.path.join(tf.compat.v1.resource_loader.get_data_files_path(),
+                 'zero_out_op_kernel_2.so'))
+zero_out = _zero_out_module.zero_out
+zero_out2 = _zero_out_module.zero_out2
+zero_out3 = _zero_out_module.zero_out3

tensorflow/examples/adding_an_op/zero_out_op_3.py

@@ -0,0 +1,23 @@
+# Copyright 2015 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.
+# ==============================================================================
+"""ZeroOut op Python library."""
+import os.path
+
+import tensorflow as tf
+
+_zero_out_module = tf.load_op_library(
+    os.path.join(tf.compat.v1.resource_loader.get_data_files_path(),
+                 'zero_out_op_kernel_3.so'))
+zero_out = _zero_out_module.zero_out

tensorflow/examples/adding_an_op/zero_out_op_kernel_1.cc

@@ -0,0 +1,96 @@
+/* Copyright 2015 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/op.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/shape_inference.h"
+
+using namespace tensorflow;  // NOLINT(build/namespaces)
+
+REGISTER_OP("ZeroOut")
+    .Input("to_zero: int32")
+    .Output("zeroed: int32")
+    .SetShapeFn([](shape_inference::InferenceContext* c) {
+      c->set_output(0, c->input(0));
+      return absl::OkStatus();
+    })
+    .Doc(R"doc(
+Zeros out all but the first value of a Tensor.
+
+zeroed: A Tensor whose first value is identical to `to_zero`, and 0
+  otherwise.
+)doc");
+
+class ZeroOutOp : public OpKernel {
+ public:
+  explicit ZeroOutOp(OpKernelConstruction* context) : OpKernel(context) {}
+
+  void Compute(OpKernelContext* context) override {
+    // Grab the input tensor
+    const Tensor& input_tensor = context->input(0);
+    auto input = input_tensor.flat<int32_t>();
+
+    // Create an output tensor
+    Tensor* output_tensor = nullptr;
+    OP_REQUIRES_OK(context, context->allocate_output(0, input_tensor.shape(),
+                                                     &output_tensor));
+    auto output = output_tensor->template flat<int32_t>();
+
+    // Set all but the first element of the output tensor to 0.
+    const int N = input.size();
+    for (int i = 1; i < N; i++) {
+      output(i) = 0;
+    }
+
+    // Preserve the first input value.
+    if (N > 0) output(0) = input(0);
+  }
+};
+
+REGISTER_KERNEL_BUILDER(Name("ZeroOut").Device(DEVICE_CPU), ZeroOutOp);
+
+REGISTER_OP("Namespace>ZeroOut")
+    .Input("to_zero: int32")
+    .Output("zeroed: int32")
+    .SetShapeFn([](shape_inference::InferenceContext* c) {
+      c->set_output(0, c->input(0));
+      return absl::OkStatus();
+    })
+    .Doc(R"doc(
+Zeros out all but the first value of a Tensor.
+
+zeroed: A Tensor whose first value is identical to `to_zero`, and 0
+  otherwise.
+)doc");
+
+REGISTER_KERNEL_BUILDER(Name("Namespace>ZeroOut").Device(DEVICE_CPU),
+                        ZeroOutOp);
+
+REGISTER_OP("Namespace>Nested>ZeroOut")
+    .Input("to_zero: int32")
+    .Output("zeroed: int32")
+    .SetShapeFn([](shape_inference::InferenceContext* c) {
+      c->set_output(0, c->input(0));
+      return absl::OkStatus();
+    })
+    .Doc(R"doc(
+Zeros out all but the first value of a Tensor.
+
+zeroed: A Tensor whose first value is identical to `to_zero`, and 0
+  otherwise.
+)doc");
+
+REGISTER_KERNEL_BUILDER(Name("Namespace>Nested>ZeroOut").Device(DEVICE_CPU),
+                        ZeroOutOp);

tensorflow/examples/adding_an_op/zero_out_op_kernel_2.cc

@@ -0,0 +1,116 @@
+/* Copyright 2015 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/op.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/register_types.h"
+#include "tensorflow/core/framework/shape_inference.h"
+
+using namespace tensorflow;  // NOLINT(build/namespaces)
+
+REGISTER_OP("ZeroOut")
+    .Attr("T: realnumbertype")
+    .Input("to_zero: T")
+    .Output("zeroed: T")
+    .SetShapeFn(shape_inference::UnchangedShape)
+    .Doc(R"doc(
+Zeros out all but the first value of a Tensor.
+
+zeroed: A Tensor whose first value is identical to `to_zero`, and 0
+  otherwise.
+)doc");
+
+REGISTER_OP("ZeroOut2")
+    .Attr("T: realnumbertype")
+    .Input("to_zero: T")
+    .Output("zeroed: T")
+    .Doc(R"doc(
+Zeros out all but the first value of a Tensor.
+
+zeroed: A Tensor whose first value is identical to `to_zero`, and 0
+  otherwise.
+)doc");
+
+REGISTER_OP("ZeroOut3")
+    .Attr("T: realnumbertype")
+    .Input("to_zero: T")
+    .Output("zeroed: T")
+    .Doc(R"doc(
+Zeros out all but the first value of a Tensor.
+
+zeroed: A Tensor whose first value is identical to `to_zero`, and 0
+  otherwise.
+)doc");
+
+template <typename T>
+class ZeroOutOp : public OpKernel {
+ public:
+  explicit ZeroOutOp(OpKernelConstruction* context) : OpKernel(context) {}
+
+  void Compute(OpKernelContext* context) override {
+    // Grab the input tensor
+    const Tensor& input_tensor = context->input(0);
+    auto input = input_tensor.flat<T>();
+
+    // Create an output tensor
+    Tensor* output = nullptr;
+    OP_REQUIRES_OK(context,
+                   context->allocate_output(0, input_tensor.shape(), &output));
+    auto output_flat = output->template flat<T>();
+
+    // Set all the elements of the output tensor to 0
+    const int N = input.size();
+    for (int i = 0; i < N; i++) {
+      output_flat(i) = T(0);
+    }
+
+    // Preserve the first input value
+    if (N > 0) output_flat(0) = input(0);
+  }
+};
+
+REGISTER_KERNEL_BUILDER(Name("ZeroOut")
+                            .Device(DEVICE_CPU)
+                            .TypeConstraint<float>("T"),
+                        ZeroOutOp<float>);
+REGISTER_KERNEL_BUILDER(Name("ZeroOut")
+                            .Device(DEVICE_CPU)
+                            .TypeConstraint<double>("T"),
+                        ZeroOutOp<double>);
+REGISTER_KERNEL_BUILDER(Name("ZeroOut")
+                            .Device(DEVICE_CPU)
+                            .TypeConstraint<int>("T"),
+                        ZeroOutOp<int>);
+
+#define REGISTER_KERNEL(type)                                       \
+  REGISTER_KERNEL_BUILDER(                                          \
+      Name("ZeroOut2").Device(DEVICE_CPU).TypeConstraint<type>("T"), \
+      ZeroOutOp<type>)
+
+REGISTER_KERNEL(float);
+REGISTER_KERNEL(double);
+REGISTER_KERNEL(int32_t);
+
+#undef REGISTER_KERNEL
+
+#define REGISTER_KERNEL(type)                                       \
+  REGISTER_KERNEL_BUILDER(                                          \
+      Name("ZeroOut3").Device(DEVICE_CPU).TypeConstraint<type>("T"), \
+      ZeroOutOp<type>)
+
+TF_CALL_REAL_NUMBER_TYPES(REGISTER_KERNEL);
+
+#undef REGISTER_KERNEL

tensorflow/examples/adding_an_op/zero_out_op_kernel_3.cc

@@ -0,0 +1,72 @@
+/* Copyright 2015 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/op.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/shape_inference.h"
+
+using namespace tensorflow;  // NOLINT(build/namespaces)
+
+REGISTER_OP("ZeroOut")
+    .Attr("preserve_index: int = 0")
+    .Input("to_zero: int32")
+    .Output("zeroed: int32")
+    .SetShapeFn([](shape_inference::InferenceContext* c) {
+      c->set_output(0, c->input(0));
+      return absl::OkStatus();
+    });
+
+class ZeroOutOp : public OpKernel {
+ public:
+  explicit ZeroOutOp(OpKernelConstruction* context) : OpKernel(context) {
+    // Get the index of the value to preserve
+    OP_REQUIRES_OK(context,
+                   context->GetAttr("preserve_index", &preserve_index_));
+    // Check that preserve\_index is positive
+    OP_REQUIRES(context, preserve_index_ >= 0,
+                absl::InvalidArgumentError(absl::StrCat(
+                    "Need preserve_index >= 0, got ", preserve_index_)));
+  }
+
+  void Compute(OpKernelContext* context) override {
+    // Grab the input tensor
+    const Tensor& input_tensor = context->input(0);
+    auto input = input_tensor.flat<int32_t>();
+
+    // Check that preserve_index is in range
+    OP_REQUIRES(context, preserve_index_ < input.dimension(0),
+                absl::InvalidArgumentError("preserve_index out of range"));
+
+    // Create an output tensor
+    Tensor* output_tensor = nullptr;
+    OP_REQUIRES_OK(context, context->allocate_output(0, input_tensor.shape(),
+                                                     &output_tensor));
+    auto output = output_tensor->template flat<int32_t>();
+
+    // Set all the elements of the output tensor to 0
+    const int N = input.size();
+    for (int i = 0; i < N; i++) {
+      output(i) = 0;
+    }
+
+    // Preserve the requested input value
+    output(preserve_index_) = input(preserve_index_);
+  }
+
+ private:
+  int preserve_index_;
+};
+
+REGISTER_KERNEL_BUILDER(Name("ZeroOut").Device(DEVICE_CPU), ZeroOutOp);

tensorflow/examples/android/.gitignore

@@ -0,0 +1,30 @@
+# This file is based on https://github.com/github/gitignore/blob/master/Android.gitignore
+*.iml
+.idea/compiler.xml
+.idea/copyright
+.idea/dictionaries
+.idea/gradle.xml
+.idea/libraries
+.idea/inspectionProfiles
+.idea/misc.xml
+.idea/modules.xml
+.idea/runConfigurations.xml
+.idea/tasks.xml
+.idea/workspace.xml
+.gradle
+local.properties
+.DS_Store
+build/
+gradleBuild/
+*.apk
+*.ap_
+*.dex
+*.class
+bin/
+gen/
+out/
+*.log
+.navigation/
+/captures
+.externalNativeBuild
+.cxx

tensorflow/examples/custom_ops_doc/README.md

@@ -0,0 +1,10 @@
+# Custom Ops Examples
+
+The following subdirectories have examples of Custom Ops.
+
+* multiplex_1: starting example, similar to np.where
+* multiplex_2: GPU (and CPU) example, similar to np.where
+* multiplex_3: dispatch to special case kernel (and sparse tensors)
+* multiplex_4: C++ backward compatibility example: lists of tensors
+* simple_hash_table: internal state using a Resource using ref-counting
+* sleep: Asynchronous (non-blocking) sleep op using AsyncOpKernel

tensorflow/examples/custom_ops_doc/multiplex_1/BUILD

@@ -0,0 +1,43 @@
+# Build multiplex_1 custom ops examples, which is similar to np.where
+
+load("//tensorflow:strict.default.bzl", "py_strict_library")
+load("//tensorflow:tensorflow.bzl", "tf_custom_op_library")
+load("//tensorflow:tensorflow.default.bzl", "tf_py_test")
+
+licenses(["notice"])
+
+tf_custom_op_library(
+    name = "multiplex_1_kernel.so",
+    srcs = [
+        "multiplex_1_kernel.cc",
+        "multiplex_1_op.cc",
+    ],
+)
+
+py_strict_library(
+    name = "multiplex_1_op",
+    srcs = ["multiplex_1_op.py"],
+    data = [":multiplex_1_kernel.so"],
+    srcs_version = "PY3",
+    deps = [
+        "//tensorflow:tensorflow_py",
+    ],
+)
+
+tf_py_test(
+    name = "multiplex_1_test",
+    size = "medium",
+    srcs = ["multiplex_1_test.py"],
+    python_version = "PY3",
+    srcs_version = "PY3",
+    tags = [
+        "no_mac",  # TODO(b/216321151): Re-enable this test.
+    ],
+    deps = [
+        ":multiplex_1_op",
+        "//tensorflow:tensorflow_py",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:test_lib",
+        "//third_party/py/numpy",
+    ],
+)

tensorflow/examples/custom_ops_doc/multiplex_1/README.md

@@ -0,0 +1,401 @@
+<!-- LINT.IfChange -->
+# Create a custom multiplexer op
+
+This page provides an end-to-end example for adding a custom multiplexer op to
+TensorFlow. For additional context,
+read the
+[OSS guide on creating custom ops](https://www.tensorflow.org/guide/create_op).
+
+## Creating a custom multiplexer op
+
+This examples demonstrates how you can create a Python custom multiplexer
+`multiplex_1_op`, similar to
+[`tf.where`](https://tensorflow.org/api_docs/python/tf/where?version=nightly)
+which you can call as:
+
+<!-- test_snippets_in_readme skip -->
+```python
+multiplex_1_op.multiplex(condition, x, y)                                        # doctest: skip
+```
+
+This custom op returns elements chosen from either of the two input tensors `x`
+or `y` depending on the `condition`.
+
+Example usage:
+
+<!-- test_snippets_in_readme skip -->
+```python
+from tensorflow.examples.custom_ops_doc.multiplex_1 import multiplex_1_op
+
+m = multiplex_1_op.multiplex([True, False, False, True], [1,2,3,4], [100,200,300,400])
+m.numpy()
+```
+
+<!-- test_snippets_in_readme skip -->
+```
+array([  1, 200, 300,   4], dtype=int32)
+```
+
+Note that this simplified `multiplex_1` op has limitations that are not present
+in `tf.where` such as:
+
+*   Support only for dense tensors
+*   Support only for CPU computations
+*   No broadcasting capabilities
+*   No extensibility through optional parameters
+
+The example below contains C++ and Python code snippets to illustrate the code
+flow. These snippets are not all complete; some are missing namespace
+declarations, imports, and test cases.
+
+### Step 1 - Define op interface
+
+Define the op interface and register it using the `REGISTER_OP` macro.
+
+<!-- test_snippets_in_readme skip -->
+```
+#include "tensorflow/core/framework/op.h"
+#include "tensorflow/core/framework/shape_inference.h"
+```
+
+```
+REGISTER_OP("Examples1>MultiplexDense")
+    .Input("cond: bool")
+    .Input("a_values: T")
+    .Input("b_values: T")
+    .Output("output_values: T")
+    .Attr("T: type")
+    .SetShapeFn([](tensorflow::shape_inference::InferenceContext* c) {
+      // Determine the output shape and also assert that inputs 0 and 1 have
+      // the same shape.
+      tensorflow::shape_inference::ShapeHandle out;
+      TF_RETURN_IF_ERROR(c->Merge(c->input(0), c->input(1), &out));
+      // Assert that inputs 0 and 2 have the same shape, i.e. that all inputs
+      // have the same shape. This is optional, but it is desirable
+      // to raise errors about inconsistent input shapes early when using
+      // graph mode.
+      tensorflow::shape_inference::ShapeHandle unused;
+      TF_RETURN_IF_ERROR(c->Merge(c->input(0), c->input(2), &unused));
+
+      c->set_output(0, out);
+      return ::tensorflow::OkStatus();
+    })
+    .Doc(R"doc(
+Return elements chosen from `a` or `b` depending on `cond`.
+
+This is similar to `np.where` and `tf.where`, but simplified to only handle
+the case of dense tensors, no optional parameters, no broadcasting, etc..
+
+cond: tf.Tensor of type bool.
+a_values: tf.Tensor with the same type and shape as `b_values`.
+b_values: tf.Tensor with the same type and shape as `a_values`.
+
+      Where True, yield `a_values`, otherwise yield `b_values`.
+output_values: A tf.Tensor with elements from `a` where `cond` is True, and
+               elements from `b` elsewhere.
+)doc");
+```
+
+Note that:
+
+*   This op has three input tensors - one boolean tensor for selecting which
+    values to choose from the two other input tensors of matching type `T`, and
+    one output tensor of type `T`.
+*   The `Attr` for this op is defined as `.Attr("T: type")` which specifies `T`
+    as an `Attr` of type `type`. In the subsequent steps, you will use `T` with
+    a template class to define the type of the contents of tensors.
+*   The docstring for this op is specified by passing a string to `.Doc()`.
+*   The shape function for this op uses the `Merge` method of the
+    [`tensorflow::shape_inference::InferenceContext`](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/framework/shape_inference.h#:~:text=class%20InferenceContext)
+    object which is a helper function to set the output shape to be the same as
+    the identical shapes of the two inputs (for example, if it is used for
+    binary ops) and has error checking that the two inputs have the same shape.
+    Since `multiplex_1` has three inputs, two calls to `Merge` are used to
+    assert that all three inputs are the same shape.
+
+### Step 2 - Register the op implementation (kernel)
+
+Register the kernel by calling the `REGISTER_KERNEL_BUILDER` macro.
+
+```
+#define REGISTER_KERNELS(type)                                  \
+  REGISTER_KERNEL_BUILDER(Name("Examples1>MultiplexDense")      \
+                              .Device(::tensorflow::DEVICE_CPU) \
+                              .TypeConstraint<type>("T"),       \
+                          MultiplexDenseOp<type>)
+
+TF_CALL_ALL_TYPES(REGISTER_KERNELS);
+#undef REGISTER_KERNELS
+```
+
+### Step 3 - Implement the op kernel(s)
+
+In the op kernel in `multiplex_1_kernel.cc`, create a class derived from
+[`OpKernel`](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/framework/op_kernel.h#:~:text=class%20OpKernel)
+that implements a `Compute` method to get and validate input tensors, perform
+computation, and create the output tensors.
+
+```
+template <typename T>
+class MultiplexDenseOp : public OpKernel {
+ public:
+  explicit MultiplexDenseOp(OpKernelConstruction* ctx) : OpKernel(ctx) {}
+  MultiplexDenseOp(const MultiplexDenseOp& other) = delete;
+  MultiplexDenseOp& operator=(const MultiplexDenseOp& other) = delete;
+  ~MultiplexDenseOp() override = default;
+
+  void Compute(OpKernelContext* ctx) override {
+    const auto& cond_tensor = ctx->input(0);
+    const auto& a_values_tensor = ctx->input(1);
+    const auto& b_values_tensor = ctx->input(2);
+
+    // Allow any shape, but require that a_values, b_values, and cond all
+    // have the same shape.
+    // Note that ::tensorflow::TensorShapeUtils has some useful functions
+    // for checking shapes.
+    OP_REQUIRES(ctx, a_values_tensor.shape() == b_values_tensor.shape(),
+                InvalidArgument(
+                    "a_values and b_values must have the same shape. "
+                    "a_values shape: ",
+                    a_values_tensor.shape().DebugString(), " b_values shape: ",
+                    b_values_tensor.shape().DebugString()));
+    OP_REQUIRES(
+        ctx, a_values_tensor.shape() == cond_tensor.shape(),
+        InvalidArgument("a_values and cond must have the same shape. "
+                        "a_values shape: ",
+                        a_values_tensor.shape().DebugString(),
+                        " cond shape: ", cond_tensor.shape().DebugString()));
+
+    const auto a_values = a_values_tensor.flat<T>();
+    const auto b_values = b_values_tensor.flat<T>();
+    const auto cond = cond_tensor.flat<bool>();
+
+    // Create an output tensor
+    Tensor* output_tensor = nullptr;
+    OP_REQUIRES_OK(
+        ctx, ctx->allocate_output(0, a_values_tensor.shape(), &output_tensor));
+    auto output = output_tensor->template flat<T>();
+    const int64_t N = a_values_tensor.NumElements();
+
+    // Here is an example of processing tensors in a simple loop directly
+    // without relying on any libraries. For intensive math operations, it is
+    // a good practice to use libraries such as Eigen that support
+    // tensors when possible, e.g. "output = cond.select(a_values, b_values);"
+    // Eigen supports chunking into blocks and multi-threading.
+    // See
+    // https://eigen.tuxfamily.org/dox/unsupported/eigen_tensors.html#title55
+    for (int64_t i = 0; i < N; i++) {
+      if (cond(i)) {
+        output(i) = a_values(i);
+      } else {
+        output(i) = b_values(i);
+      }
+    }
+  }
+};
+```
+
+A common way to access the values in tensors for manipulation is to get
+flattened rank-1
+[`Eigen::Tensor`](https://eigen.tuxfamily.org/dox/unsupported/eigen_tensors.html)
+objects. In the example code, this is done for all three inputs and the output.
+The example also processes tensors in a simple loop directly without relying on
+any libraries.
+
+Using Eigen, the `for` loop above could have been written simply as:
+
+<!-- test_snippets_in_readme skip -->
+```c++
+output = cond.select(a_values, b_values);
+```
+
+[Selection](https://eigen.tuxfamily.org/dox/unsupported/eigen_tensors.html#title55)
+from Eigen supports chunking into blocks and multi-threading.
+
+For intensive mathematical operations, it is a good practice to use libraries
+such as [Eigen](https://eigen.tuxfamily.org/index.php?title=Main_Page#Overview)
+that support tensors to do the computation when possible. Eigen is vectorized,
+avoids dynamic memory allocation and therefore is typically faster than using
+simple `for` loops.
+
+Eigen provides the following for accessing tensor values (for both inputs and
+outputs):
+
+*   `flat<T>()(index)` for element access for tensors of any rank
+*   `scalar<T>()()` for rank 0 tensors
+*   `vec<T>()(index)` for rank 1 tensors
+*   `matrix<T>()(i, j)` for rank 2 tensors
+*   `tensor<T, 3>()(i, j, k)` for tensors of known rank (e.g. 3).
+
+#### Compile the op (optional)
+
+Compile the C++ op to create a kernel library and Python wrapper that enables
+you to use the op with TensorFlow.
+
+Create a `BUILD` file for the op which declares the dependencies and the output
+build targets. Refer to
+[building for OSS](https://www.tensorflow.org/guide/create_op#build_the_op_library).
+
+### Step 4 - Create the Python wrapper (optional)
+
+To create the Python wrapper, import and implement a function that serves as the
+op's public API and provides a docstring.
+
+```
+def multiplex(cond, a, b, name=None):
+  """Return elements chosen from `a` or `b` depending on `cond`.
+
+  This is similar to `np.where` and `tf.where`, but simplified to only handle
+  the case of dense tensors, no optional parameters, no broadcasting, etc..
+
+  >>> multiplex([True, False, False, True], [1,2,3,4], [100,200,300,400])
+  <tf.Tensor: shape=(4,), dtype=int32, numpy=array([  1, 200, 300,   4], ...)>
+
+  Args:
+    cond: tf.Tensor of type bool. Where True, yield `a`, otherwise yield `b`.
+    a: tf.Tensor with the same type and shape as `b`.
+    b: tf.Tensor with the same type and shape as `a`.
+    name: An optional name for the op.
+
+  Returns:
+    A tf.Tensor with elements from `a` where `cond` is True, and elements
+    from `b` elsewhere.
+  """
+  return gen_multiplex_1_op.examples1_multiplex_dense(
+      cond=cond, a_values=a, b_values=b, name=name)
+```
+
+### Step 5 - Test the op
+
+Create op tests using classes derived from
+[`tf.test.TestCase`](https://www.tensorflow.org/api_docs/python/tf/test/TestCase).
+
+When writing tests to ensure that the op works correctly in both graph and eager
+executions, it is important to note that errors in the op code may be detected
+in two distinct phases of code execution depending on how it is executed (eager
+or graph executions). Errors may be detected early by the shape function or a
+bit later from the logic in the `Compute` method. This may lead to differing
+error types and/or messages.
+
+Below are test excerpts showing how to handle errors for different scenarios.
+The first test case demonstrates error handling when errors are common across
+eager and graph executions and the second test case demonstrates error handling
+when the errors are different in eager and graph executions.
+
+```
+  @test_util.run_in_graph_and_eager_modes
+  def test_multiplex_int(self):
+    a = tf.constant([1, 2, 3, 4, 5])
+    b = tf.constant([10, 20, 30, 40, 50])
+    cond = tf.constant([True, False, True, False, True], dtype=bool)
+    expect = np.where(self.evaluate(cond), self.evaluate(a), self.evaluate(b))
+    # expected result is [1, 20, 3, 40, 5]
+    result = multiplex_1_op.multiplex(cond, a, b)
+    self.assertAllEqual(result, expect)
+```
+
+```
+  @test_util.run_in_graph_and_eager_modes
+  def test_multiplex_bad_types(self):
+    a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0])  # float
+    b = tf.constant([10, 20, 30, 40, 50])  # int32
+    cond = tf.constant([True, False, True, False, True], dtype=bool)
+    with self.assertRaisesRegex(
+        (errors_impl.InvalidArgumentError, TypeError),
+        # Eager mode raises InvalidArgumentError with the following message
+        r'(cannot compute Examples1>MultiplexDense as input #2\(zero-based\) '
+        r'was expected to be a float tensor but is a int32 tensor '
+        r'\[Op:Examples1>MultiplexDense\]'
+        r')|('
+        # Graph mode raises TypeError with the following message
+        r"Input 'b_values' of 'Examples1>MultiplexDense' Op has type int32 that "
+        r"does not match type float32 of argument 'a_values'.)"):
+      self.evaluate(multiplex_1_op.multiplex(cond, a, b))
+```
+
+Refer to `multiplex_1_test.py` for the full source code which contains all the
+test cases.
+
+Reuse the `BUILD` file created in Step 3a above to add build
+rules for the Python API wrapper and the op test.
+
+<!-- test_snippets_in_readme skip -->
+```
+load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
+load("//third_party/tensorflow:tensorflow.default.bzl", "tf_py_test")
+```
+
+```
+py_library(
+    name = "multiplex_1_op",
+    srcs = ["multiplex_1_op.py"],
+    strict_deps = True,
+    visibility = ["//third_party/tensorflow/google/g3doc:__subpackages__"],
+    deps = [
+        ":gen_multiplex_1_op",
+        ":multiplex_1_kernel",
+        "//third_party/py/tensorflow",
+    ],
+)
+
+tf_py_strict_test(
+    name = "multiplex_1_test",
+    size = "medium",
+    srcs = ["multiplex_1_test.py"],
+    deps = [
+        ":multiplex_1_op",
+        "//third_party/py/numpy",
+        "//third_party/py/tensorflow",
+        "//third_party/tensorflow/python/framework:errors",
+        "//third_party/tensorflow/python/framework:test_lib",
+    ],
+)
+```
+
+Test the op by running:
+
+<!-- test_snippets_in_readme skip -->
+```shell
+$ bazel test //third_party/tensorflow/google/g3doc/example/multiplex_1:multiplex_1_test
+```
+
+### Use the op
+
+Use the op by importing and calling it as follows:
+
+<!-- test_snippets_in_readme skip -->
+```python
+import tensorflow as tf
+from tensorflow.examples.custom_ops_doc.multiplex_1 import multiplex_1_op
+
+a = tf.constant([1, 2, 3, 4, 5])
+b = tf.constant([10, 20, 30, 40, 50])
+cond = tf.constant([True, False, True, False, True], dtype=bool)
+
+result = multiplex_1_op.multiplex(cond, a, b)
+result.numpy()
+```
+
+<!-- test_snippets_in_readme skip -->
+```
+ array([ 1, 20,  3, 40,  5], dtype=int32)
+```
+
+Here, `multiplex_1_op` is the name of the Python wrapper that was created in
+this example.
+
+### Summary
+
+In this example, you learned how to define and use a custom multiplexer op. The
+image below summarizes the files created for this op.
+
+The table below summarizes the build rules and targets for building and testing
+the `multiplex_1` op.
+
+Op components                           | Build rule             | Build target         | Source
+--------------------------------------- | ---------------------- | -------------------- | ------
+Kernels (C++)                           | `tf_custom_op_library` | `multiplex_1_kernel` | `multiplex_1_kernel.cc`, `multiplex_1_op.cc`
+Wrapper (automatically generated)       | N/A                    | `gen_multiplex_1_op` | N/A
+Wrapper (with public API and docstring) | `py_strict_library`    | `multiplex_1_op`     | `multiplex_1_op.py`
+Tests                                   | `tf_py_strict_test`    | `multiplex_1_test`   | `multiplex_1_test.py`
+<!-- LINT.ThenChange(multiplex_1.md) -->

tensorflow/examples/custom_ops_doc/multiplex_1/multiplex_1_kernel.cc

@@ -0,0 +1,110 @@
+/* Copyright 2021 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 "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/op_requires.h"
+#include "tensorflow/core/framework/register_types.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/types.h"
+
+// Please use the appropriate namespace for your project
+namespace tensorflow {
+namespace custom_op_examples {
+
+using ::tensorflow::OpKernel;
+using ::tensorflow::OpKernelConstruction;
+using ::tensorflow::OpKernelContext;
+using ::tensorflow::Tensor;
+using ::tensorflow::errors::InvalidArgument;
+
+// Multiple types for the values inside two of the input tensors
+// (e.g. int32, float) are supported by using a template where the type is T.
+template <typename T>
+class MultiplexDenseOp : public OpKernel {
+ public:
+  explicit MultiplexDenseOp(OpKernelConstruction* ctx) : OpKernel(ctx) {}
+  MultiplexDenseOp(const MultiplexDenseOp& other) = delete;
+  MultiplexDenseOp& operator=(const MultiplexDenseOp& other) = delete;
+  ~MultiplexDenseOp() override = default;
+
+  void Compute(OpKernelContext* ctx) override {
+    const auto& cond_tensor = ctx->input(0);
+    const auto& a_values_tensor = ctx->input(1);
+    const auto& b_values_tensor = ctx->input(2);
+
+    // Allow any shape, but require that a_values, b_values, and cond all
+    // have the same shape.
+    // Note that ::tensorflow::TensorShapeUtils has some useful functions
+    // for checking shapes.
+    OP_REQUIRES(ctx, a_values_tensor.shape() == b_values_tensor.shape(),
+                InvalidArgument(
+                    "a_values and b_values must have the same shape. "
+                    "a_values shape: ",
+                    a_values_tensor.shape().DebugString(), " b_values shape: ",
+                    b_values_tensor.shape().DebugString()));
+    OP_REQUIRES(
+        ctx, a_values_tensor.shape() == cond_tensor.shape(),
+        InvalidArgument("a_values and cond must have the same shape. "
+                        "a_values shape: ",
+                        a_values_tensor.shape().DebugString(),
+                        " cond shape: ", cond_tensor.shape().DebugString()));
+
+    const auto a_values = a_values_tensor.flat<T>();
+    const auto b_values = b_values_tensor.flat<T>();
+    const auto cond = cond_tensor.flat<bool>();
+
+    // Create an output tensor
+    Tensor* output_tensor = nullptr;
+    OP_REQUIRES_OK(
+        ctx, ctx->allocate_output(0, a_values_tensor.shape(), &output_tensor));
+    auto output = output_tensor->template flat<T>();
+    const int64_t N = a_values_tensor.NumElements();
+
+    // Here is an example of processing tensors in a simple loop directly
+    // without relying on any libraries. For intensive math operations, it is
+    // a good practice to use libraries such as Eigen that support
+    // tensors when possible, e.g. "output = cond.select(a_values, b_values);"
+    // Eigen supports chunking into blocks and multi-threading.
+    // See
+    // https://eigen.tuxfamily.org/dox/unsupported/eigen_tensors.html#title55
+    for (int64_t i = 0; i < N; i++) {
+      if (cond(i)) {
+        output(i) = a_values(i);
+      } else {
+        output(i) = b_values(i);
+      }
+    }
+  }
+};
+
+// The "Name" used by REGISTER_KERNEL_BUILDER is defined by REGISTER_OP,
+// see multiplex_1_op.cc.
+// To support tensors containing different types (e.g. int32, float), one
+// kernel per type is registered and is templatized by the "T" attr value.
+// The TF_CALL_ALL_TYPES macro registers the op for all types appropriate for
+// the target platform. See go/tf-custom-ops-guide
+#define REGISTER_KERNELS(type)                                  \
+  REGISTER_KERNEL_BUILDER(Name("Examples1>MultiplexDense")      \
+                              .Device(::tensorflow::DEVICE_CPU) \
+                              .TypeConstraint<type>("T"),       \
+                          MultiplexDenseOp<type>)
+
+TF_CALL_ALL_TYPES(REGISTER_KERNELS);
+#undef REGISTER_KERNELS
+
+}  // namespace custom_op_examples
+}  // namespace tensorflow

tensorflow/examples/custom_ops_doc/multiplex_1/multiplex_1_op.cc

@@ -0,0 +1,61 @@
+/* Copyright 2021 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 "xla/tsl/platform/errors.h"
+#include "tensorflow/core/framework/op.h"
+#include "tensorflow/core/framework/shape_inference.h"
+#include "tensorflow/core/platform/status.h"
+
+// Use a namespace when registering by prepending the
+// package's name to the op’s name and separate with a '>'.
+// This is the recommendation for out-of-tree ops to avoid name collisions in
+// "Best practices for custom operations in TensorFlow"
+// https://github.com/tensorflow/community/blob/master/rfcs/20190726-custom-ops.md
+
+REGISTER_OP("Examples1>MultiplexDense")
+    .Input("cond: bool")
+    .Input("a_values: T")
+    .Input("b_values: T")
+    .Output("output_values: T")
+    .Attr("T: type")
+    .SetShapeFn([](tensorflow::shape_inference::InferenceContext* c) {
+      // Determine the output shape and also assert that inputs 0 and 1 have
+      // the same shape.
+      tensorflow::shape_inference::ShapeHandle out;
+      TF_RETURN_IF_ERROR(c->Merge(c->input(0), c->input(1), &out));
+      // Assert that inputs 0 and 2 have the same shape, i.e. that all inputs
+      // have the same shape. This is optional, but it is desirable
+      // to raise errors about inconsistent input shapes early when using
+      // graph mode.
+      tensorflow::shape_inference::ShapeHandle unused;
+      TF_RETURN_IF_ERROR(c->Merge(c->input(0), c->input(2), &unused));
+
+      c->set_output(0, out);
+      return ::tensorflow::OkStatus();
+    })
+    .Doc(R"doc(
+Return elements chosen from `a` or `b` depending on `cond`.
+
+This is similar to `np.where` and `tf.where`, but simplified to only handle
+the case of dense tensors, no optional parameters, no broadcasting, etc..
+
+cond: tf.Tensor of type bool.
+a_values: tf.Tensor with the same type and shape as `b_values`.
+b_values: tf.Tensor with the same type and shape as `a_values`.
+
+      Where True, yield `a_values`, otherwise yield `b_values`.
+output_values: A tf.Tensor with elements from `a` where `cond` is True, and
+               elements from `b` elsewhere.
+)doc");

tensorflow/examples/custom_ops_doc/multiplex_1/multiplex_1_op.py

@@ -0,0 +1,54 @@
+# Copyright 2021 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.
+# ==============================================================================
+"""A wrapper for gen_multiplex_1_op.py.
+
+This defines a public API and provides a docstring for the C++ Op defined by
+multiplex_1_kernel.cc
+"""
+
+import tensorflow as tf
+from tensorflow.python.platform import resource_loader
+
+_multiplex_1_module = tf.load_op_library(
+    resource_loader.get_path_to_datafile("multiplex_1_kernel.so"))
+
+examples_multiplex_dense = _multiplex_1_module.examples1_multiplex_dense
+
+# In this example, this Python function is a trivial wrapper for the C++ Op:
+# it provides a public API and docstring that are equivalent to the API
+# and documentation of the C++ op. The motivation for it is to be a placeholder
+# that allows a wider variety of non-breaking future changes than are possible
+# with the generated wrapper alone. Having this wrapper is optional.
+def multiplex(cond, a, b, name=None):
+  """Return elements chosen from `a` or `b` depending on `cond`.
+
+  This is similar to `np.where` and `tf.where`, but simplified to only handle
+  the case of dense tensors, no optional parameters, no broadcasting, etc..
+
+  >>> multiplex([True, False, False, True], [1,2,3,4], [100,200,300,400])
+  <tf.Tensor: shape=(4,), dtype=int32, numpy=array([  1, 200, 300,   4], ...)>
+
+  Args:
+    cond: tf.Tensor of type bool. Where True, yield `a`, otherwise yield `b`.
+    a: tf.Tensor with the same type and shape as `b`.
+    b: tf.Tensor with the same type and shape as `a`.
+    name: An optional name for the op.
+
+  Returns:
+    A tf.Tensor with elements from `a` where `cond` is True, and elements
+    from `b` elsewhere.
+  """
+  return examples_multiplex_dense(
+      cond=cond, a_values=a, b_values=b, name=name)

tensorflow/examples/custom_ops_doc/multiplex_1/multiplex_1_test.py

@@ -0,0 +1,111 @@
+# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for multiplex_1."""
+
+import numpy as np
+import tensorflow as tf
+
+from tensorflow.examples.custom_ops_doc.multiplex_1 import multiplex_1_op
+# This pylint disable is only needed for internal google users
+from tensorflow.python.framework import errors_impl  # pylint: disable=g-direct-tensorflow-import
+from tensorflow.python.framework import test_util  # pylint: disable=g-direct-tensorflow-import
+
+
+@test_util.with_eager_op_as_function
+class MultiplexOpRank1Test(tf.test.TestCase):
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_multiplex_int(self):
+    a = tf.constant([1, 2, 3, 4, 5])
+    b = tf.constant([10, 20, 30, 40, 50])
+    cond = tf.constant([True, False, True, False, True], dtype=bool)
+    expect = np.where(self.evaluate(cond), self.evaluate(a), self.evaluate(b))
+    # expected result is [1, 20, 3, 40, 5]
+    result = multiplex_1_op.multiplex(cond, a, b)
+    self.assertAllEqual(result, expect)
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_multiplex_float(self):
+    a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0])
+    b = tf.constant([10.0, 20.0, 30.0, 40.0, 50.0])
+    cond = tf.constant([True, False, True, False, True], dtype=bool)
+    # expected result is [1.0, 20.0, 3.0, 40.0, 5.0]
+    expect = np.where(self.evaluate(cond), self.evaluate(a), self.evaluate(b))
+    result = multiplex_1_op.multiplex(cond, a, b)
+    self.assertAllEqual(result, expect)
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_multiplex_bad_types(self):
+    a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0])  # float
+    b = tf.constant([10, 20, 30, 40, 50])  # int32
+    cond = tf.constant([True, False, True, False, True], dtype=bool)
+    with self.assertRaisesRegex(
+        (errors_impl.InvalidArgumentError, TypeError),
+        # Eager mode raises InvalidArgumentError with the following message
+        r'(cannot compute Examples1>MultiplexDense as input #2\(zero-based\) '
+        r'was expected to be a float tensor but is a int32 tensor '
+        r'\[Op:Examples1>MultiplexDense\]'
+        r')|('
+        # Graph mode raises TypeError with the following message
+        r"Input 'b_values' of 'Examples1>MultiplexDense' Op has type int32 that "
+        r"does not match type float32 of argument 'a_values'.)"):
+      self.evaluate(multiplex_1_op.multiplex(cond, a, b))
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_multiplex_bad_size(self):
+    a = tf.constant([1, 2, 3, 4, 5])  # longer than b
+    b = tf.constant([10, 20])  # shorter than a
+    cond = tf.constant([True, False, True, False, True], dtype=bool)
+    with self.assertRaisesRegex(
+        (errors_impl.InvalidArgumentError, ValueError),
+        # Eager mode raises InvalidArgumentError with the following message
+        r'(?s)(a_values and b_values must have the same shape. '
+        r'a_values shape: \[5\] b_values shape: \[2\].* '
+        r'\[Op:Examples1>MultiplexDense\]'
+        r')|('
+        # Graph mode raises ValueError with the following message
+        r'Dimension 0 in both shapes must be equal, but are 5 and 2\. '
+        r'Shapes are \[5\] and \[2\]\.)'):
+      self.evaluate(multiplex_1_op.multiplex(cond, a, b))
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_multiplex_2d(self):
+    a = tf.constant([[1, 2, 3], [4, 5, 6]])
+    b = tf.constant([[10, 20, 30], [40, 50, 60]])
+    cond = tf.constant([[True, False, True], [False, True, False]], dtype=bool)
+    expect = np.where(self.evaluate(cond), self.evaluate(a), self.evaluate(b))
+    # expected result is [[1, 20], [3, 40]]
+    result = multiplex_1_op.multiplex(cond, a, b)
+    self.assertAllEqual(result, expect)
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_multiplex_bad_shape(self):
+    a = tf.constant([[1, 2, 3], [4, 5, 6]])  # shape (2,3)
+    b = tf.constant([[10, 20], [30, 40], [50, 60]])  # shape (3,2)
+    cond = tf.constant([[True, False, True], [False, True, False]], dtype=bool)
+    with self.assertRaisesRegex(
+        (errors_impl.InvalidArgumentError, ValueError),
+        # Eager mode raises InvalidArgumentError with the following message
+        r'(a_values and b_values must have the same shape.'
+        r' a_values shape: \[2,3\] b_values shape: \[3,2\]'
+        r')|('
+        # Graph mode raises ValueError with the following message
+        r'Dimension 0 in both shapes must be equal, '
+        r'but are 2 and 3\. Shapes are \[2,3\] and \[3,2\])\.'):
+      self.evaluate(multiplex_1_op.multiplex(cond, a, b))
+
+
+if __name__ == '__main__':
+  tf.test.main()

tensorflow/examples/custom_ops_doc/multiplex_2/BUILD

@@ -0,0 +1,52 @@
+# Build multiplex_2 custom ops example, which is similar to np.where.
+# This example supports GPU (and CPU), in contrast to multiplex_1 which
+# only supports CPU.
+
+load("//tensorflow:strict.default.bzl", "py_strict_library")
+load("//tensorflow:tensorflow.bzl", "tf_custom_op_library")
+load("//tensorflow:tensorflow.default.bzl", "cuda_py_test")
+
+licenses(["notice"])
+
+tf_custom_op_library(
+    name = "multiplex_2_kernel.so",
+    srcs = [
+        "multiplex_2_kernel.cc",
+        "multiplex_2_kernel.h",
+        "multiplex_2_op.cc",
+    ],
+    gpu_srcs = [
+        "multiplex_2_kernel.h",
+        "multiplex_2_kernel.cu.cc",
+    ],
+    visibility = ["//tensorflow/examples/custom_ops_doc:__subpackages__"],
+)
+
+py_strict_library(
+    name = "multiplex_2_op",
+    srcs = ["multiplex_2_op.py"],
+    data = ["multiplex_2_kernel.so"],
+    srcs_version = "PY3",
+    visibility = ["//tensorflow/examples/custom_ops_doc:__subpackages__"],
+    deps = [
+        "//tensorflow:tensorflow_py",
+    ],
+)
+
+cuda_py_test(
+    name = "multiplex_2_test",
+    size = "medium",
+    srcs = ["multiplex_2_test.py"],
+    python_version = "PY3",
+    srcs_version = "PY3",
+    tags = [
+        "no_mac",  # TODO(b/216321151): Re-enable this test.
+    ],
+    deps = [
+        ":multiplex_2_op",
+        "//tensorflow:tensorflow_py",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:test_lib",
+        "//third_party/py/numpy",
+    ],
+)

tensorflow/examples/custom_ops_doc/multiplex_2/README.md

@@ -0,0 +1,415 @@
+<!-- LINT.IfChange -->
+# Create a custom multiplexer op with GPU support
+
+This guide provides an end-to-end example for adding a custom multiplexer op
+with both CPU and GPU support.
+
+For a simpler example of a TensorFlow multiplexer custom op, refer to
+`multiplex_1`. The `multiplex_2` operation builds on the `multiplex_1` operation
+in the following ways:
+
+*   This op includes support for both GPU and CPU, while `multiplex_1` only
+    supports CPU.
+*   This op uses the
+    [Eigen](https://eigen.tuxfamily.org/index.php?title=Main_Page#Overview)
+    library to access tensor values and compute the multiplex operation. The
+    `multiplex_1` op only uses Eigen to access tensor values.
+
+This example uses `multiplex_2_kernel.cc` to register the op for CPU, and
+`multiplex_2_kernel.cu.cc` to register the op for GPU. Excluding the
+`multiplex_2_kernel.cu.cc` file from this op will result in a multiplexer similar
+to multiplex_1.
+
+The content on this page assumes familiarity with the high-level process for
+adding custom ops to TensorFlow. For additional context, read the
+[OSS guide on creating custom ops](https://www.tensorflow.org/guide/create_op).
+
+## Creating a custom multiplexer op with GPU support
+
+This example demonstrates how you can create a Python custom multiplexer,
+`multiplex_2_op`, similar to
+[`tf.where`](https://tensorflow.org/api_docs/python/tf/where?version=nightly).
+It returns elements chosen from either of the two input tensors (`x` or `y`),
+depending on the `condition`. You can call the op with the following:
+
+<!-- test_snippets_in_readme skip -->
+```python
+multiplex_2_op.multiplex(condition, x, y)
+```
+
+This simplified `multiplex_2` op has the following limitations that are not
+present in `tf.where`:
+
+*   Support only for dense tensors
+*   No broadcasting capabilities
+*   No extensibility through optional parameters
+
+This example contains C++ and Python code snippets to illustrate the code flow.
+These snippets may be missing namespace declarations, imports, and test cases.
+
+### Step 1 - Define the op interface
+
+Define the op interface and register it using the `REGISTER_OP` macro.
+
+```
+REGISTER_OP("Examples>MultiplexDense")
+    .Input("cond: bool")
+    .Input("a: T")
+    .Input("b: T")
+    .Output("output_values: T")
+    .Attr("T: type")
+    .SetShapeFn([](tensorflow::shape_inference::InferenceContext* c) {
+      // Determine the output shape and also assert that inputs 0 and 1 have
+      // the same shape.
+      tensorflow::shape_inference::ShapeHandle out;
+      TF_RETURN_IF_ERROR(c->Merge(c->input(0), c->input(1), &out));
+      // Assert that inputs 0 and 2 have the same shape, i.e. that all inputs
+      // have the same shape. This is optional, but it is desirable
+      // to raise errors about inconsistent input shapes early when using
+      // graph mode.
+      tensorflow::shape_inference::ShapeHandle unused;
+      TF_RETURN_IF_ERROR(c->Merge(c->input(0), c->input(2), &unused));
+
+      c->set_output(0, out);
+      return ::tensorflow::OkStatus();
+    })
+    .Doc(R"doc(
+Return elements chosen from `a` or `b` depending on `cond`.
+
+This is similar to `np.where` and `tf.where`, but simplified to only handle
+the case of dense tensors, no optional parameters, no broadcasting, etc..
+This uses cond.select from the Eigen library and supports GPU (and CPU).
+
+cond: tf.Tensor of type bool.
+a: tf.Tensor with the same type and shape as `b`.
+b: tf.Tensor with the same type and shape as `a`.
+
+      Where True, yield `a`, otherwise yield `b`.
+output_values: A tf.Tensor with elements from `a` where `cond` is True, and
+               elements from `b` elsewhere.
+)doc");
+```
+
+Note that:
+
+*   This op has three input tensors - one boolean tensor for selecting which
+    values to choose from the two other input tensors of matching type `T`, and
+    one output tensor of type `T`.
+*   The `Attr` for this op is defined as `.Attr("T: type")` which specifies `T`
+    as an `Attr` of type `type`. In the subsequent steps, you will use `T` with
+    a template class to define the type of the contents of tensors.
+*   The docstring for this op is specified by passing a string to `.Doc()`.
+*   The shape function for this op uses the `Merge` method of the
+    [`tensorflow::shape_inference::InferenceContext`](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/framework/shape_inference.h#:~:text=class%20InferenceContext)
+    object which is a helper function to set the output shape to be the same as
+    the identical shapes of the two inputs (for example, if it is used for
+    binary ops) and has error checking to ensure that the two inputs have the
+    same shape. Since `multiplex_2` has three inputs, two calls to `Merge` are
+    used to assert that all three inputs are the same shape.
+    
+### Step 2 - Register the op implementation (kernel)
+
+This example registers the kernel for both CPU and GPU. You can register the
+kernel for only CPU using `multiplex_2_kernel.cc`.
+This will result in a kernel similar to the `multiplex_1` custom op.
+The types supported by GPU kernels are a subset of the types supported by CPU
+kernels.
+
+Register the kernel by calling the `REGISTER_KERNEL_BUILDER` macro.
+
+```
+#define REGISTER_KERNELS_GPU(type)                              \
+  REGISTER_KERNEL_BUILDER(Name("Examples>MultiplexDense")       \
+                              .Device(::tensorflow::DEVICE_GPU) \
+                              .TypeConstraint<type>("T"),       \
+                          MultiplexDenseOp<GPUDevice, type>)
+
+REGISTER_KERNELS_GPU(bool);
+REGISTER_KERNELS_GPU(Eigen::half);
+REGISTER_KERNELS_GPU(float);
+REGISTER_KERNELS_GPU(double);
+REGISTER_KERNELS_GPU(int64);
+REGISTER_KERNELS_GPU(complex64);
+REGISTER_KERNELS_GPU(complex128);
+
+#undef REGISTER_KERNELS_GPU
+```
+
+### Step 3 - Implement the op kernel(s)
+
+In the op kernel (`multiplex_2_kernel.h`), create a class derived from
+[`OpKernel`](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/framework/op_kernel.h#:~:text=class%20OpKernel)
+that implements a `Compute` method to get and validate input tensors, perform
+computation, and create the output tensors. This file is included by both
+`multiplex_2_kernel.cu.cc` (for GPU) and `multiplex_2_kernel.cc` (for CPU).
+
+```
+template <typename Device, typename T>
+class MultiplexDenseOp : public OpKernel {
+ public:
+  explicit MultiplexDenseOp(OpKernelConstruction* ctx) : OpKernel(ctx) {}
+  MultiplexDenseOp(const MultiplexDenseOp& other) = delete;
+  MultiplexDenseOp& operator=(const MultiplexDenseOp& other) = delete;
+  ~MultiplexDenseOp() override = default;
+
+  void Compute(OpKernelContext* ctx) override {
+    const auto& cond_tensor = ctx->input(0);
+    const auto& a_values_tensor = ctx->input(1);
+    const auto& b_values_tensor = ctx->input(2);
+
+    // Allow any shape, but require that a_values, b_values, and cond all
+    // have the same shape.
+    // Note that ::tensorflow::TensorShapeUtils has some useful functions
+    // for checking shapes.
+    OP_REQUIRES(ctx, a_values_tensor.shape() == b_values_tensor.shape(),
+                ::tensorflow::errors::InvalidArgument(
+                    "a and b must have the same shape. "
+                    "a shape: ",
+                    a_values_tensor.shape().DebugString(),
+                    " b shape: ", b_values_tensor.shape().DebugString()));
+    OP_REQUIRES(ctx, a_values_tensor.shape() == cond_tensor.shape(),
+                ::tensorflow::errors::InvalidArgument(
+                    "a and cond must have the same shape. "
+                    "a shape: ",
+                    a_values_tensor.shape().DebugString(),
+                    " cond shape: ", cond_tensor.shape().DebugString()));
+    OP_REQUIRES(ctx, a_values_tensor.NumElements() > 0,
+                ::tensorflow::errors::InvalidArgument(
+                    "Inputs must have at least one element."));
+
+    const auto a_values = a_values_tensor.flat<T>();
+    const auto b_values = b_values_tensor.flat<T>();
+    const auto cond = cond_tensor.flat<bool>();
+
+    // Create an output tensor
+    Tensor* output_tensor = nullptr;
+    OP_REQUIRES_OK(
+        ctx, ctx->allocate_output(0, a_values_tensor.shape(), &output_tensor));
+    auto output = output_tensor->template flat<T>();
+    // Here is an example of processing tensors using the Eigen library.
+    // This supports both CPU and GPU.
+    // For CPU, it supports chunking into blocks and multi-threading.
+    // See
+    // https://eigen.tuxfamily.org/dox/unsupported/eigen_tensors.html#title55
+    output.device(ctx->eigen_device<Device>()) =
+        cond.select(a_values, b_values);
+  }
+};
+```
+
+For intensive mathematical operations, it is a good practice to use
+[Eigen](https://eigen.tuxfamily.org/index.php?title=Main_Page#Overview) to
+perform the computation. Eigen is vectorized, avoids dynamic memory allocation
+and is faster on tensors.The definitions related to Eigen are:
+
+<!-- test_snippets_in_readme skip -->
+```c++
+#define EIGEN_USE_THREADS
+
+#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
+#define EIGEN_USE_GPU
+#endif
+```
+
+[Selection](https://eigen.tuxfamily.org/dox/unsupported/eigen_tensors.html#title55)
+from Eigen supports CPU and GPU devices, as well as chunking data into blocks
+and multi-threading. The `multiplex_2` op contains the following:
+
+<!-- test_snippets_in_readme skip -->
+```c++
+output.device(ctx->eigen_device<Device>()) =
+     cond.select(a_values, b_values);
+```
+
+Using Eigen simplified this example. Alternatively, Custom Ops may implement
+kernels for GPU directly in the `*.cu.cc` files using C++.
+
+#### Compile the op
+
+Compile the C++ op to create a kernel library and Python wrapper that enables
+you to use the op with TensorFlow.
+
+Create a `BUILD` file for the op which declares the dependencies and the output
+build targets. Refer to
+[building for OSS](https://www.tensorflow.org/guide/create_op#build_the_op_library).
+
+### Step 4 - Create the Python wrapper
+
+To create the Python wrapper, import and implement a function that serves as the
+op's public API and provides a docstring.
+
+```
+def multiplex(cond, a, b, name=None):
+  """Return elements chosen from `a` or `b` depending on `cond`.
+
+  This is similar to `np.where` and `tf.where`, but simplified to only handle
+  the case of dense tensors, no optional parameters, no broadcasting, etc..
+
+  >>> multiplex([True, False, False, True], [1,2,3,4], [100,200,300,400])
+  <tf.Tensor: shape=(4,), dtype=int32, numpy=array([  1, 200, 300,   4], ...)>
+
+  Args:
+    cond: tf.Tensor of type bool. Where True, yield `a`, otherwise yield `b`.
+    a: tf.Tensor with the same type and shape as `b`.
+    b: tf.Tensor with the same type and shape as `a`.
+    name: An optional name for the op.
+
+  Returns:
+    A tf.Tensor with elements from `a` where `cond` is True, and elements
+    from `b` elsewhere.
+  """
+  return gen_multiplex_2_op.examples_multiplex_dense(
+      cond=cond, a=a, b=b, name=name)
+```
+
+### Step 5 - Test the op
+
+Create op tests using classes derived from
+[`tf.test.TestCase`](https://www.tensorflow.org/api_docs/python/tf/test/TestCase).
+
+When writing tests to ensure that the op works correctly in both graph and eager
+executions, it is important to note that errors in the op code may be detected
+in two distinct phases of code execution depending on how it is executed (eager
+or graph executions). Errors may be detected early by the shape function or a
+bit later from the logic in the `Compute` method. This may lead to differing
+error types and/or messages.
+
+Below are test excerpts showing how to handle errors for different scenarios.
+The first test case demonstrates error handling when errors are common across
+eager and graph executions and the second test case demonstrates error handling
+when the errors are different in eager and graph executions.
+
+```
+  @test_util.run_in_graph_and_eager_modes
+  def test_multiplex_int(self):
+    a = tf.constant([1, 2, 3, 4, 5], dtype=tf.int64)
+    b = tf.constant([10, 20, 30, 40, 50], dtype=tf.int64)
+    cond = tf.constant([True, False, True, False, True], dtype=bool)
+    expect = np.where(self.evaluate(cond), self.evaluate(a), self.evaluate(b))
+    # expected result is [1, 20, 3, 40, 5]
+    result = multiplex_2_op.multiplex(cond, a, b)
+    self.assertAllEqual(result, expect)
+```
+
+```
+  @test_util.run_in_graph_and_eager_modes
+  def test_multiplex_bad_types(self):
+    a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0])  # float
+    b = tf.constant([10, 20, 30, 40, 50], dtype=tf.int64)
+    cond = tf.constant([True, False, True, False, True], dtype=bool)
+    with self.assertRaisesRegex(
+        (errors_impl.InvalidArgumentError, TypeError),
+        # Eager mode raises InvalidArgumentError with the following message
+        r'(cannot compute Examples>MultiplexDense as input #2\(zero-based\) '
+        r'was expected to be a float tensor but is a int64 tensor '
+        r'\[Op:Examples>MultiplexDense\]'
+        r')|('
+        # Graph mode raises TypeError with the following message
+        r"Input 'b' of 'Examples>MultiplexDense' Op has type int64 that "
+        r"does not match type float32 of argument 'a'.)"):
+      self.evaluate(multiplex_2_op.multiplex(cond, a, b))
+```
+
+Refer to `multiplex_2_test.py` for the full source code which contains all the
+test cases.
+
+Reuse the `BUILD` file to add build rules for the Python API wrapper and the op
+test.
+
+```
+py_strict_library(
+    name = "multiplex_2_op",
+    srcs = ["multiplex_2_op.py"],
+    data = ["multiplex_2_kernel.so"],
+    srcs_version = "PY3",
+    visibility = ["//third_party/tensorflow/examples/custom_ops_doc:__subpackages__"],
+    deps = [
+        "//third_party/py/tensorflow",
+    ],
+)
+
+cuda_py_test(
+    name = "multiplex_2_test",
+    size = "medium",
+    srcs = ["multiplex_2_test.py"],
+    python_version = "PY3",
+    srcs_version = "PY3",
+    tags = [
+        "no_mac",  # TODO(b/216321151): Re-enable this test.
+    ],
+    deps = [
+        ":multiplex_2_op",
+        "//third_party/py/numpy",
+        "//third_party/py/tensorflow",
+        "//third_party/tensorflow/python/framework:errors",
+        "//third_party/tensorflow/python/framework:test_lib",
+    ],
+)
+```
+
+Test the op in the following ways:
+
+*   Build for CPU and test on CPU
+
+    <!-- test_snippets_in_readme skip -->
+    ```shell
+    bazel test //third_party/tensorflow/google/g3doc/example/multiplex_2:multiplex_2_test
+    ```
+
+*   Build for GPU and CPU; test on CPU
+
+    <!-- test_snippets_in_readme skip -->
+    ```shell
+    $ bazel test --config=cuda //third_party/tensorflow/google/g3doc/example/multiplex_2:multiplex_2_test
+    ```
+
+*   Build for GPU and CPU; test on GPU (note the `_gpu` suffix in the target)
+
+    <!-- test_snippets_in_readme skip -->
+    ```shell
+    $ bazel test --config=cuda //third_party/tensorflow/google/g3doc/example/multiplex_2:multiplex_2_test_gpu
+    ```
+
+Testing and building exclusively on CPU only requires the multiplex_2_kernel.cc
+file when registering the op. For all other cases, include both
+multiplex_2_kernel.cc and multiplex_2_kernel.cu.cc files.
+
+### Use the op
+
+Import the op and call it using the following example:
+
+<!-- test_snippets_in_readme skip -->
+```python
+import tensorflow as tf
+
+from tensorflow.examples.custom_ops_doc.multiplex_2 import multiplex_2_op
+
+a = tf.constant([1, 2, 3, 4, 5], dtype=tf.int64)
+b = tf.constant([10, 20, 30, 40, 50], dtype=tf.int64)
+cond = tf.constant([True, False, True, False, True], dtype=bool)
+# expected result is [1, 20, 3, 40, 5]
+result = multiplex_2_op.multiplex(cond, a, b)
+```
+
+Here, `multiplex_2_op` is the name of the Python wrapper that was created in
+this example.
+
+When running an op on GPU, use inputs with types supported by the GPU kernels
+(e.g. this example uses `tf.int64` for `a` and `b` since this type was
+registered).
+
+### Summary
+
+In this example, you learned how to define and use a custom multiplexer op for
+GPU. The image below summarizes the files created for this op.
+
+The table below summarizes the build rules and targets for building and testing
+the `multiplex_2` op.
+
+Op components                           | Build rule             | Build target         | Source
+--------------------------------------- | ---------------------- | -------------------- | ------
+Kernels (C++)                           | `tf_custom_op_library` | `multiplex_2_kernel` | `multiplex_2_kernel.cu.cc`, `multiplex_2_kernel.cc`, `multiplex_2_op.cc`, `multiplex_2_kernel.h`
+Wrapper (automatically generated)       | N/A                    | `gen_multiplex_2_op` | N/A
+Wrapper (with public API and docstring) | `py_strict_library`    | `multiplex_2_op`     | `multiplex_2_op.py`
+Tests                                   | `cuda_py_test`         | `multiplex_2_test`   | `multiplex_2_test.py`
+<!-- LINT.ThenChange(multiplex_2.md) -->

tensorflow/examples/custom_ops_doc/multiplex_2/multiplex_2_kernel.cc

@@ -0,0 +1,41 @@
+/* Copyright 2021 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/examples/custom_ops_doc/multiplex_2/multiplex_2_kernel.h"
+
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/register_types.h"
+#include "tensorflow/core/framework/types.h"
+
+// Please use the appropriate namespace for your project
+namespace tensorflow {
+namespace custom_op_examples {
+
+using CPUDevice = Eigen::ThreadPoolDevice;
+
+// To support tensors containing different types (e.g. int32, float), one
+// kernel per type is registered and is templatized by the "T" attr value.
+// See go/tf-custom-ops-guide
+#define REGISTER_KERNELS_CPU(type)                              \
+  REGISTER_KERNEL_BUILDER(Name("Examples>MultiplexDense")       \
+                              .Device(::tensorflow::DEVICE_CPU) \
+                              .TypeConstraint<type>("T"),       \
+                          MultiplexDenseOp<CPUDevice, type>)
+TF_CALL_ALL_TYPES(REGISTER_KERNELS_CPU);
+
+#undef REGISTER_KERNELS_CPU
+
+}  // namespace custom_op_examples
+}  // namespace tensorflow

tensorflow/examples/custom_ops_doc/multiplex_2/multiplex_2_kernel.cu.cc

@@ -0,0 +1,48 @@
+/* Copyright 2021 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.
+==============================================================================*/
+
+#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
+
+#include "tensorflow/examples/custom_ops_doc/multiplex_2/multiplex_2_kernel.h"
+
+// Please use the appropriate namespace for your project
+namespace tensorflow {
+namespace custom_op_examples {
+
+using GPUDevice = Eigen::GpuDevice;
+
+// To support tensors containing different types (e.g. int32, float), one
+// kernel per type is registered and is templatized by the "T" attr value.
+// See go/tf-custom-ops-guide
+#define REGISTER_KERNELS_GPU(type)                              \
+  REGISTER_KERNEL_BUILDER(Name("Examples>MultiplexDense")       \
+                              .Device(::tensorflow::DEVICE_GPU) \
+                              .TypeConstraint<type>("T"),       \
+                          MultiplexDenseOp<GPUDevice, type>)
+
+REGISTER_KERNELS_GPU(bool);
+REGISTER_KERNELS_GPU(Eigen::half);
+REGISTER_KERNELS_GPU(float);
+REGISTER_KERNELS_GPU(double);
+REGISTER_KERNELS_GPU(int64);
+REGISTER_KERNELS_GPU(complex64);
+REGISTER_KERNELS_GPU(complex128);
+
+#undef REGISTER_KERNELS_GPU
+
+}  // namespace custom_op_examples
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA || TENSORFLOW_USE_ROCM

tensorflow/examples/custom_ops_doc/multiplex_2/multiplex_2_kernel.h

@@ -0,0 +1,99 @@
+/* Copyright 2021 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_EXAMPLES_CUSTOM_OPS_DOC_MULTIPLEX_2_MULTIPLEX_2_KERNEL_H_
+#define TENSORFLOW_EXAMPLES_CUSTOM_OPS_DOC_MULTIPLEX_2_MULTIPLEX_2_KERNEL_H_
+
+#define EIGEN_USE_THREADS
+
+#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
+#define EIGEN_USE_GPU
+#endif
+
+#include <algorithm>
+#include <cstdint>
+#include <limits>
+#include <utility>
+#include <vector>
+
+#include "tensorflow/core/framework/op.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/register_types.h"
+#include "tensorflow/core/framework/shape_inference.h"
+#include "tensorflow/core/platform/errors.h"
+
+// Please use the appropriate namespace for your project
+namespace tensorflow {
+namespace custom_op_examples {
+
+// Multiple devices (i.e. CPU and GPU) and multiple types for the values inside
+// two of the input tensors (e.g. int32, float) are supported by using a
+// template where the device is DEVICE and the type is T.
+template <typename Device, typename T>
+class MultiplexDenseOp : public OpKernel {
+ public:
+  explicit MultiplexDenseOp(OpKernelConstruction* ctx) : OpKernel(ctx) {}
+  MultiplexDenseOp(const MultiplexDenseOp& other) = delete;
+  MultiplexDenseOp& operator=(const MultiplexDenseOp& other) = delete;
+  ~MultiplexDenseOp() override = default;
+
+  void Compute(OpKernelContext* ctx) override {
+    const auto& cond_tensor = ctx->input(0);
+    const auto& a_values_tensor = ctx->input(1);
+    const auto& b_values_tensor = ctx->input(2);
+
+    // Allow any shape, but require that a_values, b_values, and cond all
+    // have the same shape.
+    // Note that ::tensorflow::TensorShapeUtils has some useful functions
+    // for checking shapes.
+    OP_REQUIRES(ctx, a_values_tensor.shape() == b_values_tensor.shape(),
+                ::tensorflow::errors::InvalidArgument(
+                    "a and b must have the same shape. "
+                    "a shape: ",
+                    a_values_tensor.shape().DebugString(),
+                    " b shape: ", b_values_tensor.shape().DebugString()));
+    OP_REQUIRES(ctx, a_values_tensor.shape() == cond_tensor.shape(),
+                ::tensorflow::errors::InvalidArgument(
+                    "a and cond must have the same shape. "
+                    "a shape: ",
+                    a_values_tensor.shape().DebugString(),
+                    " cond shape: ", cond_tensor.shape().DebugString()));
+    OP_REQUIRES(ctx, a_values_tensor.NumElements() > 0,
+                ::tensorflow::errors::InvalidArgument(
+                    "Inputs must have at least one element."));
+
+    const auto a_values = a_values_tensor.flat<T>();
+    const auto b_values = b_values_tensor.flat<T>();
+    const auto cond = cond_tensor.flat<bool>();
+
+    // Create an output tensor
+    Tensor* output_tensor = nullptr;
+    OP_REQUIRES_OK(
+        ctx, ctx->allocate_output(0, a_values_tensor.shape(), &output_tensor));
+    auto output = output_tensor->template flat<T>();
+    // Here is an example of processing tensors using the Eigen library.
+    // This supports both CPU and GPU.
+    // For CPU, it supports chunking into blocks and multi-threading.
+    // See
+    // https://eigen.tuxfamily.org/dox/unsupported/eigen_tensors.html#title55
+    output.device(ctx->eigen_device<Device>()) =
+        cond.select(a_values, b_values);
+  }
+};
+
+}  // namespace custom_op_examples
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_EXAMPLES_CUSTOM_OPS_DOC_MULTIPLEX_2_MULTIPLEX_2_KERNEL_H_

tensorflow/examples/custom_ops_doc/multiplex_2/multiplex_2_op.cc

@@ -0,0 +1,62 @@
+/* Copyright 2021 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 "xla/tsl/platform/errors.h"
+#include "tensorflow/core/framework/op.h"
+#include "tensorflow/core/framework/shape_inference.h"
+#include "tensorflow/core/platform/status.h"
+
+// Use a namespace when registering by prepending the
+// package's name to the op’s name and separate with a '>'.
+// This is the recommendation for out-of-tree ops to avoid name collisions in
+// "Best practices for custom operations in TensorFlow"
+// https://github.com/tensorflow/community/blob/master/rfcs/20190726-custom-ops.md
+
+REGISTER_OP("Examples>MultiplexDense")
+    .Input("cond: bool")
+    .Input("a: T")
+    .Input("b: T")
+    .Output("output_values: T")
+    .Attr("T: type")
+    .SetShapeFn([](tensorflow::shape_inference::InferenceContext* c) {
+      // Determine the output shape and also assert that inputs 0 and 1 have
+      // the same shape.
+      tensorflow::shape_inference::ShapeHandle out;
+      TF_RETURN_IF_ERROR(c->Merge(c->input(0), c->input(1), &out));
+      // Assert that inputs 0 and 2 have the same shape, i.e. that all inputs
+      // have the same shape. This is optional, but it is desirable
+      // to raise errors about inconsistent input shapes early when using
+      // graph mode.
+      tensorflow::shape_inference::ShapeHandle unused;
+      TF_RETURN_IF_ERROR(c->Merge(c->input(0), c->input(2), &unused));
+
+      c->set_output(0, out);
+      return ::tensorflow::OkStatus();
+    })
+    .Doc(R"doc(
+Return elements chosen from `a` or `b` depending on `cond`.
+
+This is similar to `np.where` and `tf.where`, but simplified to only handle
+the case of dense tensors, no optional parameters, no broadcasting, etc..
+This uses cond.select from the Eigen library and supports GPU (and CPU).
+
+cond: tf.Tensor of type bool.
+a: tf.Tensor with the same type and shape as `b`.
+b: tf.Tensor with the same type and shape as `a`.
+
+      Where True, yield `a`, otherwise yield `b`.
+output_values: A tf.Tensor with elements from `a` where `cond` is True, and
+               elements from `b` elsewhere.
+)doc");

tensorflow/examples/custom_ops_doc/multiplex_2/multiplex_2_op.py

@@ -0,0 +1,50 @@
+# Copyright 2021 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.
+# ==============================================================================
+"""A wrapper for gen_multiplex_2_op.py.
+
+This defines a public API and provides a docstring for the C++ Op defined by
+multiplex_2_kernel.cc
+"""
+
+import tensorflow as tf
+from tensorflow.python.platform import resource_loader
+
+_multiplex_2_module = tf.load_op_library(
+    resource_loader.get_path_to_datafile("multiplex_2_kernel.so"))
+
+examples_multiplex_dense = _multiplex_2_module.examples_multiplex_dense
+
+
+def multiplex(cond, a, b, name=None):
+  """Return elements chosen from `a` or `b` depending on `cond`.
+
+  This is similar to `np.where` and `tf.where`, but simplified to only handle
+  the case of dense tensors, no optional parameters, no broadcasting, etc..
+
+  >>> multiplex([True, False, False, True], [1,2,3,4], [100,200,300,400])
+  <tf.Tensor: shape=(4,), dtype=int32, numpy=array([  1, 200, 300,   4], ...)>
+
+  Args:
+    cond: tf.Tensor of type bool. Where True, yield `a`, otherwise yield `b`.
+    a: tf.Tensor with the same type and shape as `b`.
+    b: tf.Tensor with the same type and shape as `a`.
+    name: An optional name for the op.
+
+  Returns:
+    A tf.Tensor with elements from `a` where `cond` is True, and elements
+    from `b` elsewhere.
+  """
+  return examples_multiplex_dense(
+      cond=cond, a=a, b=b, name=name)

tensorflow/examples/custom_ops_doc/multiplex_2/multiplex_2_test.py

@@ -0,0 +1,112 @@
+# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for multiplex_2."""
+
+import numpy as np
+import tensorflow as tf
+
+from tensorflow.examples.custom_ops_doc.multiplex_2 import multiplex_2_op
+from tensorflow.python.framework import errors_impl
+# This pylint disable is only needed for internal google users
+from tensorflow.python.framework import test_util  # pylint: disable=g-direct-tensorflow-import
+
+
+@test_util.with_eager_op_as_function
+class MultiplexOpRank1Test(tf.test.TestCase):
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_multiplex_int(self):
+    a = tf.constant([1, 2, 3, 4, 5], dtype=tf.int64)
+    b = tf.constant([10, 20, 30, 40, 50], dtype=tf.int64)
+    cond = tf.constant([True, False, True, False, True], dtype=bool)
+    expect = np.where(self.evaluate(cond), self.evaluate(a), self.evaluate(b))
+    # expected result is [1, 20, 3, 40, 5]
+    result = multiplex_2_op.multiplex(cond, a, b)
+    self.assertAllEqual(result, expect)
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_multiplex_float(self):
+    a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0])
+    b = tf.constant([10.0, 20.0, 30.0, 40.0, 50.0])
+    cond = tf.constant([True, False, True, False, True], dtype=bool)
+    # expected result is [1.0, 20.0, 3.0, 40.0, 5.0]
+    expect = np.where(self.evaluate(cond), self.evaluate(a), self.evaluate(b))
+    result = multiplex_2_op.multiplex(cond, a, b)
+    self.assertAllEqual(result, expect)
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_multiplex_bad_types(self):
+    a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0])  # float
+    b = tf.constant([10, 20, 30, 40, 50], dtype=tf.int64)
+    cond = tf.constant([True, False, True, False, True], dtype=bool)
+    with self.assertRaisesRegex(
+        (errors_impl.InvalidArgumentError, TypeError),
+        # Eager mode raises InvalidArgumentError with the following message
+        r'(cannot compute Examples>MultiplexDense as input #2\(zero-based\) '
+        r'was expected to be a float tensor but is a int64 tensor '
+        r'\[Op:Examples>MultiplexDense\]'
+        r')|('
+        # Graph mode raises TypeError with the following message
+        r"Input 'b' of 'Examples>MultiplexDense' Op has type int64 that "
+        r"does not match type float32 of argument 'a'.)"):
+      self.evaluate(multiplex_2_op.multiplex(cond, a, b))
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_multiplex_bad_size(self):
+    a = tf.constant([1, 2, 3, 4, 5], dtype=tf.int64)  # longer than b
+    b = tf.constant([10, 20], dtype=tf.int64)  # shorter than a
+    cond = tf.constant([True, False, True, False, True], dtype=bool)
+    with self.assertRaisesRegex(
+        (errors_impl.InvalidArgumentError, ValueError),
+        # Eager mode raises InvalidArgumentError with the following message
+        r'(?s)(a and b must have the same shape. '
+        r'a shape: \[5\] b shape: \[2\].* '
+        r'\[Op:Examples>MultiplexDense\]'
+        r')|('
+        # Graph mode raises ValueError with the following message
+        r'Dimension 0 in both shapes must be equal, but are 5 and 2\. '
+        r'Shapes are \[5\] and \[2\]\.)'):
+      self.evaluate(multiplex_2_op.multiplex(cond, a, b))
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_multiplex_2d(self):
+    a = tf.constant([[1, 2, 3], [4, 5, 6]], dtype=tf.int64)
+    b = tf.constant([[10, 20, 30], [40, 50, 60]], dtype=tf.int64)
+    cond = tf.constant([[True, False, True], [False, True, False]], dtype=bool)
+    expect = np.where(self.evaluate(cond), self.evaluate(a), self.evaluate(b))
+    # expected result is [[1, 20], [3, 40]]
+    result = multiplex_2_op.multiplex(cond, a, b)
+    self.assertAllEqual(result, expect)
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_multiplex_bad_shape(self):
+    a = tf.constant([[1, 2, 3], [4, 5, 6]], dtype=tf.int64)  # shape (2,3)
+    b = tf.constant([[10, 20], [30, 40], [50, 60]],
+                    dtype=tf.int64)  # shape (3,2)
+    cond = tf.constant([[True, False, True], [False, True, False]], dtype=bool)
+    with self.assertRaisesRegex(
+        (errors_impl.InvalidArgumentError, ValueError),
+        # Eager mode raises InvalidArgumentError with the following message
+        r'(a and b must have the same shape.'
+        r' a shape: \[2,3\] b shape: \[3,2\]'
+        r')|('
+        # Graph mode raises ValueError with the following message
+        r'Dimension 0 in both shapes must be equal, '
+        r'but are 2 and 3\. Shapes are \[2,3\] and \[3,2\])\.'):
+      self.evaluate(multiplex_2_op.multiplex(cond, a, b))
+
+
+if __name__ == '__main__':
+  tf.test.main()

tensorflow/examples/custom_ops_doc/multiplex_3/BUILD

@@ -0,0 +1,48 @@
+# Build multiplex_3 custom ops example, which is similar to np.where.
+# This example shows how a Python wrapper can choose either to use "dispach
+# for custom object types" to choose an old C++ Op (that supports only dense
+# tensors) for backwards compatibility or a new C++ for new functionality
+# (that supports sparse tensors).
+
+load("//tensorflow:strict.default.bzl", "py_strict_library")
+load("//tensorflow:tensorflow.bzl", "tf_custom_op_library")
+load("//tensorflow:tensorflow.default.bzl", "tf_py_test")
+
+licenses(["notice"])
+
+tf_custom_op_library(
+    name = "multiplex_3_kernel.so",
+    srcs = [
+        "multiplex_3_kernel.cc",
+        "multiplex_3_op.cc",
+    ],
+)
+
+py_strict_library(
+    name = "multiplex_3_op",
+    srcs = ["multiplex_3_op.py"],
+    data = [":multiplex_3_kernel.so"],
+    srcs_version = "PY3",
+    deps = [
+        "//tensorflow:tensorflow_py",
+        "//tensorflow/examples/custom_ops_doc/multiplex_2:multiplex_2_op",
+    ],
+)
+
+tf_py_test(
+    name = "multiplex_3_test",
+    size = "small",
+    srcs = ["multiplex_3_test.py"],
+    python_version = "PY3",
+    srcs_version = "PY3",
+    tags = [
+        "no_mac",  # TODO(b/216321151): Re-enable this test.
+    ],
+    deps = [
+        ":multiplex_3_op",
+        "//tensorflow:tensorflow_py",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:test_lib",
+        "//third_party/py/numpy",
+    ],
+)

tensorflow/examples/custom_ops_doc/multiplex_3/README.md

@@ -0,0 +1,686 @@
+<!-- LINT.IfChange -->
+# Create a custom multiplexer op with dispatch to special case kernels
+
+This guide provides an end-to-end example of handling special cases with a new
+C++ kernel. The custom op includes a Python wrapper that uses
+[dispatch decorators](https://www.tensorflow.org/guide/extension_type#tensor_api_dispatch)
+to override the default behavior of TensorFlow operations when applied to
+tensor-like types. For more information, refer to
+[extension types](https://www.tensorflow.org/guide/extension_type).
+
+Special case kernels can add new functionality to an existing op without any
+required changes to existing kernels that have already been registered. For
+example, a special case kernel can enable an existing op to handle a different
+type of input.
+
+Optional Python wrappers can enable a variety of non-breaking future changes,
+though it is important to avoid any non-TensorFlow Python code in the
+implementation. This is because any non-Tensorflow Python code will only be used
+in eager execution and not in
+[`tf.function`](https://www.tensorflow.org/api_docs/python/tf/function)
+execution.
+
+Python wrappers can serve the following purposes:
+
+*   **Handling special cases**: The default C++ kernel handles normal cases,
+    while another C++ kernel handles special cases. For example,
+    [`tf.add`](https://www.tensorflow.org/api_docs/python/tf/math/add) calls one
+    of two different C++ kernels, depending on whether the inputs are strings or
+    not.
+
+*   **Decomposing operations into multiple kernels**: Some operations at the
+    Python level are decomposed into multiple C++ kernels, rather than
+    implemented through a single kernel. For example, there is no
+    `ReduceVariance` kernel for the
+    [`tf.reduce_variance`](https://www.tensorflow.org/api_docs/python/tf/math/reduce_variance)
+    op. Instead, the Python `reduce_variance` function computes the variance
+    based on squared deviations from the mean.
+
+*   **Adding an argument to an existing op**: Since `name` is always the last
+    argument wrapper, adding a new, optional argument requires a new wrapper.
+    This prevents the op from mistaking the new argument for the `name`
+    argument.
+
+*   **Changing the order of arguments**: Similar to the point above, a wrapper
+    can be used to change the order of arguments.
+
+The content on this page assumes familiarity with the high-level process for
+adding custom ops to TensorFlow. For additional context, read the
+[OSS guide on creating custom ops](https://www.tensorflow.org/guide/create_op).
+
+## Dispatch to special case kernels using sparse tensors
+
+This example demonstrates how you can create a Python custom multiplexer,
+`multiplex_3_op`, to register a new kernel. If an existing op already handles
+certain kinds of inputs, a special case kernel can extend the op to handle a
+different kind of input without changing the existing op.
+
+The special kernel (`multiplex_3`) in this example extends an existing op
+(`multiplex_2`) so that it can handle
+[sparse tensors](https://www.tensorflow.org/guide/sparse_tensor) as inputs. This
+provides the custom op with the following two kernels:
+
+*   **Default kernel**: registers the multiplex op with dense tensors
+    (`MultiplexDense`).
+*   **Special case kernel**: registers the multiplex op with sparse tensors
+    (`MultiplexSparse`).
+
+The sparse tensor object (`tf.SparseTensor`) is appropriate for tensors that
+contain missing values. Storing sparse values in sparse tensors is more
+memory-efficient than storing in a dense tensor.
+
+In this example, the default kernel is the
+`multiplex_2` (multiplex_2_kernel.cc) kernel, and the new kernel
+(multiplex_3_kernel.cc) implements the multiplex with sparse tensors.
+
+Like other multiplex custom ops, `multiplex_3` is similar to
+[`tf.where`](https://tensorflow.org/api_docs/python/tf/where?version=nightly).
+It returns elements chosen from either of the two input tensors (`x` or `y`),
+depending on the `condition`. You can call the op with the following:
+
+<!-- test_snippets_in_readme skip -->
+```python
+multiplex_3_op.multiplex(condition, x, y)
+```
+
+This simplified `multiplex_3` op has the following limitations that are not
+present in `tf.where`:
+
+*   Support only for CPU computations
+*   No broadcasting capabilities
+*   No extensibility through optional parameters
+
+This example contains C++ and Python code snippets to illustrate the code flow.
+These snippets may be missing namespace declarations, imports, and test cases.
+
+### Step 1 - Define the op interface
+
+Define the op interface and register it using the `REGISTER_OP` macro.
+
+```
+REGISTER_OP("Examples>MultiplexSparse")
+    .Input("cond_indices: int64")
+    .Input("cond_values: bool")
+    .Input("cond_shape: int64")
+    .Input("a_indices: int64")
+    .Input("a_values: T")
+    .Input("a_shape: int64")
+    .Input("b_indices: int64")
+    .Input("b_values: T")
+    .Input("b_shape: int64")
+    .Output("output_indices: int64")
+    .Output("output_values: T")
+    .Output("output_shape: int64")
+    .Attr("T: type")
+    .SetShapeFn([](tensorflow::shape_inference::InferenceContext* c) {
+      tensorflow::shape_inference::ShapeHandle unused;
+      TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 2, &unused));  // cond_indices
+      TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 1, &unused));  // cond_values
+      TF_RETURN_IF_ERROR(c->WithRank(c->input(2), 1, &unused));  // cond_shape
+      TF_RETURN_IF_ERROR(c->WithRank(c->input(3), 2, &unused));  // a_indices
+      TF_RETURN_IF_ERROR(c->WithRank(c->input(4), 1, &unused));  // a_values
+      TF_RETURN_IF_ERROR(c->WithRank(c->input(5), 1, &unused));  // a_shape
+      TF_RETURN_IF_ERROR(c->WithRank(c->input(6), 2, &unused));  // b_indices
+      TF_RETURN_IF_ERROR(c->WithRank(c->input(7), 1, &unused));  // b_values
+      TF_RETURN_IF_ERROR(c->WithRank(c->input(8), 1, &unused));  // b_shape
+      const auto num_rows = c->UnknownDim();
+      const auto dense_rank = c->UnknownDim();
+      c->set_output(0, c->Matrix(num_rows, dense_rank));
+      c->set_output(1, c->Vector(num_rows));
+      c->set_output(2, c->Vector(dense_rank));
+      return ::tensorflow::OkStatus();
+    })
+    .Doc(R"doc(
+Return elements chosen from `a` or `b` depending on `cond`.
+
+This is similar to `np.where` and `tf.where`, but simplified to only handle
+the case of sparse tensors that are vectors, no optional parameters,
+no broadcasting, etc.. Elements for `a` are chosen if there is a `true` `cond`
+value at the same position. Elements for `b` are chosen if there is not a `true`
+`cond` value at the same position, i.e., if either there is a `false` `cond`
+value or the `cond` value is not specified.
+
+Indices must be ordered as described by tf.sparse_reorder.
+
+cond_indices: a rank-2 tensor of sparse indices.
+cond_values: a rank-1 tensor of sparse values.
+cond_shape: a rank-1 tensor representing the dense shape.
+a_indices: a rank-2 tensor of sparse indices.
+a_values: a rank-1 tensor of sparse values.
+a_shape: a rank-1 tensor representing the dense shape.
+b_indices: a rank-2 tensor of sparse indices.
+b_values: a rank-1 tensor of sparse values.
+b_shape: a rank-1 tensor representing the dense shape.
+output_indices: a rank-2 tensor of sparse indices.
+output_values: a rank-1 tensor of sparse values.
+output_shape: a rank-1 tensor representing the dense shape.
+)doc");
+```
+
+#### Inputs and outputs
+
+This op contains a total of nine input tensors. This is made up of three sparse
+tensors (`a`, `b`, and `cond`), where each sparse tensor is encoded using the
+coordinate list (COO) format:
+
+*   `values`: 1D tensor with shape `[N]` containing all non-zero values.
+*   `indices`: 2D tensor with shape `[N, rank]`, containing the indices of the
+    non-zero values
+*   `dense_shape`: 1D tensor with shape `[rank]`, specifying the shape of the
+    tensor.
+
+The `cond` tensor accepts a boolean value to select between `a` and `b`, and the
+`a` and `b` tensors accept a value of type `T`. The output tensor also contains
+a value of type `T`.
+
+#### Shape function
+
+Unlike dense tensors, which have a fixed shape, the shape of sparse tensors
+depend on the number of non-missing values in the output. Since this can not be
+determined by the shape of the inputs, the shape function (`SetShapeFn`) uses
+`UnknownDim()`.
+
+<!-- test_snippets_in_readme skip -->
+```c++
+  .SetShapeFn([](tensorflow::shape_inference::InferenceContext* c) {
+      // Error checking omitted, see source file.
+      const auto num_rows = c->UnknownDim();
+      const auto dense_rank = c->UnknownDim();
+      c->set_output(0, c->Matrix(num_rows, dense_rank));
+      c->set_output(1, c->Vector(num_rows));
+      c->set_output(2, c->Vector(dense_rank));
+      return tensorflow::OkStatus();
+    })
+```
+
+#### Attributes and docstrings
+
+The `Attr` for this op is defined as `.Attr("T: type")`, which specifies `T` as
+an `Attr` of type `type`. In the subsequent steps, you will use `T` with a
+template class to define the type of the contents of tensors.
+
+The docstring for this op is specified by passing a string to `.Doc()`.
+
+### Step 2 - Register the op implementation (kernel)
+
+The C++ kernel in `multiplex_3_kernel.cc` implements a multiplex for sparse
+tensors. For simplicity, this example only supports rank 1 sparse tensors
+(sparse vectors).
+
+Register the kernel by calling the `REGISTER_KERNEL_BUILDER` macro.
+
+```
+#define REGISTER_KERNELS_CPU(type)                              \
+  REGISTER_KERNEL_BUILDER(Name("Examples>MultiplexSparse")      \
+                              .Device(::tensorflow::DEVICE_CPU) \
+                              .TypeConstraint<type>("T"),       \
+                          MultiplexSparseOp<type>)
+TF_CALL_ALL_TYPES(REGISTER_KERNELS_CPU);
+
+#undef REGISTER_KERNELS_CPU
+```
+
+### Step 3 - Implement the op kernel(s)
+
+In the `multiplex_3_kernel.cc` op kernel, create a class derived from
+[`OpKernel`](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/framework/op_kernel.h#:~:text=class%20OpKernel)
+that implements a `Compute` method. This method retrieves and validates input
+tensors, performs computation, and creates output tensors.
+
+```
+template <typename T>
+class MultiplexSparseOp : public OpKernel {
+ public:
+  explicit MultiplexSparseOp(OpKernelConstruction* ctx) : OpKernel(ctx) {}
+  MultiplexSparseOp(const MultiplexSparseOp& other) = delete;
+  MultiplexSparseOp& operator=(const MultiplexSparseOp& other) = delete;
+  ~MultiplexSparseOp() override = default;
+
+  void Compute(OpKernelContext* ctx) override {
+    const auto& cond_indices_tensor = ctx->input(0);
+    const auto& cond_values_tensor = ctx->input(1);
+    const auto& cond_shape_tensor = ctx->input(2);
+    const auto& a_indices_tensor = ctx->input(3);
+    const auto& a_values_tensor = ctx->input(4);
+    const auto& a_shape_tensor = ctx->input(5);
+    const auto& b_indices_tensor = ctx->input(6);
+    const auto& b_values_tensor = ctx->input(7);
+    const auto& b_shape_tensor = ctx->input(8);
+    OP_REQUIRES_OK(ctx,
+                   ValidateSparseTensor(cond_indices_tensor, cond_values_tensor,
+                                        cond_shape_tensor, "cond"));
+    OP_REQUIRES_OK(ctx, ValidateSparseTensor(a_indices_tensor, a_values_tensor,
+                                             a_shape_tensor, "a"));
+    OP_REQUIRES_OK(ctx, ValidateSparseTensor(b_indices_tensor, b_values_tensor,
+                                             b_shape_tensor, "b"));
+    OP_REQUIRES(
+        ctx, cond_shape_tensor.shape() == a_shape_tensor.shape(),
+        InvalidArgument("Sparse tensors must be the same shape. cond_shape: ",
+                        cond_shape_tensor.shape().DebugString(),
+                        " vs a_shape: ", a_shape_tensor.shape().DebugString()));
+    OP_REQUIRES(
+        ctx, a_shape_tensor.shape() == b_shape_tensor.shape(),
+        InvalidArgument("Sparse tensors must be the same shape. a_shape: ",
+                        a_shape_tensor.shape().DebugString(),
+                        " vs b_shape: ", b_shape_tensor.shape().DebugString()));
+    const int rank = a_shape_tensor.dim_size(0);
+    OP_REQUIRES(
+        ctx, rank == 1,
+        InvalidArgument("Sorry, multiplex for sparse tensors only "
+                        "supports rank 1 tensors to simplify this example."));
+    const int cond_elements = cond_indices_tensor.dim_size(0);
+    const int a_elements = a_indices_tensor.dim_size(0);
+    const int b_elements = b_indices_tensor.dim_size(0);
+    const auto cond_indices = cond_indices_tensor.matrix<int64_t>();
+    const auto cond_values = cond_values_tensor.flat<bool>();
+    const auto cond_shape = cond_shape_tensor.flat<int64_t>();
+    const auto a_indices = a_indices_tensor.matrix<int64_t>();
+    const auto a_values = a_values_tensor.flat<T>();
+    const auto a_shape = a_shape_tensor.flat<int64_t>();
+    const auto b_indices = b_indices_tensor.matrix<int64_t>();
+    const auto b_values = b_values_tensor.flat<T>();
+    const auto b_shape = b_shape_tensor.flat<int64_t>();
+    int cond_index = 0;
+    int a_index = 0;
+    int b_index = 0;
+    // This vector is a list of source tensors (a = true, b = false) and source
+    // indices.
+    std::vector<std::pair<bool, int>> merged_output;
+    merged_output.reserve(std::min(cond_elements, a_elements) + b_elements);
+    while (a_index < a_elements || b_index < b_elements) {
+      // Determine the whether the current location with values has a value
+      // for `a`, for `b` or for both `a` and `b`.
+      int64_t cur_row;
+      bool is_a_at_cur = false;
+      bool is_b_at_cur = false;
+      if (a_index < a_elements && b_index < b_elements) {
+        const int64_t a_row = a_indices(a_index, 0);
+        const int64_t b_row = b_indices(b_index, 0);
+        cur_row = std::min(a_row, b_row);
+        if (a_row == cur_row) {
+          is_a_at_cur = true;
+        }
+        if (b_row == cur_row) {
+          is_b_at_cur = true;
+        }
+      } else if (a_index < a_elements) {
+        cur_row = a_indices(a_index, 0);
+        is_a_at_cur = true;
+      } else {  // b_index < b_elements
+        cur_row = b_indices(b_index, 0);
+        is_b_at_cur = true;
+      }
+      // Determine if `cond` has a value at the current location
+      bool cond_flag = false;
+      while (cond_index < cond_elements) {
+        const int64_t cond_row = cond_indices(cond_index, 0);
+        if (cond_row > cur_row) {
+          break;
+        }
+        if (cond_row == cur_row) {
+          cond_flag = cond_values(cond_index);
+          break;
+        }
+        ++cond_index;
+      }
+      // Add `a` or `b` to the merged output based on the condition
+      if (is_a_at_cur) {
+        if (cond_flag) {
+          merged_output.emplace_back(true, a_index);
+        }
+        ++a_index;
+      }
+      if (is_b_at_cur) {
+        if (!cond_flag) {
+          merged_output.emplace_back(false, b_index);
+        }
+        ++b_index;
+      }
+    }
+
+    // Allocate output tensors.
+    Tensor* output_indices_tensor;
+    Tensor* output_values_tensor;
+    Tensor* output_dense_shape_tensor;
+    const int num_values = merged_output.size();
+    OP_REQUIRES_OK(ctx, ctx->allocate_output(0, TensorShape({num_values, rank}),
+                                             &output_indices_tensor));
+    OP_REQUIRES_OK(ctx, ctx->allocate_output(1, TensorShape({num_values}),
+                                             &output_values_tensor));
+    OP_REQUIRES_OK(ctx, ctx->allocate_output(2, TensorShape({rank}),
+                                             &output_dense_shape_tensor));
+    auto output_indices = output_indices_tensor->matrix<int64_t>();
+    auto output_values = output_values_tensor->flat<T>();
+    auto output_shape = output_dense_shape_tensor->flat<int64_t>();
+    for (int row = 0; row < num_values; ++row) {
+      const auto& source_flag = merged_output[row].first;
+      const auto& source_row = merged_output[row].second;
+      const auto& indices = source_flag ? a_indices : b_indices;
+      const auto& values = source_flag ? a_values : b_values;
+      for (int column = 0; column < rank; ++column) {
+        output_indices(row, column) = indices(source_row, column);
+      }
+      output_values(row) = values(source_row);
+    }
+    // Expand the shape of the output sparse tensor so that it is as large
+    // as the shape of the largest input in each dimension.
+    // An alternative behavior would be to require that the shapes be the
+    // same and implement error checking that all the corresponding values
+    // in the shape tensors are the same (e.g.
+    // `cond_shape(i) == a_shape(i)` and `a_shape(i) == b_shape(i)` in
+    // OP_REQUIRES above and `output_shape(i) = a_shape(i)` here).
+    for (int i = 0; i < rank; ++i) {
+      output_shape(i) =
+          std::max(cond_shape(i), std::max(a_shape(i), b_shape(i)));
+    }
+  }
+
+ private:
+  Status ValidateSparseTensor(const ::tensorflow::Tensor& indices_tensor,
+                              const ::tensorflow::Tensor& values_tensor,
+                              const ::tensorflow::Tensor& shape_tensor,
+                              const string label) {
+    if (!TensorShapeUtils::IsMatrix(indices_tensor.shape())) {
+      return InvalidArgument(
+          "Sparse indices for ", label,
+          " must be rank 2, not shape: ", indices_tensor.shape().DebugString());
+    }
+    if (!TensorShapeUtils::IsVector(values_tensor.shape())) {
+      return InvalidArgument("Sparse values for ", label,
+                             " must be a vector, not shape: ",
+                             values_tensor.shape().DebugString());
+    }
+    if (!TensorShapeUtils::IsVector(shape_tensor.shape())) {
+      return InvalidArgument(
+          "Sparse shape for ", label,
+          " must be a vector, not shape: ", shape_tensor.shape().DebugString());
+    }
+    if (indices_tensor.dim_size(0) != values_tensor.dim_size(0)) {
+      return InvalidArgument("Sparse indices and values for " + label +
+                                 " must have the same "
+                                 "number of rows. indices: ",
+                             indices_tensor.shape().DebugString(),
+                             " values: ", values_tensor.shape().DebugString());
+    }
+    return OkStatus();
+  }
+};
+```
+
+The following snippet shows how the kernel accesses one of the inputs.
+
+<!-- test_snippets_in_readme skip -->
+```c++
+  void Compute(OpKernelContext* ctx) override {
+    const auto& cond_indices_tensor = ctx->input(0);
+    const auto& cond_values_tensor = ctx->input(1);
+    const auto& cond_shape_tensor = ctx->input(2);
+    // Error checking omitted, see source file.
+    const int cond_elements = cond_indices_tensor.dim_size(0);
+    const auto cond_indices = cond_indices_tensor.matrix<int64_t>();
+    const auto cond_values = cond_values_tensor.flat<bool>();
+    const auto cond_shape = cond_shape_tensor.flat<int64_t>();
+  }
+```
+
+The kernel implements the multiplex operation through the following steps:
+
+1.  **Create an empty list**: The list contains elements that refer to values of
+    `a` or `b` in merged order. The elements are pairs of a `bool` to indicate
+    the source tensor (a = true, b = false) and an `int` to indicate the source
+    index.
+
+2.  **Append values in `a` and `b` to the list**: Looping through all
+    non-missing values in `a` and `b`, determine whether each position has a
+    non-missing value in `a`, `b`, or both.
+
+    If `cond` is true and `a` has a non-missing value at that position, append
+    this element to the list. If `cond` is false or missing and `b` has a
+    non-missing value at that position, append this element to the list.
+
+    Note: Assume that indices are sorted in canonical row-major order (e.g.
+    using
+    [`tf.sparse.reorder`](https://www.tensorflow.org/api_docs/python/tf/sparse/reorder)).
+
+3.  **Allocate the output**: The size of the output is based on the length of
+    the list.
+
+4.  **Add indices and values of `a` and `b` to the output**: Iterate through the
+    elements in the list and copy the indices and values of `a` and `b` to the
+    output.
+
+5.  **Set the shape of the output**: Set the (dense) shape of the output to
+    match the largest of the inputs. Shaping the output is accomplished with
+    following snippet:
+
+    <!-- test_snippets_in_readme skip -->
+    ```c++
+    for (int i = 0; i < rank; ++i) {
+      output_shape(i) =
+          std::max(cond_shape(i), std::max(a_shape(i), b_shape(i)));
+    }
+    ```
+
+##### Considerations when setting the output shape
+
+This implementation is specific to sparse tensors. The expansion is simple
+because the concept of "missing values" for sparse tensors is well-defined.
+There is no exact equivalent for dense tensors.
+
+In many cases, sparse tensors are just sparse encodings of a dense tensor. In
+these cases, all inputs should have the same dense shape, and the output shape
+would be identical to the shape of the inputs.
+
+In these cases, you can replace the `std::max calculation` with
+`output_shape(i) = a_shape(i)` after verifying the following conditions in
+`OP_REQUIRES`:
+
+<!-- test_snippets_in_readme skip -->
+```c++
+cond_shape(i) == a_shape(i)
+a_shape(i) == b_shape(i)
+```
+
+### Step 4 - Create the Python wrapper
+
+To create the Python wrapper, import and implement a function that serves as the
+op's public API and provides a docstring.
+
+When the inputs are
+[`tf.SparseTensor`](https://www.tensorflow.org/api_docs/python/tf/sparse/SparseTensor),
+the
+[dispatch decorator](https://www.tensorflow.org/guide/extension_type#tensor_api_dispatch)
+in the snippet below prompts the previous
+`gen_multiplex_2_op.examples_multiplex_dense` op to use the new C++ kernel
+wrapped by `gen_multiplex_3_op.examples_multiplex_sparse`.
+
+Note: This optional Python wrapper depends on the `multiplex_2` op in addition
+to the `multiplex_3` dependencies. See `deps` for `multiplex_3_op` in the BUILD
+file.
+
+```
+@tf.experimental.dispatch_for_api(gen_multiplex_2_op.examples_multiplex_dense)
+def multiplex_sparse(cond: tf.SparseTensor,
+                     a: tf.SparseTensor,
+                     b: tf.SparseTensor,
+                     name=None):
+  """Return elements chosen from `a` or `b` depending on `cond`.
+
+
+  This is similar to `np.where` and `tf.where`, but simplified to only handle
+  the case of rank 1 sparse tensors, no optional parameters, no broadcasting,
+  etc..
+
+  >>> cond = tf.SparseTensor(
+  ...     indices=[[1], [3], [6]], values=[True, False, True], dense_shape=[7])
+  >>> a = tf.sparse.from_dense(['', 'a0', '', 'a1', '', 'a2', ''])
+  >>> b = tf.sparse.from_dense(['b0', '', 'b1', 'b2', '', '', 'b3'])
+  >>> multiplex_3_op.multiplex_sparse(cond, a, b)
+  SparseTensorValue(indices=array([[0],
+    [1],
+    [2],
+    [3]]), values=array([b'b0', b'a0', b'b1', b'b2'], dtype=object),
+    dense_shape=array([7]))
+  Args:
+    cond: tf.SparseTensor of type bool. Where True, yield `a`, otherwise yield
+      `b`.
+    a: tf.SparseTensor with the same type and shape as `b`.
+    b: tf.SparseTensor with the same type and shape as `a`.
+    name: An optional name for the op.
+
+  Returns:
+    A tf.SparseTensor with elements from `a` where `cond` is True, and elements
+    from `b` elsewhere.
+  """
+  (indices, values, shape) = examples_multiplex_sparse(
+      cond_indices=cond.indices,
+      cond_values=cond.values,
+      cond_shape=cond.dense_shape,
+      a_indices=a.indices,
+      a_values=a.values,
+      a_shape=a.dense_shape,
+      b_indices=b.indices,
+      b_values=b.values,
+      b_shape=b.dense_shape,
+      name=name)
+  return tf.SparseTensor(indices, values, shape)
+```
+
+### Step 5 - Test the op
+
+Create op tests using classes derived from
+[`tf.test.TestCase`](https://www.tensorflow.org/api_docs/python/tf/test/TestCase).
+
+When writing tests to ensure that the op works correctly in both graph and eager
+executions, it is important to note that errors in the op code may be detected
+in two distinct phases of code execution depending on how it is executed (eager
+or graph executions). Errors may be detected early by the shape function or a
+bit later from the logic in the `Compute` method. This may lead to differing
+error types and/or messages.
+
+The following tests use the `multiplex_2_op.multiplex` custom op from the
+previous `multiplex_2` example, which now supports sparse tensors (while
+continuing to support dense tensors). The `test_sparse_op_different` test inputs
+are sparse tensors, so it uses the new `multiplex_3_kernel` C++ kernel.
+
+```
+  @test_util.run_in_graph_and_eager_modes
+  def test_sparse_op_different(self):
+    cond = tf.SparseTensor(
+        indices=[[1], [3], [6]], values=[True, False, True], dense_shape=[7])
+    a = tf.SparseTensor(
+        indices=[[1], [3], [5]], values=['a0', 'a1', 'a2'], dense_shape=[6])
+    b = tf.SparseTensor(
+        indices=[[0], [2], [3], [6]],
+        values=['b0', 'b1', 'b2', 'b3'],
+        dense_shape=[7])
+    result = self.evaluate(multiplex_2_op.multiplex(cond, a, b))
+    self.assertAllEqual([7], result.dense_shape)
+    self.assertAllEqual([[0], [1], [2], [3]], result.indices)
+    self.assertAllEqual([b'b0', b'a0', b'b1', b'b2'], result.values)
+```
+
+The `test_multiplex_int` test inputs are dense tensors, so it uses the old
+`multiplex_2_kernel` C++ kernel.
+
+```
+  @test_util.run_in_graph_and_eager_modes
+  def test_multiplex_int(self):
+    a = tf.constant([1, 2, 3, 4, 5], dtype=tf.int64)
+    b = tf.constant([10, 20, 30, 40, 50], dtype=tf.int64)
+    cond = tf.constant([True, False, True, False, True], dtype=bool)
+    expect = np.where(self.evaluate(cond), self.evaluate(a), self.evaluate(b))
+    # expected result is [1, 20, 3, 40, 5]
+    result = multiplex_2_op.multiplex(cond, a, b)
+    self.assertAllEqual(result, expect)
+```
+
+Refer to `multiplex_3_test.py` for the full source code which contains all the
+test cases.
+
+Reuse the `BUILD` file to add build rules for the Python API wrapper and the op
+test.
+
+```
+tf_custom_op_library(
+    name = "multiplex_3_kernel.so",
+    srcs = [
+        "multiplex_3_kernel.cc",
+        "multiplex_3_op.cc",
+    ],
+)
+
+py_strict_library(
+    name = "multiplex_3_op",
+    srcs = ["multiplex_3_op.py"],
+    data = [":multiplex_3_kernel.so"],
+    srcs_version = "PY3",
+    deps = [
+        "//third_party/py/tensorflow",
+        "//third_party/tensorflow/examples/custom_ops_doc/multiplex_2:multiplex_2_op",
+    ],
+)
+```
+
+Test the op with the following:
+
+<!-- test_snippets_in_readme skip -->
+```shell
+bazel test //third_party/tensorflow/google/g3doc/example/multiplex_3:multiplex_3_test
+```
+
+### Use the op
+
+Import the op and call it using the following example:
+
+<!-- test_snippets_in_readme skip -->
+```python
+import tensorflow as tf
+
+from tensorflow.examples.custom_ops_doc.multiplex_2 import multiplex_2_op
+from tensorflow.examples.custom_ops_doc.multiplex_3 import multiplex_3_op
+
+cond = tf.SparseTensor(indices=[[1], [3], [6]],
+                    values=[True, False, True], dense_shape=[7])
+a = tf.SparseTensor(indices=[[1], [3], [5]],
+                    values=['a0', 'a1', 'a2'], dense_shape=[6])
+b = tf.SparseTensor(indices=[[0], [2], [3], [6]],
+                    values=['b0', 'b1', 'b2', 'b3'], dense_shape=[7])
+result = multiplex_2_op.multiplex(cond, a, b)
+```
+
+Here, `multiplex_2_op` is the name of the Python wrapper that was created in the
+multiplex_2 example. Importing the `multiplex_3_op` Python wrapper created in
+this example extends `multiplex_2_op.multiplex` to handle sparse tensors.
+
+Build the op with the following:
+
+<!-- test_snippets_in_readme skip -->
+```shell
+bazel build //third_party/tensorflow/examples/custom_ops_doc/multiplex_3:multiplex_3_op
+```
+
+### Summary
+
+In this example, you learned how implement a new multiplexer kernel to handle
+special cases. With a Python wrapper that uses
+[dispatch decorators](https://www.tensorflow.org/guide/extension_type#tensor_api_dispatch)
+to override the default kernel, this custom op uses a new kernel to handle
+sparse tensors.
+
+The table below summarizes the build rules and targets for building and testing
+the `multiplex_3` op.
+
+Op components                           | Build rule             | Build target         | Source
+--------------------------------------- | ---------------------- | -------------------- | ------
+Kernels (C++)                           | `tf_custom_op_library` | `multiplex_3_kernel` | `multiplex_3_kernel.cc`, `multiplex_3_op.cc`
+Wrapper (automatically generated)       | N/A                    | `gen_multiplex_3_op` | N/A
+Wrapper (with public API and docstring) | `py_strict_library`    | `multiplex_3_op`     | `multiplex_3_op.py`
+Tests                                   | `tf_py_test`           | `multiplex_3_test`   | `multiplex_3_test.py`
+
+## Resources
+
+*   [OSS custom ops guide](https://www.tensorflow.org/guide/create_op)
+*   [Extension types and dispatch decorators](https://www.tensorflow.org/guide/extension_type#tensor_api_dispatch)
+*   [Working with sparse tensors](https://www.tensorflow.org/guide/sparse_tensor)
+<!-- LINT.ThenChange(multiplex_3.md) -->

tensorflow/examples/custom_ops_doc/multiplex_3/multiplex_3_kernel.cc

@@ -0,0 +1,228 @@
+/* Copyright 2021 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 <cstdint>
+#include <utility>
+#include <vector>
+
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/op_requires.h"
+#include "tensorflow/core/framework/register_types.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/platform/types.h"
+
+// Please use the appropriate namespace for your project
+namespace tensorflow {
+namespace custom_op_examples {
+
+using CPUDevice = Eigen::ThreadPoolDevice;
+using ::tensorflow::errors::InvalidArgument;
+
+template <typename T>
+class MultiplexSparseOp : public OpKernel {
+ public:
+  explicit MultiplexSparseOp(OpKernelConstruction* ctx) : OpKernel(ctx) {}
+  MultiplexSparseOp(const MultiplexSparseOp& other) = delete;
+  MultiplexSparseOp& operator=(const MultiplexSparseOp& other) = delete;
+  ~MultiplexSparseOp() override = default;
+
+  void Compute(OpKernelContext* ctx) override {
+    const auto& cond_indices_tensor = ctx->input(0);
+    const auto& cond_values_tensor = ctx->input(1);
+    const auto& cond_shape_tensor = ctx->input(2);
+    const auto& a_indices_tensor = ctx->input(3);
+    const auto& a_values_tensor = ctx->input(4);
+    const auto& a_shape_tensor = ctx->input(5);
+    const auto& b_indices_tensor = ctx->input(6);
+    const auto& b_values_tensor = ctx->input(7);
+    const auto& b_shape_tensor = ctx->input(8);
+    OP_REQUIRES_OK(ctx,
+                   ValidateSparseTensor(cond_indices_tensor, cond_values_tensor,
+                                        cond_shape_tensor, "cond"));
+    OP_REQUIRES_OK(ctx, ValidateSparseTensor(a_indices_tensor, a_values_tensor,
+                                             a_shape_tensor, "a"));
+    OP_REQUIRES_OK(ctx, ValidateSparseTensor(b_indices_tensor, b_values_tensor,
+                                             b_shape_tensor, "b"));
+    OP_REQUIRES(
+        ctx, cond_shape_tensor.shape() == a_shape_tensor.shape(),
+        InvalidArgument("Sparse tensors must be the same shape. cond_shape: ",
+                        cond_shape_tensor.shape().DebugString(),
+                        " vs a_shape: ", a_shape_tensor.shape().DebugString()));
+    OP_REQUIRES(
+        ctx, a_shape_tensor.shape() == b_shape_tensor.shape(),
+        InvalidArgument("Sparse tensors must be the same shape. a_shape: ",
+                        a_shape_tensor.shape().DebugString(),
+                        " vs b_shape: ", b_shape_tensor.shape().DebugString()));
+    const int rank = a_shape_tensor.dim_size(0);
+    OP_REQUIRES(
+        ctx, rank == 1,
+        InvalidArgument("Sorry, multiplex for sparse tensors only "
+                        "supports rank 1 tensors to simplify this example."));
+    const int cond_elements = cond_indices_tensor.dim_size(0);
+    const int a_elements = a_indices_tensor.dim_size(0);
+    const int b_elements = b_indices_tensor.dim_size(0);
+    const auto cond_indices = cond_indices_tensor.matrix<int64_t>();
+    const auto cond_values = cond_values_tensor.flat<bool>();
+    const auto cond_shape = cond_shape_tensor.flat<int64_t>();
+    const auto a_indices = a_indices_tensor.matrix<int64_t>();
+    const auto a_values = a_values_tensor.flat<T>();
+    const auto a_shape = a_shape_tensor.flat<int64_t>();
+    const auto b_indices = b_indices_tensor.matrix<int64_t>();
+    const auto b_values = b_values_tensor.flat<T>();
+    const auto b_shape = b_shape_tensor.flat<int64_t>();
+    int cond_index = 0;
+    int a_index = 0;
+    int b_index = 0;
+    // This vector is a list of source tensors (a = true, b = false) and source
+    // indices.
+    std::vector<std::pair<bool, int>> merged_output;
+    merged_output.reserve(std::min(cond_elements, a_elements) + b_elements);
+    while (a_index < a_elements || b_index < b_elements) {
+      // Determine the whether the current location with values has a value
+      // for `a`, for `b` or for both `a` and `b`.
+      int64_t cur_row;
+      bool is_a_at_cur = false;
+      bool is_b_at_cur = false;
+      if (a_index < a_elements && b_index < b_elements) {
+        const int64_t a_row = a_indices(a_index, 0);
+        const int64_t b_row = b_indices(b_index, 0);
+        cur_row = std::min(a_row, b_row);
+        if (a_row == cur_row) {
+          is_a_at_cur = true;
+        }
+        if (b_row == cur_row) {
+          is_b_at_cur = true;
+        }
+      } else if (a_index < a_elements) {
+        cur_row = a_indices(a_index, 0);
+        is_a_at_cur = true;
+      } else {  // b_index < b_elements
+        cur_row = b_indices(b_index, 0);
+        is_b_at_cur = true;
+      }
+      // Determine if `cond` has a value at the current location
+      bool cond_flag = false;
+      while (cond_index < cond_elements) {
+        const int64_t cond_row = cond_indices(cond_index, 0);
+        if (cond_row > cur_row) {
+          break;
+        }
+        if (cond_row == cur_row) {
+          cond_flag = cond_values(cond_index);
+          break;
+        }
+        ++cond_index;
+      }
+      // Add `a` or `b` to the merged output based on the condition
+      if (is_a_at_cur) {
+        if (cond_flag) {
+          merged_output.emplace_back(true, a_index);
+        }
+        ++a_index;
+      }
+      if (is_b_at_cur) {
+        if (!cond_flag) {
+          merged_output.emplace_back(false, b_index);
+        }
+        ++b_index;
+      }
+    }
+
+    // Allocate output tensors.
+    Tensor* output_indices_tensor;
+    Tensor* output_values_tensor;
+    Tensor* output_dense_shape_tensor;
+    const int num_values = merged_output.size();
+    OP_REQUIRES_OK(ctx, ctx->allocate_output(0, TensorShape({num_values, rank}),
+                                             &output_indices_tensor));
+    OP_REQUIRES_OK(ctx, ctx->allocate_output(1, TensorShape({num_values}),
+                                             &output_values_tensor));
+    OP_REQUIRES_OK(ctx, ctx->allocate_output(2, TensorShape({rank}),
+                                             &output_dense_shape_tensor));
+    auto output_indices = output_indices_tensor->matrix<int64_t>();
+    auto output_values = output_values_tensor->flat<T>();
+    auto output_shape = output_dense_shape_tensor->flat<int64_t>();
+    for (int row = 0; row < num_values; ++row) {
+      const auto& source_flag = merged_output[row].first;
+      const auto& source_row = merged_output[row].second;
+      const auto& indices = source_flag ? a_indices : b_indices;
+      const auto& values = source_flag ? a_values : b_values;
+      for (int column = 0; column < rank; ++column) {
+        output_indices(row, column) = indices(source_row, column);
+      }
+      output_values(row) = values(source_row);
+    }
+    // Expand the shape of the output sparse tensor so that it is as large
+    // as the shape of the largest input in each dimension.
+    // An alternative behavior would be to require that the shapes be the
+    // same and implement error checking that all the corresponding values
+    // in the shape tensors are the same (e.g.
+    // `cond_shape(i) == a_shape(i)` and `a_shape(i) == b_shape(i)` in
+    // OP_REQUIRES above and `output_shape(i) = a_shape(i)` here).
+    for (int i = 0; i < rank; ++i) {
+      output_shape(i) =
+          std::max(cond_shape(i), std::max(a_shape(i), b_shape(i)));
+    }
+  }
+
+ private:
+  Status ValidateSparseTensor(const ::tensorflow::Tensor& indices_tensor,
+                              const ::tensorflow::Tensor& values_tensor,
+                              const ::tensorflow::Tensor& shape_tensor,
+                              const string label) {
+    if (!TensorShapeUtils::IsMatrix(indices_tensor.shape())) {
+      return InvalidArgument(
+          "Sparse indices for ", label,
+          " must be rank 2, not shape: ", indices_tensor.shape().DebugString());
+    }
+    if (!TensorShapeUtils::IsVector(values_tensor.shape())) {
+      return InvalidArgument("Sparse values for ", label,
+                             " must be a vector, not shape: ",
+                             values_tensor.shape().DebugString());
+    }
+    if (!TensorShapeUtils::IsVector(shape_tensor.shape())) {
+      return InvalidArgument(
+          "Sparse shape for ", label,
+          " must be a vector, not shape: ", shape_tensor.shape().DebugString());
+    }
+    if (indices_tensor.dim_size(0) != values_tensor.dim_size(0)) {
+      return InvalidArgument("Sparse indices and values for " + label +
+                                 " must have the same "
+                                 "number of rows. indices: ",
+                             indices_tensor.shape().DebugString(),
+                             " values: ", values_tensor.shape().DebugString());
+    }
+    return OkStatus();
+  }
+};
+
+// To support tensors containing different types (e.g. int32, float), one
+// kernel per type is registered and is templatized by the "T" attr value.
+// See go/tf-custom-ops-guide
+#define REGISTER_KERNELS_CPU(type)                              \
+  REGISTER_KERNEL_BUILDER(Name("Examples>MultiplexSparse")      \
+                              .Device(::tensorflow::DEVICE_CPU) \
+                              .TypeConstraint<type>("T"),       \
+                          MultiplexSparseOp<type>)
+TF_CALL_ALL_TYPES(REGISTER_KERNELS_CPU);
+
+#undef REGISTER_KERNELS_CPU
+
+}  // namespace custom_op_examples
+}  // namespace tensorflow

tensorflow/examples/custom_ops_doc/multiplex_3/multiplex_3_op.cc

@@ -0,0 +1,83 @@
+/* Copyright 2021 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 "xla/tsl/platform/errors.h"
+#include "tensorflow/core/framework/op.h"
+#include "tensorflow/core/framework/shape_inference.h"
+#include "tensorflow/core/platform/status.h"
+
+// Use a namespace when registering by prepending the
+// package's name to the op’s name and separate with a '>'.
+// This is the recommendation for out-of-tree ops to avoid name collisions in
+// "Best practices for custom operations in TensorFlow"
+// https://github.com/tensorflow/community/blob/master/rfcs/20190726-custom-ops.md
+
+REGISTER_OP("Examples>MultiplexSparse")
+    .Input("cond_indices: int64")
+    .Input("cond_values: bool")
+    .Input("cond_shape: int64")
+    .Input("a_indices: int64")
+    .Input("a_values: T")
+    .Input("a_shape: int64")
+    .Input("b_indices: int64")
+    .Input("b_values: T")
+    .Input("b_shape: int64")
+    .Output("output_indices: int64")
+    .Output("output_values: T")
+    .Output("output_shape: int64")
+    .Attr("T: type")
+    .SetShapeFn([](tensorflow::shape_inference::InferenceContext* c) {
+      tensorflow::shape_inference::ShapeHandle unused;
+      TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 2, &unused));  // cond_indices
+      TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 1, &unused));  // cond_values
+      TF_RETURN_IF_ERROR(c->WithRank(c->input(2), 1, &unused));  // cond_shape
+      TF_RETURN_IF_ERROR(c->WithRank(c->input(3), 2, &unused));  // a_indices
+      TF_RETURN_IF_ERROR(c->WithRank(c->input(4), 1, &unused));  // a_values
+      TF_RETURN_IF_ERROR(c->WithRank(c->input(5), 1, &unused));  // a_shape
+      TF_RETURN_IF_ERROR(c->WithRank(c->input(6), 2, &unused));  // b_indices
+      TF_RETURN_IF_ERROR(c->WithRank(c->input(7), 1, &unused));  // b_values
+      TF_RETURN_IF_ERROR(c->WithRank(c->input(8), 1, &unused));  // b_shape
+      const auto num_rows = c->UnknownDim();
+      const auto dense_rank = c->UnknownDim();
+      c->set_output(0, c->Matrix(num_rows, dense_rank));
+      c->set_output(1, c->Vector(num_rows));
+      c->set_output(2, c->Vector(dense_rank));
+      return ::tensorflow::OkStatus();
+    })
+    .Doc(R"doc(
+Return elements chosen from `a` or `b` depending on `cond`.
+
+This is similar to `np.where` and `tf.where`, but simplified to only handle
+the case of sparse tensors that are vectors, no optional parameters,
+no broadcasting, etc.. Elements for `a` are chosen if there is a `true` `cond`
+value at the same position. Elements for `b` are chosen if there is not a `true`
+`cond` value at the same position, i.e., if either there is a `false` `cond`
+value or the `cond` value is not specified.
+
+Indices must be ordered as described by tf.sparse_reorder.
+
+cond_indices: a rank-2 tensor of sparse indices.
+cond_values: a rank-1 tensor of sparse values.
+cond_shape: a rank-1 tensor representing the dense shape.
+a_indices: a rank-2 tensor of sparse indices.
+a_values: a rank-1 tensor of sparse values.
+a_shape: a rank-1 tensor representing the dense shape.
+b_indices: a rank-2 tensor of sparse indices.
+b_values: a rank-1 tensor of sparse values.
+b_shape: a rank-1 tensor representing the dense shape.
+output_indices: a rank-2 tensor of sparse indices.
+output_values: a rank-1 tensor of sparse values.
+output_shape: a rank-1 tensor representing the dense shape.
+)doc");

tensorflow/examples/custom_ops_doc/multiplex_3/multiplex_3_op.py

@@ -0,0 +1,84 @@
+# Copyright 2021 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.
+# ==============================================================================
+"""A wrapper for gen_multiplex_3_op.py.
+
+This defines a public API and provides a docstring for the C++ Op defined by
+multiplex_3_kernel.cc
+"""
+
+import tensorflow as tf
+from tensorflow.examples.custom_ops_doc.multiplex_2.multiplex_2_op import examples_multiplex_dense
+from tensorflow.python.platform import resource_loader
+
+
+_multiplex_3_module = tf.load_op_library(
+    resource_loader.get_path_to_datafile("multiplex_3_kernel.so"))
+
+examples_multiplex_sparse = _multiplex_3_module.examples_multiplex_sparse
+
+
+# This Python function uses a decorator for "tensor API dispatch" so that it
+# specializes the "gen_multiplex_2_op.examples_multiplex_dense" function that
+# previously supported only dense tensors (using the C++ kernel from the
+# previous multiplex_2 example) to now support sparse tensors (using the new
+# C++ kernel from this example). Since multiplex_2_op.multiplex uses this,
+# it now supports sparse tensors. Behavior for dense tensors is unchanged so
+# this is a non-breaking change. See
+# https://www.tensorflow.org/guide/extension_type#tensor_api_dispatch
+@tf.experimental.dispatch_for_api(examples_multiplex_dense)
+def multiplex_sparse(cond: tf.SparseTensor,
+                     a: tf.SparseTensor,
+                     b: tf.SparseTensor,
+                     name=None):
+  """Return elements chosen from `a` or `b` depending on `cond`.
+
+
+  This is similar to `np.where` and `tf.where`, but simplified to only handle
+  the case of rank 1 sparse tensors, no optional parameters, no broadcasting,
+  etc..
+
+  >>> cond = tf.SparseTensor(
+  ...     indices=[[1], [3], [6]], values=[True, False, True], dense_shape=[7])
+  >>> a = tf.sparse.from_dense(['', 'a0', '', 'a1', '', 'a2', ''])
+  >>> b = tf.sparse.from_dense(['b0', '', 'b1', 'b2', '', '', 'b3'])
+  >>> multiplex_3_op.multiplex_sparse(cond, a, b)
+  SparseTensorValue(indices=array([[0],
+    [1],
+    [2],
+    [3]]), values=array([b'b0', b'a0', b'b1', b'b2'], dtype=object),
+    dense_shape=array([7]))
+  Args:
+    cond: tf.SparseTensor of type bool. Where True, yield `a`, otherwise yield
+      `b`.
+    a: tf.SparseTensor with the same type and shape as `b`.
+    b: tf.SparseTensor with the same type and shape as `a`.
+    name: An optional name for the op.
+
+  Returns:
+    A tf.SparseTensor with elements from `a` where `cond` is True, and elements
+    from `b` elsewhere.
+  """
+  (indices, values, shape) = examples_multiplex_sparse(
+      cond_indices=cond.indices,
+      cond_values=cond.values,
+      cond_shape=cond.dense_shape,
+      a_indices=a.indices,
+      a_values=a.values,
+      a_shape=a.dense_shape,
+      b_indices=b.indices,
+      b_values=b.values,
+      b_shape=b.dense_shape,
+      name=name)
+  return tf.SparseTensor(indices, values, shape)

tensorflow/examples/custom_ops_doc/multiplex_3/multiplex_3_test.py

@@ -0,0 +1,180 @@
+# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for multiplex_3."""
+
+import numpy as np
+import tensorflow as tf
+
+from tensorflow.examples.custom_ops_doc.multiplex_2 import multiplex_2_op
+from tensorflow.examples.custom_ops_doc.multiplex_3 import multiplex_3_op
+# This pylint disable is only needed for internal google users
+from tensorflow.python.framework import test_util  # pylint: disable=g-direct-tensorflow-import
+
+
+@test_util.with_eager_op_as_function
+class MultiplexOpRank1Test(tf.test.TestCase):
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_sparse_kernel(self):
+    idx0 = tf.constant([], dtype=tf.int64, shape=[0, 1])
+    val0 = tf.constant([], dtype=tf.int64)
+    val5a = tf.constant([1, 2, 3, 4, 5], dtype=tf.int64)
+    idx5b = tf.constant([[10], [20], [30], [40], [50]], dtype=tf.int64)
+    val5b = tf.constant([10, 20, 30, 40, 50], dtype=tf.int64)
+    cond0 = tf.constant([], dtype=bool)
+    cond5 = tf.constant([True, False, True, False, True], dtype=bool)
+    val3a = tf.constant([1, 2, 3], dtype=tf.int64)
+    val3b = tf.constant([4, 5, 6], dtype=tf.int64)
+    idx3c = tf.constant([[10], [20], [30]], dtype=tf.int64)
+    idx3d = tf.constant([[30], [40], [50]], dtype=tf.int64)
+    idx3e = tf.constant([[10], [30], [50]], dtype=tf.int64)
+    cond3 = tf.constant([True, False, True], dtype=bool)
+    shape = tf.constant([100], dtype=tf.int64)
+
+    # all inputs empty
+    (result_index, result_values,
+     result_shape) = multiplex_3_op.examples_multiplex_sparse(
+         idx0, cond0, shape, idx0, val0, shape, idx0, val0, shape)
+    self.assertAllEqual(idx0, result_index)
+    self.assertAllEqual(val0, result_values)
+    self.assertAllEqual(shape, result_shape)
+
+    # only b is not empty
+    (result_index, result_values,
+     result_shape) = multiplex_3_op.examples_multiplex_sparse(
+         idx0, cond0, shape, idx0, val0, shape, idx5b, val5a, shape)
+    self.assertAllEqual(idx5b, result_index)
+    self.assertAllEqual(val5a, result_values)
+    self.assertAllEqual(shape, result_shape)
+
+    # all indices the same
+    (result_index, result_values,
+     result_shape) = multiplex_3_op.examples_multiplex_sparse(
+         idx5b, cond5, shape, idx5b, val5a, shape, idx5b, val5b, shape)
+    expect_values = tf.constant([1, 20, 3, 40, 5], dtype=tf.int64)
+    self.assertAllEqual(idx5b, result_index)
+    self.assertAllEqual(expect_values, result_values)
+    self.assertAllEqual(shape, result_shape)
+
+    # cond and a have same positions, b values after a values
+    (result_index, result_values,
+     result_shape) = multiplex_3_op.examples_multiplex_sparse(
+         idx3c, cond3, shape, idx3c, val3a, shape, idx3d, val3b, shape)
+    expect_index = tf.constant([[10], [30], [40], [50]], dtype=tf.int64)
+    expect_values = tf.constant([1, 3, 5, 6], dtype=tf.int64)
+    self.assertAllEqual(expect_index, result_index)
+    self.assertAllEqual(expect_values, result_values)
+    self.assertAllEqual(shape, result_shape)
+
+    # cond and a have same positions, b values before a values
+    (result_index, result_values,
+     result_shape) = multiplex_3_op.examples_multiplex_sparse(
+         idx3d, cond3, shape, idx3d, val3a, shape, idx3c, val3b, shape)
+    expect_index = tf.constant([[10], [20], [30], [50]], dtype=tf.int64)
+    expect_values = tf.constant([4, 5, 1, 3], dtype=tf.int64)
+    self.assertAllEqual(expect_index, result_index)
+    self.assertAllEqual(expect_values, result_values)
+    self.assertAllEqual(shape, result_shape)
+
+    # cond and b have same positions, b values after a values
+    (result_index, result_values,
+     result_shape) = multiplex_3_op.examples_multiplex_sparse(
+         idx3d, cond3, shape, idx3c, val3a, shape, idx3d, val3b, shape)
+    expect_index = tf.constant([[30], [40]], dtype=tf.int64)
+    expect_values = tf.constant([3, 5], dtype=tf.int64)
+    self.assertAllEqual(expect_index, result_index)
+    self.assertAllEqual(expect_values, result_values)
+    self.assertAllEqual(shape, result_shape)
+
+    # cond and b have same positions, b values before a values
+    (result_index, result_values,
+     result_shape) = multiplex_3_op.examples_multiplex_sparse(
+         idx3c, cond3, shape, idx3d, val3a, shape, idx3c, val3b, shape)
+    expect_index = tf.constant([[20], [30]], dtype=tf.int64)
+    expect_values = tf.constant([5, 1], dtype=tf.int64)
+    self.assertAllEqual(expect_index, result_index)
+    self.assertAllEqual(expect_values, result_values)
+    self.assertAllEqual(shape, result_shape)
+
+    # cond and a and b all have different positions
+    (result_index, result_values,
+     result_shape) = multiplex_3_op.examples_multiplex_sparse(
+         idx3e, cond3, shape, idx3c, val3a, shape, idx3d, val3b, shape)
+    expect_index = tf.constant([[10], [30], [40]], dtype=tf.int64)
+    expect_values = tf.constant([1, 4, 5], dtype=tf.int64)
+    self.assertAllEqual(expect_index, result_index)
+    self.assertAllEqual(expect_values, result_values)
+    self.assertAllEqual(shape, result_shape)
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_sparse_op_only(self):
+    cond = tf.SparseTensor(
+        indices=[[1], [3], [6]], values=[True, False, True], dense_shape=[7])
+    a = tf.SparseTensor(
+        indices=[[1], [3], [5]], values=['a0', 'a1', 'a2'], dense_shape=[6])
+    b = tf.SparseTensor(
+        indices=[[0], [2], [3], [6]],
+        values=['b0', 'b1', 'b2', 'b3'],
+        dense_shape=[7])
+    result = self.evaluate(multiplex_3_op.multiplex_sparse(cond, a, b))
+    self.assertAllEqual([7], result.dense_shape)
+    self.assertAllEqual([[0], [1], [2], [3]], result.indices)
+    self.assertAllEqual([b'b0', b'a0', b'b1', b'b2'], result.values)
+
+  # The following tests use multiplex_2_op.multiplex, which now supports
+  # sparse tensors.
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_sparse_op_same(self):
+    cond = tf.SparseTensor(
+        indices=[[1], [3], [6]], values=[True, False, True], dense_shape=[7])
+    a = tf.SparseTensor(
+        indices=[[1], [3], [6]], values=['a0', 'a1', 'a2'], dense_shape=[6])
+    b = tf.SparseTensor(
+        indices=[[1], [3], [6]], values=['b0', 'b1', 'b2'], dense_shape=[7])
+    result = self.evaluate(multiplex_2_op.multiplex(cond, a, b))
+    self.assertAllEqual([7], result.dense_shape)
+    self.assertAllEqual([[1], [3], [6]], result.indices)
+    self.assertAllEqual([b'a0', b'b1', b'a2'], result.values)
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_sparse_op_different(self):
+    cond = tf.SparseTensor(
+        indices=[[1], [3], [6]], values=[True, False, True], dense_shape=[7])
+    a = tf.SparseTensor(
+        indices=[[1], [3], [5]], values=['a0', 'a1', 'a2'], dense_shape=[6])
+    b = tf.SparseTensor(
+        indices=[[0], [2], [3], [6]],
+        values=['b0', 'b1', 'b2', 'b3'],
+        dense_shape=[7])
+    result = self.evaluate(multiplex_2_op.multiplex(cond, a, b))
+    self.assertAllEqual([7], result.dense_shape)
+    self.assertAllEqual([[0], [1], [2], [3]], result.indices)
+    self.assertAllEqual([b'b0', b'a0', b'b1', b'b2'], result.values)
+
+  # muliplex still works on dense tensors
+  @test_util.run_in_graph_and_eager_modes
+  def test_multiplex_int(self):
+    a = tf.constant([1, 2, 3, 4, 5], dtype=tf.int64)
+    b = tf.constant([10, 20, 30, 40, 50], dtype=tf.int64)
+    cond = tf.constant([True, False, True, False, True], dtype=bool)
+    expect = np.where(self.evaluate(cond), self.evaluate(a), self.evaluate(b))
+    # expected result is [1, 20, 3, 40, 5]
+    result = multiplex_2_op.multiplex(cond, a, b)
+    self.assertAllEqual(result, expect)
+
+
+if __name__ == '__main__':
+  tf.test.main()

tensorflow/examples/custom_ops_doc/multiplex_4/BUILD

@@ -0,0 +1,79 @@
+# Build multiplex_4 custom ops examples, which is similar to np.where
+
+load("//tensorflow:strict.default.bzl", "py_strict_binary", "py_strict_library")
+load("//tensorflow:tensorflow.bzl", "tf_custom_op_library")
+load("//tensorflow:tensorflow.default.bzl", "tf_py_test")
+
+licenses(["notice"])
+
+tf_custom_op_library(
+    name = "multiplex_4_kernel.so",
+    srcs = [
+        "multiplex_4_kernel.cc",
+        "multiplex_4_op.cc",
+    ],
+    deps = [
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/status",
+    ],
+)
+
+py_strict_library(
+    name = "multiplex_4_op",
+    srcs = ["multiplex_4_op.py"],
+    data = ["multiplex_4_kernel.so"],
+    srcs_version = "PY3",
+    deps = [
+        "//tensorflow:tensorflow_py",
+    ],
+)
+
+py_strict_library(
+    name = "model_using_multiplex",
+    srcs = ["model_using_multiplex.py"],
+    srcs_version = "PY3",
+    deps = [
+        "//tensorflow:tensorflow_py",
+    ],
+)
+
+tf_py_test(
+    name = "multiplex_4_test",
+    size = "medium",  # This test blocks because it writes and reads a file,
+    srcs = ["multiplex_4_test.py"],
+    python_version = "PY3",
+    srcs_version = "PY3",
+    tags = [
+        "no_mac",  # TODO(b/216321151): Re-enable this test.
+    ],
+    deps = [
+        ":model_using_multiplex",
+        ":multiplex_4_op",
+        "//tensorflow:tensorflow_py",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:test_lib",
+        "//third_party/py/numpy",
+    ],
+)
+
+py_strict_binary(
+    name = "multiplex_2_save",
+    srcs = ["multiplex_2_save.py"],
+    srcs_version = "PY3",
+    deps = [
+        "model_using_multiplex",
+        "//tensorflow/examples/custom_ops_doc/multiplex_2:multiplex_2_op",
+        "@absl_py//absl:app",
+    ],
+)
+
+py_strict_binary(
+    name = "multiplex_4_load_use",
+    srcs = ["multiplex_4_load_use.py"],
+    srcs_version = "PY3",
+    deps = [
+        "model_using_multiplex",
+        "multiplex_4_op",
+        "@absl_py//absl:app",
+    ],
+)

tensorflow/examples/custom_ops_doc/multiplex_4/README.md

@@ -0,0 +1,612 @@
+<!-- LINT.IfChange -->
+# Create a custom multiplexer op with C++ backward compatibility
+
+This guide provides an end-to-end implementation of a new custom op that is
+backwards compatible with an existing custom op.
+
+The example in this guide implements a new custom op that handles inputs that
+are Python lists of tensors, and is backwards compatible with an existing op
+that only handles inputs that are single tensors.
+
+The existing op is a multiplexer that returns elements chosen from either of two
+single input tensors `x` or `y` depending on a single `condition` tensor.
+
+The content on this page assumes familiarity with the high-level process for
+adding custom ops to TensorFlow. For additional context, read the
+[OSS guide on creating custom ops](https://www.tensorflow.org/guide/create_op).
+
+## A backwards compatible kernel that handles lists of tensors
+
+This example demonstrates how you can create a custom multiplexer,
+`multiplex_4`, to register a new kernel that is backward compatible with an
+existing multiplex_2` op.
+
+The new custom op registers a kernel
+(multiplex_4_kernel.cc) that takes lists of tensors as inputs, and is backwards
+compatible with the existing kernel (multiplex_2_kernel.cc) that takes only
+single tensors as inputs.
+
+The `multiplex_4` op is similar to
+[numpy.select](https://numpy.org/doc/stable/reference/generated/numpy.select.html),
+while the `multiplex_2` op is similar to
+[numpy.where](https://numpy.org/doc/stable/reference/generated/numpy.where.html).
+
+The lists of tensors that the new op takes as inputs are of a particular fixed
+size. Since the list size is defined in `Attr`, it is fixed at graph
+construction time when the constructor of the C++ kernel is called. Therefore,
+the size of the list cannot be data dependent. See
+[Ragged tensors](https://www.tensorflow.org/guide/ragged_tensor) for
+variable length lists.
+
+This example contains C++ and Python code snippets to illustrate the code flow.
+These snippets may be missing namespace declarations, imports, and test cases.
+
+### Prerequisites - Implement `multiplex_2` and `SavedModel`
+
+This example uses a [`SavedModel`](https://www.tensorflow.org/guide/saved_model)
+from an existing `multiplex_2` custom op.
+
+The `muliplex_2_save.py` file uses `save` from `model_using_muliplex.py` to
+create a `SavedModel` named `model_using_multiplex` in the current working
+directory.
+
+```
+def save(multiplex_op, path):
+  """Save a model that contains the given `multiplex_op`.
+
+  Args:
+    multiplex_op: A multiplex Custom Op, e.g. multiplex_4_op.multiplex. This is
+      parameterized so it can also be used to create an "old" model with an
+      older version of the op, e.g. multiplex_2_op.multiplex.
+    path: Directory to save model to.
+  """
+  example_cond, example_a, example_b = _get_example_tensors()
+
+  class UseMultiplex(tf.Module):
+
+    @tf.function(input_signature=[
+        tf.TensorSpec.from_tensor(example_cond),
+        tf.TensorSpec.from_tensor(example_a),
+        tf.TensorSpec.from_tensor(example_b)
+    ])
+    def use_multiplex(self, cond, a, b):
+      return multiplex_op(cond, a, b)
+
+  model = UseMultiplex()
+  tf.saved_model.save(
+      model,
+      path,
+      signatures=model.use_multiplex.get_concrete_function(
+          tf.TensorSpec.from_tensor(example_cond),
+          tf.TensorSpec.from_tensor(example_a),
+          tf.TensorSpec.from_tensor(example_b)))
+```
+
+This `SavedModel` has the old version of the custom op (`multiplex_2`) that only
+supports individual tensors as inputs. The following steps will register a
+kernel that accepts lists of tensors as inputs, while maintaining backward
+compatibility with the previous op.
+
+### Step 1 - Define the op interface
+
+Define the op interface and register it using the `REGISTER_OP` macro.
+
+```
+REGISTER_OP("Examples>MultiplexDense")
+    .Input("cond: N * bool")
+    .Input("a_values: N * T")
+    .Input("b_values: T")
+    .Output("output_values: T")
+    .Attr("T: type")
+    .Attr("N: int = 1")
+    .SetShapeFn(MultiplexShapeFunction)
+    .Doc(R"doc(
+Return elements chosen from `a_values` or `b_values` depending on `cond`.
+
+When `a_values` and `cond` are tenors (i.e. N=1), this is similar to `np.where`
+and `tf.where`. When `a_values` and `cond` are lists of tensors (i.e. N>1),
+this is similar to `np.select`. In either case these are simplified to only
+handle dense tensors, no optional parameters, no broadcasting, etc..
+
+cond: tf.Tensor or list of tf.Tensor of type bool. If it is a list, `a_values`
+      must be a list of the same length. Where True, yield the corresponding
+      element from `a_values` (with priority to the first one encountered in
+      lists), otherwise yield `b_values`.
+a_values: tf.Tensor or list of tf.Tensor. Each tensor has the same type and
+          shape as `b_values`. If it is a list, `cond` must be a list of the
+          same length.
+b_values: tf.Tensor with the same type and shape as the `a_values` if it is a
+          tensor or as every element of `a_values` if `a_values` is a list.
+output_values: A tf.Tensor with elements from `a_values` where `cond` is True,
+               and elements from `b` elsewhere.
+)doc");
+```
+
+While the `multiplex_2` op defined inputs as single tensors, such as `cond:
+bool` and `a_values: T`, this op supports lists of tensors by adding `N*`, where
+`N` is the length of the lists.
+
+The default list size (`N`) is set to 1 with the following: `.Attr("N: int =
+1")`. If the inputs are single tensors, then `N` is equal to 1, which is
+backwards compatible with a previous definition of `.Input("x: T")`.
+
+All lists in this example are of equal length (`N`). To support lists of
+different lengths, define an attribute for each unique length. For example:
+
+<!-- test_snippets_in_readme skip -->
+```c++
+.Input("short_list: short_len * float")
+.Input("long_list: long_len * float")
+.Attr("short_len: int = 1")
+.Attr("long_len: int >= 10")
+```
+
+### Step 2 - Register the op implementation (kernel)
+
+The C++ kernel in `multiplex_4_kernel.cc` implements a multiplexer that accepts
+lists of tensors as inputs. Register the kernel by calling the
+`REGISTER_KERNEL_BUILDER` macro.
+
+```
+#define REGISTER_KERNELS(type)                                  \
+  REGISTER_KERNEL_BUILDER(Name("Examples>MultiplexDense")       \
+                              .Device(::tensorflow::DEVICE_CPU) \
+                              .TypeConstraint<type>("T"),       \
+                          MultiplexDenseOp<type>)
+
+TF_CALL_ALL_TYPES(REGISTER_KERNELS);
+#undef REGISTER_KERNELS
+```
+
+### Step 3 - Implement the op kernel
+
+In the `multiplex_4_kernel.cc` op kernel, create a class derived from
+[`OpKernel`](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/framework/op_kernel.h#:~:text=class%20OpKernel)
+that implements a `Compute` method. This method retrieves and validates input
+tensors, performs computation, and creates output tensors.
+
+```
+template <typename T>
+class MultiplexDenseOp : public OpKernel {
+ public:
+  explicit MultiplexDenseOp(OpKernelConstruction* ctx) : OpKernel(ctx) {
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("N", &num_cond_a_));
+  }
+
+  MultiplexDenseOp(const MultiplexDenseOp& other) = delete;
+  MultiplexDenseOp& operator=(const MultiplexDenseOp& other) = delete;
+  ~MultiplexDenseOp() override = default;
+
+  void Compute(OpKernelContext* ctx) override {
+    // Optional error checking: cond and a_values are lists of N, so there are
+    // a total of 2N+1 inputs. Check that the  number of inputs and the
+    // `N` Attr is consistent.
+    const int64_t expected_inputs = 2 * num_cond_a_ + 1;
+    OP_REQUIRES(ctx, expected_inputs == ctx->num_inputs(),
+                Internal("expected_inputs != num_inputs(): ", expected_inputs,
+                         " != ", ctx->num_inputs()));
+    VLOG(1) << "N " << num_cond_a_;
+
+    const auto& first_cond_tensor = ctx->input(0);
+    const auto& first_a_values_tensor = ctx->input(num_cond_a_);
+    const auto& b_values_tensor = ctx->input(2 * num_cond_a_);
+
+    // Allow any shape, but require that a_values, b_values, and cond all
+    // have the same shape.
+    // Note that ::tensorflow::TensorShapeUtils has some useful functions
+    // for checking shapes.
+    for (int64_t i = 0; i < num_cond_a_; i++) {
+      const auto& cond_tensor_i = ctx->input(i);
+      const auto& a_values_tensor_i = ctx->input(num_cond_a_ + i);
+      OP_REQUIRES(
+          ctx, a_values_tensor_i.shape() == b_values_tensor.shape(),
+          InvalidArgument(
+              "a_values[", i,
+              "] and b_values must have the same shape. "
+              "a_values[",
+              i, "] shape: ", a_values_tensor_i.DebugString(),
+              " b_values shape: ", b_values_tensor.shape().DebugString()));
+      OP_REQUIRES(
+          ctx, cond_tensor_i.shape() == b_values_tensor.shape(),
+          InvalidArgument(
+              "cond_values[", i,
+              "] and b_valuesmust have the same shape. "
+              "cond_values[",
+              i, "] shape: ", first_a_values_tensor.shape().DebugString(),
+              " b_values shape: ", first_cond_tensor.shape().DebugString()));
+    }
+
+    // Create an output tensor
+    Tensor* output_tensor = nullptr;
+    OP_REQUIRES_OK(
+        ctx, ctx->allocate_output(0, b_values_tensor.shape(), &output_tensor));
+    auto output = output_tensor->template flat<T>();
+
+    const auto b_values = b_values_tensor.template flat<T>();
+    // np.select style behavior, `cond` and `a_values` are lists of tensors.
+    // Also works for the np.where style case where there is only one `cond`
+    // and one `a_values` tensor.
+    const int64_t N = first_a_values_tensor.NumElements();
+    for (int64_t i = 0; i < N; i++) {
+      bool flag = false;
+      for (int64_t list_index = 0; list_index < num_cond_a_; list_index++) {
+        const auto& cond_tensor = ctx->input(list_index);
+        const auto& a_values_tensor = ctx->input(num_cond_a_ + list_index);
+        const auto cond = cond_tensor.template flat<bool>();
+        const auto a_values = a_values_tensor.template flat<T>();
+        if (cond(i)) {
+          output(i) = a_values(i);
+          flag = true;
+          VLOG(1) << "A " << list_index << " for " << i;
+          break;
+        }
+      }
+      if (!flag) {
+        output(i) = b_values(i);
+        VLOG(1) << "B for " << i;
+      }
+    }
+  }
+
+ private:
+  int64_t num_cond_a_;  // the number of `cond` and `a` input tensors
+};
+```
+
+The kernel uses a private member variable (`num_cond_a_`) to hold the length of
+`cond` and `a`. The constructor saves the `N` attribute into the variable.
+
+<!-- test_snippets_in_readme skip -->
+```c++
+private:
+  int64_t num_cond_a_;  // the number of cond and a input tensors
+```
+
+<!-- test_snippets_in_readme skip -->
+```c++
+explicit MultiplexDenseOp(OpKernelConstruction* ctx) : OpKernel(ctx) {
+  OP_REQUIRES_OK(ctx, ctx->GetAttr("N", &num_cond_a_));
+}
+```
+
+The `num_cond_a_` variable is used to index the inputs in the following order:
+`cond`, `a`, `b`. The op interfaces specify that `cond` and `a` are tensor lists
+of length `N`, and `b` is a single tensor. The inputs are indexed as follows:
+
+1.  `cond`: [0 ... N-1]
+2.  `a`: [N ... 2*N-1]
+3.  `b`: [2*N]
+
+When `num_cond_a_` is equal to 1, the kernel implements `numpy.where` as it
+would in the `multiplex_2` op. When `num_cond_a_` is greater than 1, the kernel
+implements `numpy.select`. This is achieved with the following `for` loop.
+
+<!-- test_snippets_in_readme skip -->
+```c++
+for (int64_t i = 0; i < N; i++) {
+  bool flag = false;
+  for (int64_t list_index = 0; list_index < num_cond_a_; list_index++) {
+    const auto& cond_tensor = ctx->input(list_index);
+    const auto& a_values_tensor = ctx->input(num_cond_a_ + list_index);
+    const auto cond = cond_tensor.flat<bool>();
+    const auto a_values = a_values_tensor.flat<T>();
+    if (cond(i)) {
+      output(i) = a_values(i);
+      flag = true;
+      break;
+    }
+  }
+  if (!flag) {
+    output(i) = b_values(i);
+  }
+}
+```
+
+#### Compile the op
+
+Compile the C++ op to create a kernel library and Python wrapper that enables
+you to use the op with TensorFlow.
+
+Create a `BUILD` file for the op which declares the dependencies and the output
+build targets. Refer to
+[building for OSS](https://www.tensorflow.org/guide/create_op#build_the_op_library).
+
+### Step 4 - Create the Python wrapper
+
+To create the Python wrapper, import and implement a function that serves as the
+op's public API and provides a docstring.
+
+If `cond` and `a` are not already lists, the wrapper in `multiplex_4_op.py`
+puts the variables in lists before the `numpy.where` implementation.
+
+Note: The generated Python wrapper automatically sets the `N` attribute based on
+the length of the input lists.
+
+```
+def multiplex(cond, a, b, name=None):
+  """Return elements chosen from `a` or `b` depending on `cond`.
+
+  This is similar to `np.where` and `tf.where` if `cond` and `a` are tensors.
+  This is similar to `np.select` if `cond` and `a` are lists of tensors.
+  In either case, this is simplified to only handle the case of dense tensors,
+  no optional parameters, no broadcasting, etc..
+
+  >>> multiplex([True, False, False, True], [1,2,3,4], [100,200,300,400])
+  <tf.Tensor: shape=(4,), dtype=int32, numpy=array([  1, 200, 300,   4], ...)>
+
+  >>> a1 = tf.constant([1, 2, 3, 4, 5], dtype=tf.int64)
+  >>> a2 = tf.constant([6, 7, 8, 9, 10], dtype=tf.int64)
+  >>> a3 = tf.constant([11, 12, 13, 14, 15], dtype=tf.int64)
+  >>> b = tf.constant([101, 102, 103, 104, 105], dtype=tf.int64)
+  >>> cond1 = tf.constant([False, False, True, False, False], dtype=bool)
+  >>> cond2 = tf.constant([False, False, False, False, True], dtype=bool)
+  >>> cond3 = tf.constant([True, False, True, False, True], dtype=bool)
+  >>> multiplex_4_op.multiplex([cond1, cond2, cond3], [a1, a2, a3], b)
+  <tf.Tensor: shape=(5,), ... numpy=array([ 11, 102,   3, 104,  10], ...)>
+
+  Args:
+    cond: tf.Tensor or list of tf.Tensor of type bool. Where True, yield `a`.
+      When multiple corresponding `cond` elements are true, the first one yield
+      based on the first one encountered.
+    a: tf.Tensor or list of tf.Tensor, each with the same type and shape as `b`.
+    b: tf.Tensor or list of tf.Tensor with the same type and shape as `a`. Yield
+      `b` if all corresponding `cond` values is False.
+    name: An optional name for the op.
+
+  Returns:
+    A tf.Tensor with elements from `a` where `cond` is True, and elements
+    from `b` elsewhere.
+  """
+  if not isinstance(cond, (list, tuple)):
+    # Support "old" use of multiplex where `cond` and `a` are tensors,
+    # not lists of tensors.
+    return gen_multiplex_4_op.examples_multiplex_dense(
+        cond=[cond], a_values=[a], b_values=b, name=name)
+  return gen_multiplex_4_op.examples_multiplex_dense(
+      cond=cond, a_values=a, b_values=b, name=name)
+```
+
+### Step 5 - Test the op
+
+Create op tests using classes derived from
+[`tf.test.TestCase`](https://www.tensorflow.org/api_docs/python/tf/test/TestCase).
+
+When writing tests to ensure that the op works correctly in both graph and eager
+executions, it is important to note that errors in the op code may be detected
+in two distinct phases of code execution depending on how it is executed (eager
+or graph executions). Errors may be detected early by the shape function or a
+bit later from the logic in the `Compute` method. This may lead to differing
+error types and messages.
+
+```
+@test_util.with_eager_op_as_function
+class MultiplexOpTest(tf.test.TestCase):
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_multiplex_int(self):
+    a = tf.constant([1, 2, 3, 4, 5], dtype=tf.int64)
+    b = tf.constant([10, 20, 30, 40, 50], dtype=tf.int64)
+    cond = tf.constant([True, False, True, False, True], dtype=bool)
+    expect = np.where(self.evaluate(cond), self.evaluate(a), self.evaluate(b))
+    # expected result is [1, 20, 3, 40, 5]
+    result = multiplex_4_op.multiplex(cond, a, b)
+    self.assertAllEqual(result, expect)
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_multiplex_select(self):
+    a1 = tf.constant([1, 2, 3, 4, 5], dtype=tf.int64)
+    a2 = tf.constant([6, 7, 8, 9, 10], dtype=tf.int64)
+    a3 = tf.constant([11, 12, 13, 14, 15], dtype=tf.int64)
+    a = [a1, a2, a3]
+    b = tf.constant([101, 102, 103, 104, 105], dtype=tf.int64)
+    cond1 = tf.constant([False, False, True, False, False], dtype=bool)
+    cond2 = tf.constant([False, False, False, False, True], dtype=bool)
+    cond3 = tf.constant([True, False, True, False, True], dtype=bool)
+    cond = [cond1, cond2, cond3]
+    expect = np.select([self.evaluate(i) for i in cond],
+                       [self.evaluate(i) for i in a], self.evaluate(b))
+    # expected result is [11, 102, 3, 104, 10]
+    result = multiplex_4_op.multiplex(cond, a, b)
+    self.assertAllEqual(result, expect)
+
+  def test_multiplex_saved_model(self):
+    path = os.path.join(self.create_tempdir(), 'model')
+    model_using_multiplex.save(multiplex_4_op.multiplex, path)
+    result = model_using_multiplex.load_and_use(path)
+    self.assertAllEqual(result, tf.constant([1, 20, 3, 40, 5], dtype=tf.int64))
+
+  # One tf.function that uses both multiplex with single tensors for `cond`
+  # and `a` and with lists of tensors for `cond` and `a`, i.e. a graph
+  # with two example_multiplex_dense kernels that have different numbers
+  # of inputs.
+  @tf.function
+  def _both(self):
+    a1 = tf.constant([1, 2, 3, 4, 5], dtype=tf.int64)
+    a2 = tf.constant([6, 7, 8, 9, 10], dtype=tf.int64)
+    a3 = tf.constant([11, 12, 13, 14, 15], dtype=tf.int64)
+    a_123 = [a1, a2, a3]
+    b_123 = tf.constant([101, 102, 103, 104, 105], dtype=tf.int64)
+    cond1 = tf.constant([False, False, True, False, False], dtype=bool)
+    cond2 = tf.constant([False, False, False, False, True], dtype=bool)
+    cond3 = tf.constant([True, False, True, False, True], dtype=bool)
+    cond_123 = [cond1, cond2, cond3]
+    mux_123 = multiplex_4_op.multiplex(cond_123, a_123, b_123)
+    b4 = tf.constant([201, 202, 203, 204, 205], dtype=tf.int64)
+    cond4 = tf.constant([True, True, True, False, False], dtype=bool)
+    result = multiplex_4_op.multiplex(cond4, mux_123, b4)
+    return result
+
+  def test_both_single_and_list(self):
+    result = self._both()
+    self.assertAllEqual(result,
+                        tf.constant([11, 102, 3, 204, 205], dtype=tf.int64))
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_inconsistent_inputs_error(self):
+    a1 = tf.constant([1, 2, 3, 4, 5], dtype=tf.int64)
+    a2 = tf.constant([6, 7, 8, 9, 10], dtype=tf.int64)
+    a = [a1, a2]
+    b = tf.constant([101, 102, 103, 104, 105], dtype=tf.int64)
+    cond = tf.constant([False, False, True, False, False], dtype=bool)
+    with self.assertRaisesRegex(
+        (errors_impl.InvalidArgumentError, ValueError),
+        # Eager mode raises InvalidArgumentError with the following message
+        r'(a_values\[0\] and b_values must have the same shape'
+        r')|('
+        # Graph mode raises ValueError with the following message
+        r'Shapes must be equal rank, but are 2 and 1)'):
+      self.evaluate(multiplex_4_op.multiplex(cond, a, b))
+```
+
+The following `tf.function` in multiplex_4_test.py has two multiplex custom ops:
+one that takes lists for its `cond` and `a` inputs, and another that takes
+single tensors.
+
+```
+  @tf.function
+  def _both(self):
+    a1 = tf.constant([1, 2, 3, 4, 5], dtype=tf.int64)
+    a2 = tf.constant([6, 7, 8, 9, 10], dtype=tf.int64)
+    a3 = tf.constant([11, 12, 13, 14, 15], dtype=tf.int64)
+    a_123 = [a1, a2, a3]
+    b_123 = tf.constant([101, 102, 103, 104, 105], dtype=tf.int64)
+    cond1 = tf.constant([False, False, True, False, False], dtype=bool)
+    cond2 = tf.constant([False, False, False, False, True], dtype=bool)
+    cond3 = tf.constant([True, False, True, False, True], dtype=bool)
+    cond_123 = [cond1, cond2, cond3]
+    mux_123 = multiplex_4_op.multiplex(cond_123, a_123, b_123)
+    b4 = tf.constant([201, 202, 203, 204, 205], dtype=tf.int64)
+    cond4 = tf.constant([True, True, True, False, False], dtype=bool)
+    result = multiplex_4_op.multiplex(cond4, mux_123, b4)
+    return result
+```
+
+The model_using_multiplex.py file has functions for creating and using a saved
+custom op model `SavedModel`. In this test, the `multiplex_4` op is used to both
+save and use models.
+
+```
+  def test_multiplex_saved_model(self):
+    path = os.path.join(self.create_tempdir(), 'model')
+    model_using_multiplex.save(multiplex_4_op.multiplex, path)
+    result = model_using_multiplex.load_and_use(path)
+    self.assertAllEqual(result, tf.constant([1, 20, 3, 40, 5], dtype=tf.int64))
+```
+
+Test the op with the following:
+
+<!-- test_snippets_in_readme skip -->
+```shell
+bazel test //third_party/tensorflow/google/g3doc/example/multiplex_4:multiplex_4_test
+```
+
+Reuse the `BUILD` file to add build rules for the Python API wrapper and the op
+test.
+
+```
+py_strict_library(
+    name = "multiplex_4_op",
+    srcs = ["multiplex_4_op.py"],
+    data = ["multiplex_4_kernel.so"],
+    srcs_version = "PY3",
+    deps = [
+        "//third_party/py/tensorflow",
+    ],
+)
+```
+<!-- test_snippets_in_readme skip -->
+```
+tf_py_test(
+    name = "multiplex_4_test",
+    size = "medium",  # This test blocks because it writes and reads a file,
+    srcs = ["multiplex_4_test.py"],
+    python_version = "PY3",
+    srcs_version = "PY3",
+    tags = [
+        "no_mac",
+        "no_pip",
+    ],
+    deps = [
+        ":model_using_multiplex",
+        ":multiplex_4_op",
+        "//third_party/py/numpy",
+        "//third_party/py/tensorflow",
+        "//third_party/tensorflow/python/framework:errors",
+        "//third_party/tensorflow/python/framework:test_lib",
+    ],
+)
+```
+
+### Use the op
+
+Build the op with the following:
+
+<!-- test_snippets_in_readme skip -->
+```shell
+bazel build //third_party/tensorflow/examples/custom_ops_doc/multiplex_4:multiplex_4_op
+```
+
+Import the op and call it using the following example:
+
+<!-- test_snippets_in_readme skip -->
+```python
+import tensorflow as tf
+
+from tensorflow.examples.custom_ops_doc.multiplex_4 import multiplex_4_op
+
+a1 = tf.constant([1, 2, 3, 4, 5], dtype=tf.int64)
+a2 = tf.constant([6, 7, 8, 9, 10], dtype=tf.int64)
+a3 = tf.constant([11, 12, 13, 14, 15], dtype=tf.int64)
+a = [a1, a2, a3]
+b = tf.constant([101, 102, 103, 104, 105], dtype=tf.int64)
+cond1 = tf.constant([False, False, True, False, False], dtype=bool)
+cond2 = tf.constant([False, False, False, False, True], dtype=bool)
+cond3 = tf.constant([True, False, True, False, True], dtype=bool)
+cond = [cond1, cond2, cond3]
+# expected result is [11, 102, 3, 104, 10]
+result = multiplex_4_op.multiplex(cond, a, b)
+```
+
+The `multiplex_4_load_use.py` file uses `load_and_use` from
+`model_using_muliplex.py` to load a saved model from a `multiplex_2` op. The
+saved model can be executed using the new kernel, (`multiplex_4`), which
+supports both lists of tensors and single tensors for `cond` and `a` inputs.
+
+Since `Examples>MultiplexDense` can only be defined once in a binary, there must
+be two separate binaries. A binary can either depend on `multiplex_2_op` or
+`multiplex_4_op`, but not both. The custom ops are backward compatible, so we
+can use `save` on `multiplex_2` and `load_and_use` on `multiplex_4`.
+
+### Summary
+
+In this example, you learned how to implement a new multiplexer kernel that is
+backwards compatible with an existing multiplexer kernel. This custom op handles
+inputs that are lists of tensors, while continuing to handle inputs of single
+tensors.
+
+The tables below summarize the build rules and targets for building and testing
+the `multiplex_4` op.
+
+#### Kernel components
+
+Op components                           | Build rule             | Build target         | Source
+--------------------------------------- | ---------------------- | -------------------- | ------
+Kernels (C++)                           | `tf_custom_op_library` | `multiplex_4_kernel` | `multiplex_4_kernel.cc`, `multiplex_4_op.cc`
+Wrapper (automatically generated)       | N/A                    | `gen_multiplex_4_op` | N/A
+Wrapper (with public API and docstring) | `py_strict_library`    | `multiplex_4_op`     | `multiplex_4_op.py`
+Tests                                   | `tf_py_test`           | `multiplex_4_test`   | `multiplex_4_test.py`
+
+##### Usage example
+
+Op components            | Build rule          | Build target               | Source
+------------------------ | ------------------- | -------------------------- | ------
+Common library           | `py_strict_library` | `model_using_multiplex`    | `model_using_multiplex.py`
+Old op (with SavedModel) | `py_strict_binary`  | `multiplex_2_save`         | `multiplex_2_save.py`
+New op (with SavedModel) | `py_strict_binary`  | `multiplex_4_load_and_use` | `multiplex_4_load_and_use.py`
+
+## Resources
+
+*   [OSS custom ops guide](https://www.tensorflow.org/guide/create_op)
+*   [SavedModel](https://www.tensorflow.org/guide/saved_model)
+*   [Numpy Select](https://numpy.org/doc/stable/reference/generated/numpy.select.html)
+<!-- LINT.ThenChange(multiplex_4.md) -->

tensorflow/examples/custom_ops_doc/multiplex_4/model_using_multiplex.py

@@ -0,0 +1,77 @@
+# Copyright 2021 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.
+# ==============================================================================
+"""Example of using multiplex op in a SavedModel.
+
+multiplex_2_save.py and multiplex_4_load_use.py are programs that use this.
+
+https://www.tensorflow.org/guide/saved_model
+https://www.tensorflow.org/api_docs/python/tf/saved_model/save
+"""
+
+import tensorflow as tf
+
+
+def _get_example_tensors():
+  cond = tf.constant([True, False, True, False, True], dtype=bool)
+  a = tf.constant([1, 2, 3, 4, 5], dtype=tf.int64)
+  b = tf.constant([10, 20, 30, 40, 50], dtype=tf.int64)
+  return cond, a, b
+
+
+def save(multiplex_op, path):
+  """Save a model that contains the given `multiplex_op`.
+
+  Args:
+    multiplex_op: A multiplex Custom Op, e.g. multiplex_4_op.multiplex. This is
+      parameterized so it can also be used to create an "old" model with an
+      older version of the op, e.g. multiplex_2_op.multiplex.
+    path: Directory to save model to.
+  """
+  example_cond, example_a, example_b = _get_example_tensors()
+
+  class UseMultiplex(tf.Module):
+
+    @tf.function(input_signature=[
+        tf.TensorSpec.from_tensor(example_cond),
+        tf.TensorSpec.from_tensor(example_a),
+        tf.TensorSpec.from_tensor(example_b)
+    ])
+    def use_multiplex(self, cond, a, b):
+      return multiplex_op(cond, a, b)
+
+  model = UseMultiplex()
+  tf.saved_model.save(
+      model,
+      path,
+      signatures=model.use_multiplex.get_concrete_function(
+          tf.TensorSpec.from_tensor(example_cond),
+          tf.TensorSpec.from_tensor(example_a),
+          tf.TensorSpec.from_tensor(example_b)))
+
+
+def load_and_use(path):
+  """Load and used a model that was previously created by `save()`.
+
+  Args:
+    path: Directory to load model from, typically the same directory that was
+      used by save().
+
+  Returns:
+    A tensor that is the result of using the multiplex op that is
+    tf.constant([1, 20, 3, 40, 5], dtype=tf.int64).
+  """
+  example_cond, example_a, example_b = _get_example_tensors()
+  restored = tf.saved_model.load(path)
+  return restored.use_multiplex(example_cond, example_a, example_b)

tensorflow/examples/custom_ops_doc/multiplex_4/multiplex_2_save.py

@@ -0,0 +1,41 @@
+# Copyright 2021 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.
+# ==============================================================================
+"""Binary for showing C++ backward compatibility.
+
+This creates a SavedModel using the "old" op and C++ kernel from multiplex_2.
+
+https://www.tensorflow.org/guide/saved_model
+https://www.tensorflow.org/api_docs/python/tf/saved_model/save
+"""
+
+import os
+import shutil
+
+from absl import app
+from tensorflow.examples.custom_ops_doc.multiplex_2 import multiplex_2_op
+from tensorflow.examples.custom_ops_doc.multiplex_4 import model_using_multiplex
+
+
+def main(argv):
+  del argv  # not used
+  path = 'model_using_multiplex'
+  if os.path.exists(path):
+    shutil.rmtree(path, ignore_errors=True)
+  model_using_multiplex.save(multiplex_2_op.multiplex, path)
+  print('Saved model to', path)
+
+
+if __name__ == '__main__':
+  app.run(main)

tensorflow/examples/custom_ops_doc/multiplex_4/multiplex_4_kernel.cc

@@ -0,0 +1,137 @@
+/* Copyright 2021 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 "absl/log/log.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/register_types.h"
+
+// Please use the appropriate namespace for your project
+namespace tensorflow {
+namespace custom_op_examples {
+
+using ::tensorflow::OpKernel;
+using ::tensorflow::OpKernelConstruction;
+using ::tensorflow::OpKernelContext;
+using ::tensorflow::Tensor;
+using ::tensorflow::errors::Internal;
+using ::tensorflow::errors::InvalidArgument;
+
+// Multiple types for the values inside two of the input tensors
+// (e.g. int32, float) are supported by using a template where the type is T.
+template <typename T>
+class MultiplexDenseOp : public OpKernel {
+ public:
+  explicit MultiplexDenseOp(OpKernelConstruction* ctx) : OpKernel(ctx) {
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("N", &num_cond_a_));
+  }
+
+  MultiplexDenseOp(const MultiplexDenseOp& other) = delete;
+  MultiplexDenseOp& operator=(const MultiplexDenseOp& other) = delete;
+  ~MultiplexDenseOp() override = default;
+
+  void Compute(OpKernelContext* ctx) override {
+    // Optional error checking: cond and a_values are lists of N, so there are
+    // a total of 2N+1 inputs. Check that the  number of inputs and the
+    // `N` Attr is consistent.
+    const int64_t expected_inputs = 2 * num_cond_a_ + 1;
+    OP_REQUIRES(ctx, expected_inputs == ctx->num_inputs(),
+                Internal("expected_inputs != num_inputs(): ", expected_inputs,
+                         " != ", ctx->num_inputs()));
+    VLOG(1) << "N " << num_cond_a_;
+
+    const auto& first_cond_tensor = ctx->input(0);
+    const auto& first_a_values_tensor = ctx->input(num_cond_a_);
+    const auto& b_values_tensor = ctx->input(2 * num_cond_a_);
+
+    // Allow any shape, but require that a_values, b_values, and cond all
+    // have the same shape.
+    // Note that ::tensorflow::TensorShapeUtils has some useful functions
+    // for checking shapes.
+    for (int64_t i = 0; i < num_cond_a_; i++) {
+      const auto& cond_tensor_i = ctx->input(i);
+      const auto& a_values_tensor_i = ctx->input(num_cond_a_ + i);
+      OP_REQUIRES(
+          ctx, a_values_tensor_i.shape() == b_values_tensor.shape(),
+          InvalidArgument(
+              "a_values[", i,
+              "] and b_values must have the same shape. "
+              "a_values[",
+              i, "] shape: ", a_values_tensor_i.DebugString(),
+              " b_values shape: ", b_values_tensor.shape().DebugString()));
+      OP_REQUIRES(
+          ctx, cond_tensor_i.shape() == b_values_tensor.shape(),
+          InvalidArgument(
+              "cond_values[", i,
+              "] and b_valuesmust have the same shape. "
+              "cond_values[",
+              i, "] shape: ", first_a_values_tensor.shape().DebugString(),
+              " b_values shape: ", first_cond_tensor.shape().DebugString()));
+    }
+
+    // Create an output tensor
+    Tensor* output_tensor = nullptr;
+    OP_REQUIRES_OK(
+        ctx, ctx->allocate_output(0, b_values_tensor.shape(), &output_tensor));
+    auto output = output_tensor->template flat<T>();
+
+    const auto b_values = b_values_tensor.template flat<T>();
+    // np.select style behavior, `cond` and `a_values` are lists of tensors.
+    // Also works for the np.where style case where there is only one `cond`
+    // and one `a_values` tensor.
+    const int64_t N = first_a_values_tensor.NumElements();
+    for (int64_t i = 0; i < N; i++) {
+      bool flag = false;
+      for (int64_t list_index = 0; list_index < num_cond_a_; list_index++) {
+        const auto& cond_tensor = ctx->input(list_index);
+        const auto& a_values_tensor = ctx->input(num_cond_a_ + list_index);
+        const auto cond = cond_tensor.template flat<bool>();
+        const auto a_values = a_values_tensor.template flat<T>();
+        if (cond(i)) {
+          output(i) = a_values(i);
+          flag = true;
+          VLOG(1) << "A " << list_index << " for " << i;
+          break;
+        }
+      }
+      if (!flag) {
+        output(i) = b_values(i);
+        VLOG(1) << "B for " << i;
+      }
+    }
+  }
+
+ private:
+  int64_t num_cond_a_;  // the number of `cond` and `a` input tensors
+};
+
+// The "Name" used by REGISTER_KERNEL_BUILDER is defined by REGISTER_OP,
+// see multiplex_4_op.cc.
+// To support tensors containing different types (e.g. int32, float), one
+// kernel per type is registered and is templatized by the "T" attr value.
+// The TF_CALL_ALL_TYPES macro registers the op for all types appropriate for
+// the target platform. See go/tf-custom-ops-guide
+#define REGISTER_KERNELS(type)                                  \
+  REGISTER_KERNEL_BUILDER(Name("Examples>MultiplexDense")       \
+                              .Device(::tensorflow::DEVICE_CPU) \
+                              .TypeConstraint<type>("T"),       \
+                          MultiplexDenseOp<type>)
+
+TF_CALL_ALL_TYPES(REGISTER_KERNELS);
+#undef REGISTER_KERNELS
+
+}  // namespace custom_op_examples
+}  // namespace tensorflow

tensorflow/examples/custom_ops_doc/multiplex_4/multiplex_4_load_use.py

@@ -0,0 +1,37 @@
+# Copyright 2021 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.
+# ==============================================================================
+"""Binary for showing C++ backward compatibility.
+
+This loads a previously created SavedModel (esp. a model created by
+multiplex_2_save.py which uses the "old" op and C++ kernel from multiplex_2)
+and runs the model using the "new" multiplex_4 C++ kernel.
+
+https://www.tensorflow.org/guide/saved_model
+https://www.tensorflow.org/api_docs/python/tf/saved_model/save
+"""
+
+from absl import app
+from tensorflow.examples.custom_ops_doc.multiplex_4 import model_using_multiplex
+
+
+def main(argv):
+  del argv  # not used
+  path = 'model_using_multiplex'
+  result = model_using_multiplex.load_and_use(path)
+  print('Result:', result)
+
+
+if __name__ == '__main__':
+  app.run(main)

tensorflow/examples/custom_ops_doc/multiplex_4/multiplex_4_op.cc

@@ -0,0 +1,79 @@
+/* Copyright 2021 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 "absl/status/status.h"
+#include "tensorflow/core/framework/op.h"
+#include "tensorflow/core/framework/shape_inference.h"
+
+// Use a namespace when registering by prepending the
+// package's name to the op’s name and separate with a '>'.
+// This is the recommendation for out-of-tree ops to avoid name collisions in
+// "Best practices for custom operations in TensorFlow"
+// https://github.com/tensorflow/community/blob/master/rfcs/20190726-custom-ops.md
+
+// Please use the appropriate namespace for your project
+namespace tensorflow {
+namespace custom_op_examples {
+
+using ::tensorflow::shape_inference::InferenceContext;
+
+absl::Status MultiplexShapeFunction(InferenceContext* c) {
+  int64_t num_cond_a;
+  TF_RETURN_IF_ERROR(c->GetAttr("N", &num_cond_a));
+  tensorflow::shape_inference::ShapeHandle unused;
+  // There are N `cond` inputs, N `a` inputs and 1 `b` input in this order.
+  // Check that all `cond` inputs and all `a` inputs have the same shape
+  // as the `b` input.
+  int64_t last = 2 * num_cond_a;
+  for (int64_t i = 0; i < last; i++) {
+    TF_RETURN_IF_ERROR(c->Merge(c->input(i), c->input(last), &unused));
+  }
+  c->set_output(0, c->input(last));
+  return absl::OkStatus();
+}
+
+REGISTER_OP("Examples>MultiplexDense")
+    .Input("cond: N * bool")
+    .Input("a_values: N * T")
+    .Input("b_values: T")
+    .Output("output_values: T")
+    .Attr("T: type")
+    .Attr("N: int = 1")
+    .SetShapeFn(MultiplexShapeFunction)
+    .Doc(R"doc(
+Return elements chosen from `a_values` or `b_values` depending on `cond`.
+
+When `a_values` and `cond` are tenors (i.e. N=1), this is similar to `np.where`
+and `tf.where`. When `a_values` and `cond` are lists of tensors (i.e. N>1),
+this is similar to `np.select`. In either case these are simplified to only
+handle dense tensors, no optional parameters, no broadcasting, etc..
+
+cond: tf.Tensor or list of tf.Tensor of type bool. If it is a list, `a_values`
+      must be a list of the same length. Where True, yield the corresponding
+      element from `a_values` (with priority to the first one encountered in
+      lists), otherwise yield `b_values`.
+a_values: tf.Tensor or list of tf.Tensor. Each tensor has the same type and
+          shape as `b_values`. If it is a list, `cond` must be a list of the
+          same length.
+b_values: tf.Tensor with the same type and shape as the `a_values` if it is a
+          tensor or as every element of `a_values` if `a_values` is a list.
+output_values: A tf.Tensor with elements from `a_values` where `cond` is True,
+               and elements from `b` elsewhere.
+)doc");
+
+}  // namespace custom_op_examples
+}  // namespace tensorflow

tensorflow/examples/custom_ops_doc/multiplex_4/multiplex_4_op.py

@@ -0,0 +1,70 @@
+# Copyright 2021 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.
+# ==============================================================================
+"""A wrapper for gen_multiplex_4_op.py.
+
+This defines a public API (and provides a docstring for it) for the C++ Op
+defined by multiplex_4_kernel.cc
+"""
+
+import tensorflow as tf
+from tensorflow.python.platform import resource_loader
+
+_multiplex_4_module = tf.load_op_library(
+    resource_loader.get_path_to_datafile("multiplex_4_kernel.so"))
+
+examples_multiplex_dense = _multiplex_4_module.examples_multiplex_dense
+
+
+def multiplex(cond, a, b, name=None):
+  """Return elements chosen from `a` or `b` depending on `cond`.
+
+  This is similar to `np.where` and `tf.where` if `cond` and `a` are tensors.
+  This is similar to `np.select` if `cond` and `a` are lists of tensors.
+  In either case, this is simplified to only handle the case of dense tensors,
+  no optional parameters, no broadcasting, etc..
+
+  >>> multiplex([True, False, False, True], [1,2,3,4], [100,200,300,400])
+  <tf.Tensor: shape=(4,), dtype=int32, numpy=array([  1, 200, 300,   4], ...)>
+
+  >>> a1 = tf.constant([1, 2, 3, 4, 5], dtype=tf.int64)
+  >>> a2 = tf.constant([6, 7, 8, 9, 10], dtype=tf.int64)
+  >>> a3 = tf.constant([11, 12, 13, 14, 15], dtype=tf.int64)
+  >>> b = tf.constant([101, 102, 103, 104, 105], dtype=tf.int64)
+  >>> cond1 = tf.constant([False, False, True, False, False], dtype=bool)
+  >>> cond2 = tf.constant([False, False, False, False, True], dtype=bool)
+  >>> cond3 = tf.constant([True, False, True, False, True], dtype=bool)
+  >>> multiplex_4_op.multiplex([cond1, cond2, cond3], [a1, a2, a3], b)
+  <tf.Tensor: shape=(5,), ... numpy=array([ 11, 102,   3, 104,  10], ...)>
+
+  Args:
+    cond: tf.Tensor or list of tf.Tensor of type bool. Where True, yield `a`.
+      When multiple corresponding `cond` elements are true, the first one yield
+      based on the first one encountered.
+    a: tf.Tensor or list of tf.Tensor, each with the same type and shape as `b`.
+    b: tf.Tensor or list of tf.Tensor with the same type and shape as `a`. Yield
+      `b` if all corresponding `cond` values is False.
+    name: An optional name for the op.
+
+  Returns:
+    A tf.Tensor with elements from `a` where `cond` is True, and elements
+    from `b` elsewhere.
+  """
+  if not isinstance(cond, (list, tuple)):
+    # Support "old" use of multiplex where `cond` and `a` are tensors,
+    # not lists of tensors.
+    return examples_multiplex_dense(
+        cond=[cond], a_values=[a], b_values=b, name=name)
+  return examples_multiplex_dense(
+      cond=cond, a_values=a, b_values=b, name=name)

tensorflow/examples/custom_ops_doc/multiplex_4/multiplex_4_test.py

@@ -0,0 +1,109 @@
+# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for multiplex_4."""
+
+import os
+
+import numpy as np
+import tensorflow as tf
+
+from tensorflow.examples.custom_ops_doc.multiplex_4 import model_using_multiplex
+from tensorflow.examples.custom_ops_doc.multiplex_4 import multiplex_4_op
+# This pylint disable is only needed for internal google users
+from tensorflow.python.framework import errors_impl  # pylint: disable=g-direct-tensorflow-import
+from tensorflow.python.framework import test_util  # pylint: disable=g-direct-tensorflow-import
+
+
+@test_util.with_eager_op_as_function
+class MultiplexOpTest(tf.test.TestCase):
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_multiplex_int(self):
+    a = tf.constant([1, 2, 3, 4, 5], dtype=tf.int64)
+    b = tf.constant([10, 20, 30, 40, 50], dtype=tf.int64)
+    cond = tf.constant([True, False, True, False, True], dtype=bool)
+    expect = np.where(self.evaluate(cond), self.evaluate(a), self.evaluate(b))
+    # expected result is [1, 20, 3, 40, 5]
+    result = multiplex_4_op.multiplex(cond, a, b)
+    self.assertAllEqual(result, expect)
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_multiplex_select(self):
+    a1 = tf.constant([1, 2, 3, 4, 5], dtype=tf.int64)
+    a2 = tf.constant([6, 7, 8, 9, 10], dtype=tf.int64)
+    a3 = tf.constant([11, 12, 13, 14, 15], dtype=tf.int64)
+    a = [a1, a2, a3]
+    b = tf.constant([101, 102, 103, 104, 105], dtype=tf.int64)
+    cond1 = tf.constant([False, False, True, False, False], dtype=bool)
+    cond2 = tf.constant([False, False, False, False, True], dtype=bool)
+    cond3 = tf.constant([True, False, True, False, True], dtype=bool)
+    cond = [cond1, cond2, cond3]
+    expect = np.select([self.evaluate(i) for i in cond],
+                       [self.evaluate(i) for i in a], self.evaluate(b))
+    # expected result is [11, 102, 3, 104, 10]
+    result = multiplex_4_op.multiplex(cond, a, b)
+    self.assertAllEqual(result, expect)
+
+  def test_multiplex_saved_model(self):
+    path = os.path.join(self.create_tempdir(), 'model')
+    model_using_multiplex.save(multiplex_4_op.multiplex, path)
+    result = model_using_multiplex.load_and_use(path)
+    self.assertAllEqual(result, tf.constant([1, 20, 3, 40, 5], dtype=tf.int64))
+
+  # One tf.function that uses both multiplex with single tensors for `cond`
+  # and `a` and with lists of tensors for `cond` and `a`, i.e. a graph
+  # with two example_multiplex_dense kernels that have different numbers
+  # of inputs.
+  @tf.function
+  def _both(self):
+    a1 = tf.constant([1, 2, 3, 4, 5], dtype=tf.int64)
+    a2 = tf.constant([6, 7, 8, 9, 10], dtype=tf.int64)
+    a3 = tf.constant([11, 12, 13, 14, 15], dtype=tf.int64)
+    a_123 = [a1, a2, a3]
+    b_123 = tf.constant([101, 102, 103, 104, 105], dtype=tf.int64)
+    cond1 = tf.constant([False, False, True, False, False], dtype=bool)
+    cond2 = tf.constant([False, False, False, False, True], dtype=bool)
+    cond3 = tf.constant([True, False, True, False, True], dtype=bool)
+    cond_123 = [cond1, cond2, cond3]
+    mux_123 = multiplex_4_op.multiplex(cond_123, a_123, b_123)
+    b4 = tf.constant([201, 202, 203, 204, 205], dtype=tf.int64)
+    cond4 = tf.constant([True, True, True, False, False], dtype=bool)
+    result = multiplex_4_op.multiplex(cond4, mux_123, b4)
+    return result
+
+  def test_both_single_and_list(self):
+    result = self._both()
+    self.assertAllEqual(result,
+                        tf.constant([11, 102, 3, 204, 205], dtype=tf.int64))
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_inconsistent_inputs_error(self):
+    a1 = tf.constant([1, 2, 3, 4, 5], dtype=tf.int64)
+    a2 = tf.constant([6, 7, 8, 9, 10], dtype=tf.int64)
+    a = [a1, a2]
+    b = tf.constant([101, 102, 103, 104, 105], dtype=tf.int64)
+    cond = tf.constant([False, False, True, False, False], dtype=bool)
+    with self.assertRaisesRegex(
+        (errors_impl.InvalidArgumentError, ValueError),
+        # Eager mode raises InvalidArgumentError with the following message
+        r'(a_values\[0\] and b_values must have the same shape'
+        r')|('
+        # Graph mode raises ValueError with the following message
+        r'Shapes must be equal rank, but are 2 and 1)'):
+      self.evaluate(multiplex_4_op.multiplex(cond, a, b))
+
+
+if __name__ == '__main__':
+  tf.test.main()

tensorflow/examples/custom_ops_doc/simple_hash_table/BUILD

@@ -0,0 +1,60 @@
+# Build simple_hash_table custom ops example, which is similar to,
+# but simpler than, tf.lookup.experimental.MutableHashTable
+
+load("//tensorflow:strict.default.bzl", "py_strict_library")
+load("//tensorflow:tensorflow.bzl", "tf_custom_op_library")
+load("//tensorflow:tensorflow.default.bzl", "tf_py_test")
+
+licenses(["notice"])
+
+tf_custom_op_library(
+    name = "simple_hash_table_kernel.so",
+    srcs = [
+        "simple_hash_table_kernel.cc",
+        "simple_hash_table_op.cc",
+    ],
+    deps = [
+        "//tensorflow/core/lib/gtl:map_util",
+        "//tensorflow/core/platform:strcat",
+        "@com_google_absl//absl/container:flat_hash_map",
+    ],
+)
+
+py_strict_library(
+    name = "simple_hash_table_op",
+    srcs = ["simple_hash_table_op.py"],
+    data = ["simple_hash_table_kernel.so"],
+    srcs_version = "PY3",
+    deps = [
+        "//tensorflow:tensorflow_py",
+    ],
+)
+
+py_strict_library(
+    name = "simple_hash_table",
+    srcs = ["simple_hash_table.py"],
+    srcs_version = "PY3",
+    deps = [
+        ":simple_hash_table_op",
+        "//tensorflow:tensorflow_py",
+    ],
+)
+
+tf_py_test(
+    name = "simple_hash_table_test",
+    size = "medium",  # This test blocks because it writes and reads a file,
+    timeout = "short",  # but it still runs quickly.
+    srcs = ["simple_hash_table_test.py"],
+    python_version = "PY3",
+    srcs_version = "PY3",
+    tags = [
+        "no_mac",  # TODO(b/216321151): Re-enable this test.
+    ],
+    deps = [
+        ":simple_hash_table",
+        "//tensorflow:tensorflow_py",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:test_lib",
+        "//third_party/py/numpy",
+    ],
+)

tensorflow/examples/custom_ops_doc/simple_hash_table/README.md

@@ -0,0 +1,899 @@
+<!-- LINT.IfChange -->
+# Create a stateful custom op
+
+This example shows you how to create TensorFlow custom ops with internal state.
+These custom ops use resources to hold their state. You do this by implementing
+a stateful C++ data structure instead of using tensors to hold the state. This
+example also covers:
+
+*   Using `tf.Variable`s for state as an alternative to using a resource (which
+    is recommended for all cases where storing the state in fixed size tensors
+    is reasonable)
+*   Saving and restoring the state of the custom op by using `SavedModel`s
+*   The `IsStateful` flag which is true for ops with state and other unrelated
+    cases
+*   A set of related custom ops and ops with various signatures (number of
+    inputs and outputs)
+
+For additional context,
+read the
+[OSS guide on creating custom ops](https://www.tensorflow.org/guide/create_op).
+
+## Background
+
+### Overview of resources and ref-counting
+
+TensorFlow C++ resource classes are derived from the
+[`ResourceBase` class](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/framework/resource_base.h).
+A resource is referenced in the Tensorflow Graph as a
+[`ResourceHandle`](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/framework/resource_handle.h)
+Tensor, of the type
+[`DT_RESOURCE`](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/framework/types.proto).
+A resource can be owned by a ref-counting `ResourceHandle`, or by a
+[`ResourceMgr`](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/framework/resource_mgr.h)
+(deprecated for new ops).
+
+A handle-owned resource using ref-counting is automatically destroyed when all
+resource handles pointing to it go out of scope. If the resource needs to be
+looked up from a name, for example, if the resource handle is serialized and
+deserialized, the resource must be published as an unowned entry with
+[`ResourceMgr::CreateUnowned`](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/framework/resource_mgr.h#L641)
+upon creation. The entry is a weak reference to the resource, and is
+automatically removed when the resource goes out of scope.
+
+In contrast, with the deprecated `ResourceMgr` owned resources, a resource
+handle behaves like a weak ref - it does not destroy the underlying resource
+when its lifetime ends.
+[`DestroyResourceOp`](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/kernels/resource_variable_ops.cc#L339)
+must be explicitly called to destroy the resource. An example of why requiring
+calling `DestroyResourceOp` is problematic is that it is easy to introduce a bug
+when adding a new code path to an op that returns without calling
+`DestroyResourceOp`.
+
+While it would be possible to have a data structure that has a single operation
+implemented by a single custom op (for example, something that just has a 'next
+state' operation that requires no initialization), typically a related set of
+custom ops are used with a resource. **Separating creation and use into
+different custom ops is recommended.** Typically, one custom op creates the
+resource and one or more additional custom ops implement functionality to access
+or modify it.
+
+### Using `tf.Variable`s for state
+
+An alternative to custom ops with internal state is to store the state
+externally in one or more
+[`tf.Variable`](https://www.tensorflow.org/api_docs/python/tf/Variable)s as
+tensors (as detailed [here](https://www.tensorflow.org/guide/variable)) and have
+one or more (normal, stateless) ops that use tensors stored in these
+`tf.Variable`s. One example is `tf.random.Generator`.
+
+For cases where using variables is possible and efficient, using them is
+preferred since the implementation does not require adding a new C++ resource
+type. An example of a case where using variables is not possible and custom ops
+using a resource must be used is where the amount of space for data grows
+dynamically.
+
+Here is a toy example of using the
+`multiplex_2`
+custom op with a `tf.Variable` in the same manner as an in-built TensorFlow op
+is used with variables. The variable is initialized to 1 for every value.
+Indices from a
+[`tf.data.Dataset`](https://www.tensorflow.org/api_docs/python/tf/data/Dataset)
+cause the corresponding element of the variable to be doubled.
+
+<!-- test_snippets_in_readme skip -->
+```python
+import tensorflow as tf
+from tensorflow.examples.custom_ops_doc.multiplex_2 import multiplex_2_op
+
+def variable_and_stateless_op():
+  n = 10
+  v = tf.Variable(tf.ones(n, dtype=tf.int64), trainable=False)
+  dataset = tf.data.Dataset.from_tensor_slices([5, 1, 7, 5])
+  for position in dataset:
+    print(v.numpy())
+    cond = tf.one_hot(
+        position, depth=n, on_value=True, off_value=False, dtype=bool)
+    v.assign(multiplex_2_op.multiplex(cond, v*2, v))
+  print(v.numpy())
+```
+
+This outputs:
+
+<!-- test_snippets_in_readme skip -->
+```
+[1 1 1 1 1 1 1 1 1 1]
+[1 1 1 1 1 2 1 1 1 1]
+[1 2 1 1 1 2 1 1 1 1]
+[1 2 1 1 1 2 1 2 1 1]
+[1 2 1 1 1 4 1 2 1 1]
+```
+
+It is also possible to pass a Python `tf.Variable` handle as a
+[`DT_RESOURCE`](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/framework/types.proto)
+input to a custom op where its C++ kernel uses the handle to access the variable
+using the variable's internal C++ API. This is not a common case because the
+variable's C++ API provides little extra functionality. It can be appropriate
+for cases that only sparsely update the variable.
+
+### The `IsStateful` flag
+
+All TensorFlow ops have an
+[`IsStateful`](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/framework/resource_mgr.h#L505)
+flag. It is set to `True` for ops that have internal state, both for ops that
+follow the recommended pattern of using a resource and those that have internal
+state in some other way. In addition, it is also set to `True` for ops with side
+effects (for example, I/O) and for disabling certain optimizations for an op.
+
+### Save and restore state with `SavedModel`
+
+A [`SavedModel`](https://www.tensorflow.org/guide/saved_model) contains a
+complete TensorFlow program, including trained parameters (like `tf.Variable`s)
+and computation. In some cases, you can save and restore the state in a custom
+op by using `SavedModel`. This is optional, but is important for some cases, for
+example if a custom op is used during training and the state needs to persist
+when training is restarted from a checkpoint. The hash table in this example
+supports `SavedModel` with a custom Python wrapper that implements the
+`_serialize_to_tensors` and `_restore_from_tensors` methods of
+[`tf.saved_model.experimental.TrackableResource`](https://www.tensorflow.org/api_docs/python/tf/saved_model/experimental/TrackableResource)
+and ops that are used by these methods.
+
+## Creating a hash table with a stateful custom op
+
+This example demonstrates how to implement custom ops that use ref-counting
+resources to track state. This example creates a `simple_hash_table` which is
+similar to
+[`tf.lookup.experimental.MutableHashTable`](https://www.tensorflow.org/api_docs/python/tf/lookup/experimental/MutableHashTable).
+
+The hash table implements a set of 4 CRUD (Create/Read/Update/Delete) style
+custom ops. Each of these ops has a different signature and they illustrate
+cases such as custom ops with more than one output.
+
+The hash table supports `SavedModel` through the `export` and `import` ops to
+save/restore state. For an actual hash table use case, it is preferable to use
+(or extend) the existing
+[`tf.lookup`](https://www.tensorflow.org/api_docs/python/tf/lookup) ops. In this
+simple example, `insert`, `find`, and `remove` use only a single key-value pair
+per call. In contrast, existing `tf.lookup` ops can use multiple key-value
+pairs in a single call.
+
+The table below summarizes the six custom ops implemented in this example.
+
+Operation | Purpose                     | Resource class method | Kernel class                    | Custom op                      | Python class member
+--------- | --------------------------- | --------------------- | ------------------------------- | ------------------------------ | -------------------
+create    | CRUD and SavedModel: create | Default constructor   | `SimpleHashTableCreateOpKernel` | Examples>SimpleHashTableCreate | `__init__`
+find      | CRUD: read                  | Find                  | `SimpleHashTableFindOpKernel`   | Examples>SimpleHashTableFind   | `find`
+insert    | CRUD: update                | Insert                | `SimpleHashTableInsertOpKernel` | Examples>SimpleHashTableInsert | `insert`
+remove    | CRUD: delete                | Remove                | `SimpleHashTableRemoveOpKernel` | Examples>SimpleHashTableRemove | `remove`
+import    | SavedModel: restore         | Import                | `SimpleHashTableImportOpKernel` | Examples>SimpleHashTableImport | `do_import`
+export    | SavedModel: save            | Export                | `SimpleHashTableExportOpKernel` | Examples>SimpleHashTableExport | `export`
+
+You can use this hash table as:
+
+<!-- test_snippets_in_readme skip -->
+```python
+hash_table = simple_hash_table_op.SimpleHashTable(tf.int32, float,
+                                                  default_value=-999.0)
+result1 = hash_table.find(key=1, dynamic_default_value=-999.0)
+# -999.0
+hash_table.insert(key=1, value=100.0)
+result2 = hash_table.find(key=1, dynamic_default_value=-999.0)
+# 100.0
+```
+
+The example below contains C++ and Python code snippets to illustrate the code
+flow. These snippets are not all complete; some are missing namespace
+declarations, imports, and test cases.
+
+This example deviates slightly from the general recipe for creating TensorFlow
+custom ops. The most significant differences are noted in each step below.
+
+### Step 0 - Implement the resource class
+
+Implement a
+`SimpleHashTableResource`
+resource class derived from
+[`ResourceBase`](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/framework/resource_base.h).
+
+```live-snippet
+template <class K, class V>
+class SimpleHashTableResource : public ::tensorflow::ResourceBase {
+ public:
+  Status Insert(const Tensor& key, const Tensor& value) {
+    const K key_val = key.flat<K>()(0);
+    const V value_val = value.flat<V>()(0);
+
+    mutex_lock l(mu_);
+    table_[key_val] = value_val;
+    return OkStatus();
+  }
+
+  Status Find(const Tensor& key, Tensor* value, const Tensor& default_value) {
+    // Note that tf_shared_lock could be used instead of mutex_lock
+    // in ops that do not not modify data protected by a mutex, but
+    // go/totw/197 recommends using exclusive lock instead of a shared
+    // lock when the lock is not going to be held for a significant amount
+    // of time.
+    mutex_lock l(mu_);
+
+    const V default_val = default_value.flat<V>()(0);
+    const K key_val = key.flat<K>()(0);
+    auto value_val = value->flat<V>();
+    value_val(0) = gtl::FindWithDefault(table_, key_val, default_val);
+    return OkStatus();
+  }
+
+  Status Remove(const Tensor& key) {
+    mutex_lock l(mu_);
+
+    const K key_val = key.flat<K>()(0);
+    if (table_.erase(key_val) != 1) {
+      return errors::NotFound("Key for remove not found: ", key_val);
+    }
+    return OkStatus();
+  }
+
+  // Save all key, value pairs to tensor outputs to support SavedModel
+  Status Export(OpKernelContext* ctx) {
+    mutex_lock l(mu_);
+    int64_t size = table_.size();
+    Tensor* keys;
+    Tensor* values;
+    TF_RETURN_IF_ERROR(
+        ctx->allocate_output("keys", TensorShape({size}), &keys));
+    TF_RETURN_IF_ERROR(
+        ctx->allocate_output("values", TensorShape({size}), &values));
+    auto keys_data = keys->flat<K>();
+    auto values_data = values->flat<V>();
+    int64_t i = 0;
+    for (auto it = table_.begin(); it != table_.end(); ++it, ++i) {
+      keys_data(i) = it->first;
+      values_data(i) = it->second;
+    }
+    return OkStatus();
+  }
+
+  // Load all key, value pairs from tensor inputs to support SavedModel
+  Status Import(const Tensor& keys, const Tensor& values) {
+    const auto key_values = keys.flat<K>();
+    const auto value_values = values.flat<V>();
+
+    mutex_lock l(mu_);
+    table_.clear();
+    for (int64_t i = 0; i < key_values.size(); ++i) {
+      gtl::InsertOrUpdate(&table_, key_values(i), value_values(i));
+    }
+    return OkStatus();
+  }
+
+  // Create a debug string with the content of the map if this is small,
+  // or some example data if this is large, handling both the cases where the
+  // hash table has many entries and where the entries are long strings.
+  std::string DebugString() const override { return DebugString(3); }
+  std::string DebugString(int num_pairs) const {
+    std::string rval = "SimpleHashTable {";
+    size_t count = 0;
+    const size_t max_kv_str_len = 100;
+    mutex_lock l(mu_);
+    for (const auto& pair : table_) {
+      if (count >= num_pairs) {
+        absl::StrAppend(&rval, "...");
+        break;
+      }
+      std::string kv_str = absl::StrCat(pair.first, ": ", pair.second);
+      absl::StrAppend(&rval, kv_str.substr(0, max_kv_str_len));
+      if (kv_str.length() > max_kv_str_len) absl::StrAppend(&rval, " ...");
+      absl::StrAppend(&rval, ", ");
+      count += 1;
+    }
+    absl::StrAppend(&rval, "}");
+    return rval;
+  }
+
+ private:
+  mutable mutex mu_;
+  absl::flat_hash_map<K, V> table_ TF_GUARDED_BY(mu_);
+};
+```
+
+Note that this class provides:
+
+*   Helper methods to access the hash table. These methods correspond to the
+    `find`, `insert`, and `remove` ops.
+*   Helper methods to `import`/`export` the complete internal state of the hash
+    table. These methods help support `SavedModel`.
+*   A [mutex](https://en.wikipedia.org/wiki/Lock_\(computer_science\)) for the
+    helper methods to use for exclusive access to the
+    [`absl::flat_hash_map`](https://abseil.io/docs/cpp/guides/container#abslflat_hash_map-and-abslflat_hash_set).
+    This ensures thread safety by ensuring that only one thread at a time can
+    access the data in the hash table.
+
+### Step 1 - Define the op interface
+
+Define op interfaces and register all the custom ops you create for the hash
+table. You typically define one custom op to create the resource. You also
+define one or more custom ops that correspond to operations on the data
+structure. You also define custom ops to perform import and export operations
+that input/output the whole internal state. These ops are optional; you define
+them in this example to support `SavedModel`. As the resource is automatically
+deleted based on ref-counting, there is no custom op required to delete the
+resource.
+
+The `simple_hash_table` has kernels for the `create`, `insert`, `find`,
+`remove`, `import`, and `export` ops which use the resource object that actually
+stores and manipulates data. The resource is passed between ops using a resource
+handle. The `create` op has an `Output` of type `resource`. The other ops have
+an `Input` of type `resource`. The interface definitions for all the ops along
+with their shape functions are in
+`simple_hash_table_op.cc`.
+
+Note that these definitions do not explicitly use `SetIsStateful`. The
+`IsStateful` flag is set automatically for any op with an input or output of
+type `DT_RESOURCE`.
+
+The definitions below and their corresponding kernels and generated Python
+wrappers illustrate the following cases:
+
+*   0, 1, 2, and 3 inputs
+*   0, 1, and 2 outputs
+*   `Attr` for types where none are used by `Input` or `Output` and all have to
+    be explicitly passed into the generated wrapper (for example,
+    `SimpleHashTableCreate`)
+*   `Attr` for types where all are used by `Input` and/or `Output` (so they are
+    set implicitly inside the generated wrapper) and none are passed into the
+    generated wrapper (for example, `SimpleHashTableFind`)
+*   `Attr` for types where some but not all are used by `Input` or `Output` and
+    only those that are not used are explicitly passed into the generated
+    wrapper (for example, `SimpleHashTableRemove` where there is an `Input` that
+    uses `key_dtype` but the `value_type` `Attr` is a parameter to the generated
+    wrapper).
+
+```
+REGISTER_OP("Examples>SimpleHashTableCreate")
+    .Output("output: resource")
+    .Attr("key_dtype: type")
+    .Attr("value_dtype: type")
+    .SetShapeFn(ScalarOutput);
+```
+```
+REGISTER_OP("Examples>SimpleHashTableFind")
+    .Input("resource_handle: resource")
+    .Input("key: key_dtype")
+    .Input("default_value: value_dtype")
+    .Output("value: value_dtype")
+    .Attr("key_dtype: type")
+    .Attr("value_dtype: type")
+    .SetShapeFn(ThreeScalarInputsScalarOutput);
+```
+```
+REGISTER_OP("Examples>SimpleHashTableInsert")
+    .Input("resource_handle: resource")
+    .Input("key: key_dtype")
+    .Input("value: value_dtype")
+    .Attr("key_dtype: type")
+    .Attr("value_dtype: type")
+    .SetShapeFn(ThreeScalarInputs);
+```
+```
+REGISTER_OP("Examples>SimpleHashTableRemove")
+    .Input("resource_handle: resource")
+    .Input("key: key_dtype")
+    .Attr("key_dtype: type")
+    .Attr("value_dtype: type")
+    .SetShapeFn(TwoScalarInputs);
+```
+```
+REGISTER_OP("Examples>SimpleHashTableExport")
+    .Input("table_handle: resource")
+    .Output("keys: key_dtype")
+    .Output("values: value_dtype")
+    .Attr("key_dtype: type")
+    .Attr("value_dtype: type")
+    .SetShapeFn(ExportShapeFunction);
+```
+```
+REGISTER_OP("Examples>SimpleHashTableImport")
+    .Input("table_handle: resource")
+    .Input("keys: key_dtype")
+    .Input("values: value_dtype")
+    .Attr("key_dtype: type")
+    .Attr("value_dtype: type")
+    .SetShapeFn(ImportShapeFunction);
+```
+
+### Step 2 - Register the op implementation (kernel)
+
+Declare kernels for specific types of key-value pairs. Register the kernel by
+calling the `REGISTER_KERNEL_BUILDER` macro.
+
+```
+#define REGISTER_KERNEL(key_dtype, value_dtype)               \
+  REGISTER_KERNEL_BUILDER(                                    \
+      Name("Examples>SimpleHashTableCreate")                  \
+          .Device(DEVICE_CPU)                                 \
+          .TypeConstraint<key_dtype>("key_dtype")             \
+          .TypeConstraint<value_dtype>("value_dtype"),        \
+      SimpleHashTableCreateOpKernel<key_dtype, value_dtype>); \
+  REGISTER_KERNEL_BUILDER(                                    \
+      Name("Examples>SimpleHashTableFind")                    \
+          .Device(DEVICE_CPU)                                 \
+          .TypeConstraint<key_dtype>("key_dtype")             \
+          .TypeConstraint<value_dtype>("value_dtype"),        \
+      SimpleHashTableFindOpKernel<key_dtype, value_dtype>);   \
+  REGISTER_KERNEL_BUILDER(                                    \
+      Name("Examples>SimpleHashTableInsert")                  \
+          .Device(DEVICE_CPU)                                 \
+          .TypeConstraint<key_dtype>("key_dtype")             \
+          .TypeConstraint<value_dtype>("value_dtype"),        \
+      SimpleHashTableInsertOpKernel<key_dtype, value_dtype>)  \
+  REGISTER_KERNEL_BUILDER(                                    \
+      Name("Examples>SimpleHashTableRemove")                  \
+          .Device(DEVICE_CPU)                                 \
+          .TypeConstraint<key_dtype>("key_dtype")             \
+          .TypeConstraint<value_dtype>("value_dtype"),        \
+      SimpleHashTableRemoveOpKernel<key_dtype, value_dtype>)  \
+  REGISTER_KERNEL_BUILDER(                                    \
+      Name("Examples>SimpleHashTableExport")                  \
+          .Device(DEVICE_CPU)                                 \
+          .TypeConstraint<key_dtype>("key_dtype")             \
+          .TypeConstraint<value_dtype>("value_dtype"),        \
+      SimpleHashTableExportOpKernel<key_dtype, value_dtype>)  \
+  REGISTER_KERNEL_BUILDER(                                    \
+      Name("Examples>SimpleHashTableImport")                  \
+          .Device(DEVICE_CPU)                                 \
+          .TypeConstraint<key_dtype>("key_dtype")             \
+          .TypeConstraint<value_dtype>("value_dtype"),        \
+      SimpleHashTableImportOpKernel<key_dtype, value_dtype>);
+```
+```
+REGISTER_KERNEL(int32, double);
+REGISTER_KERNEL(int32, float);
+REGISTER_KERNEL(int32, int32);
+REGISTER_KERNEL(int32, tstring);
+REGISTER_KERNEL(int64_t, double);
+REGISTER_KERNEL(int64_t, float);
+REGISTER_KERNEL(int64_t, int32);
+REGISTER_KERNEL(int64_t, int64_t);
+REGISTER_KERNEL(int64_t, tstring);
+REGISTER_KERNEL(tstring, bool);
+REGISTER_KERNEL(tstring, double);
+REGISTER_KERNEL(tstring, float);
+REGISTER_KERNEL(tstring, int32);
+REGISTER_KERNEL(tstring, int64_t);
+REGISTER_KERNEL(tstring, tstring);
+```
+
+### Step 3 - Implement the op kernel(s)
+
+The implementation of kernels for the ops in the hash table all use helper
+functions from the `SimpleHashTableResource` resource class.
+
+You implement the op kernels in two phases:
+
+1.  Implement the kernel for the `create` op that has a `Compute` method that
+    creates a resource object and a ref-counted handle for it.
+
+    <!-- test_snippets_in_readme skip -->
+    ```c++
+    handle_tensor.scalar<ResourceHandle>()() =
+        ResourceHandle::MakeRefCountingHandle(
+            new SimpleHashTableResource<K, V>(), /* … */);
+    ```
+
+1.  Implement the kernels(s) for each of the other operations on the resource
+    that have Compute methods that get a resource object and use one or more of
+    its helper methods.
+
+    <!-- test_snippets_in_readme skip -->
+    ```c++
+    MyResource* resource;
+    OP_REQUIRES_OK(ctx, GetResource(ctx, &resource));
+    // The GetResource local function uses handle.GetResource<resource_type>()
+    OP_REQUIRES_OK(ctx, resource->Find(key, out, default_value));
+    ```
+
+#### Creating the resource
+
+The `create` op creates a resource object of type `SimpleHashTableResource` and
+then uses `MakeRefCountingHandle` to pass the ownership to a resource handle.
+This op outputs a `resource` handle.
+
+```
+template <class K, class V>
+class SimpleHashTableCreateOpKernel : public OpKernel {
+ public:
+  explicit SimpleHashTableCreateOpKernel(OpKernelConstruction* ctx)
+      : OpKernel(ctx) {}
+
+  void Compute(OpKernelContext* ctx) override {
+    Tensor handle_tensor;
+    AllocatorAttributes attr;
+    OP_REQUIRES_OK(ctx, ctx->allocate_temp(DT_RESOURCE, TensorShape({}),
+                                           &handle_tensor, attr));
+    handle_tensor.scalar<ResourceHandle>()() =
+        ResourceHandle::MakeRefCountingHandle(
+            new SimpleHashTableResource<K, V>(), ctx->device()->name(),
+            /*dtypes_and_shapes=*/{}, ctx->stack_trace());
+    ctx->set_output(0, handle_tensor);
+  }
+
+ private:
+  // Just to be safe, avoid accidentally copying the kernel.
+  SimpleHashTableCreateOpKernel(const SimpleHashTableCreateOpKernel&) = delete;
+  void operator=(const SimpleHashTableCreateOpKernel&) = delete;
+};
+```
+
+#### Getting the resource
+
+In
+`simple_hash_table_kernel.cc`,
+the `GetResource` helper function uses an input `resource` handle to retrieve
+the corresponding `SimpleHashTableResource` object. It is used by all the custom
+ops that use the resource (that is, all the custom ops in the set other than
+`create`).
+
+```
+template <class K, class V>
+Status GetResource(OpKernelContext* ctx,
+                   SimpleHashTableResource<K, V>** resource) {
+  const Tensor& handle_tensor = ctx->input(0);
+  const ResourceHandle& handle = handle_tensor.scalar<ResourceHandle>()();
+  typedef SimpleHashTableResource<K, V> resource_type;
+  TF_ASSIGN_OR_RETURN(*resource, handle.GetResource<resource_type>());
+  return OkStatus();
+}
+```
+
+#### Using the resource
+
+The ops that use the resource use `GetResource` to get a pointer to the resource
+object and call the corresponding helper function for that object. Below is the
+source code for the `find` op which uses `resource->Find`. The other ops that
+use the resource similarly use the corresponding helper method in the resource
+class.
+
+```
+template <class K, class V>
+class SimpleHashTableFindOpKernel : public OpKernel {
+ public:
+  explicit SimpleHashTableFindOpKernel(OpKernelConstruction* ctx)
+      : OpKernel(ctx) {}
+
+  void Compute(OpKernelContext* ctx) override {
+    DataTypeVector expected_inputs = {DT_RESOURCE, DataTypeToEnum<K>::v(),
+                                      DataTypeToEnum<V>::v()};
+    DataTypeVector expected_outputs = {DataTypeToEnum<V>::v()};
+    OP_REQUIRES_OK(ctx, ctx->MatchSignature(expected_inputs, expected_outputs));
+    SimpleHashTableResource<K, V>* resource;
+    OP_REQUIRES_OK(ctx, GetResource(ctx, &resource));
+    // Note that ctx->input(0) is the Resource handle
+    const Tensor& key = ctx->input(1);
+    const Tensor& default_value = ctx->input(2);
+    TensorShape output_shape = default_value.shape();
+    Tensor* out;
+    OP_REQUIRES_OK(ctx, ctx->allocate_output("value", output_shape, &out));
+    OP_REQUIRES_OK(ctx, resource->Find(key, out, default_value));
+  }
+};
+```
+
+#### Compile the op
+
+Compile the C++ op to create a kernel library and Python wrapper that enables
+you to use the op with TensorFlow.
+
+Create a `BUILD` file for the op which declares the dependencies and the output
+build targets. Refer to
+[building for OSS](https://www.tensorflow.org/guide/create_op#build_the_op_library).
+
+Note
+that you will be reusing this `BUILD` file later on in this example.
+
+```
+tf_custom_op_library(
+    name = "simple_hash_table_kernel.so",
+    srcs = [
+        "simple_hash_table_kernel.cc",
+        "simple_hash_table_op.cc",
+    ],
+    deps = [
+        "//third_party/absl/container:flat_hash_map",
+        "//third_party/tensorflow/core/lib/gtl:map_util",
+        "//third_party/tensorflow/core/platform:strcat",
+    ],
+)
+
+py_strict_library(
+    name = "simple_hash_table_op",
+    srcs = ["simple_hash_table_op.py"],
+    data = ["simple_hash_table_kernel.so"],
+    srcs_version = "PY3",
+    deps = [
+        "//third_party/py/tensorflow",
+    ],
+)
+
+py_strict_library(
+    name = "simple_hash_table",
+    srcs = ["simple_hash_table.py"],
+    srcs_version = "PY3",
+    deps = [
+        ":simple_hash_table_op",
+        "//third_party/py/tensorflow",
+    ],
+)
+```
+
+### Step 4 - Create the Python wrapper
+
+To create the Python wrapper, import and implement a function that serves as the
+op's public API and provides a docstring.
+
+The Python wrapper for this example, `simple_hash_table.py` uses the
+`SimpleHashTable` class to provide methods that allow access to the data
+structure. The custom ops are used to implement these methods. This class also
+supports `SavedModel` which allows the state of the resource to be saved and
+restored for checkpointing. The class is derived from the
+`tf.saved_model.experimental.TrackableResource` base class and implements the
+`_serialize_to_tensors` (which uses the `export` op) and `_restore_from_tensors`
+(which uses the `import` op) methods defined by this base class. The `__init__`
+method of this class calls the `_create_resource` method (which is an override
+of a base class method required for `SavedModel`) which in turn calls the C++
+`examples_simple_hash_table_create` kernel. The resource handle returned by this
+kernel is stored in the private `self._resource_handle` object member. Methods
+that use this handle access it using the public `self.resource_handle` property
+provided by the `tf.saved_model.experimental.TrackableResource` base class.
+
+```
+  def _create_resource(self):
+    """Create the resource tensor handle.
+
+    `_create_resource` is an override of a method in base class
+    `TrackableResource` that is required for SavedModel support. It can be
+    called by the `resource_handle` property defined by `TrackableResource`.
+
+    Returns:
+      A tensor handle to the lookup table.
+    """
+    assert self._default_value.get_shape().ndims == 0
+    table_ref = gen_simple_hash_table_op.examples_simple_hash_table_create(
+        key_dtype=self._key_dtype,
+        value_dtype=self._value_dtype,
+        name=self._name)
+    return table_ref
+
+  def _serialize_to_tensors(self):
+    """Implements checkpointing protocols for `Trackable`."""
+    tensors = self.export()
+    return {"table-keys": tensors[0], "table-values": tensors[1]}
+
+  def _restore_from_tensors(self, restored_tensors):
+    """Implements checkpointing protocols for `Trackable`."""
+    return gen_simple_hash_table_op.examples_simple_hash_table_import(
+        self.resource_handle, restored_tensors["table-keys"],
+        restored_tensors["table-values"])
+```
+
+The `find` method in this class calls the `examples_simple_hash_table_find`
+custom op using the reference handle (from the public `self.resource_handle`
+property) and a key and default value. The other methods are similar. It is
+recommended that methods for using a resource avoid logic other than a call to a
+generated Python wrapper to avoid eager/`tf.function` inconsistencies (avoid
+Python logic that is lost when a `tf.function` is created). In the case of
+`find`, using `tf.convert_to_tensor` cannot be avoided and is not lost during
+`tf.function` creation.
+
+```
+  def find(self, key, dynamic_default_value=None, name=None):
+    """Looks up `key` in a table, outputs the corresponding value.
+
+    The `default_value` is used if key not present in the table.
+
+    Args:
+      key: Key to look up. Must match the table's key_dtype.
+      dynamic_default_value: The value to use if the key is missing in the
+        table. If None (by default), the `table.default_value` will be used.
+      name: A name for the operation (optional).
+
+    Returns:
+      A tensor containing the value in the same shape as `key` using the
+        table's value type.
+
+    Raises:
+      TypeError: when `key` do not match the table data types.
+    """
+    with tf.name_scope(name or "%s_lookup_table_find" % self._name):
+      key = tf.convert_to_tensor(key, dtype=self._key_dtype, name="key")
+      if dynamic_default_value is not None:
+        dynamic_default_value = tf.convert_to_tensor(
+            dynamic_default_value,
+            dtype=self._value_dtype,
+            name="default_value")
+      value = gen_simple_hash_table_op.examples_simple_hash_table_find(
+          self.resource_handle, key, dynamic_default_value
+          if dynamic_default_value is not None else self._default_value)
+    return value
+```
+
+The Python wrapper specifies that gradients are not implemented in this example.
+For an example of a differentiable map. For general information about gradients,
+read
+[Implement the gradient in Python](https://www.tensorflow.org/guide/create_op#implement_the_gradient_in_python)
+in the OSS guide.
+
+```
+tf.no_gradient("Examples>SimpleHashTableCreate")
+tf.no_gradient("Examples>SimpleHashTableFind")
+tf.no_gradient("Examples>SimpleHashTableInsert")
+tf.no_gradient("Examples>SimpleHashTableRemove")
+```
+
+The full source code for the Python wrapper is in
+`simple_hash_table_op.py`]
+and
+`simple_hash_table.py`.
+
+### Step 5 - Test the op
+
+Create op tests using classes derived from
+[`tf.test.TestCase`](https://www.tensorflow.org/api_docs/python/tf/test/TestCase)
+and the
+[parameterized tests provided by Abseil](https://github.com/abseil/abseil-py/blob/main/absl/testing/parameterized.py).
+
+Create tests using three or more custom ops to create `SimpleHashTable` objects
+and:
+
+1.  Update the state using the `insert`, `remove`, and/or `import` methods
+1.  Observe the state using the `find` and/or `export` methods
+
+Here is an example test:
+
+```
+  def test_find_insert_find_strings_eager(self):
+    default = 'Default'
+    foo = 'Foo'
+    bar = 'Bar'
+    hash_table = simple_hash_table.SimpleHashTable(tf.string, tf.string,
+                                                   default)
+    result1 = hash_table.find(foo, default)
+    self.assertEqual(result1, default)
+    hash_table.insert(foo, bar)
+    result2 = hash_table.find(foo, default)
+    self.assertEqual(result2, bar)
+```
+
+Create a helper function to work with the hash table:
+
+```
+  def _use_table(self, key_dtype, value_dtype):
+    hash_table = simple_hash_table.SimpleHashTable(key_dtype, value_dtype, 111)
+    result1 = hash_table.find(1, -999)
+    hash_table.insert(1, 100)
+    result2 = hash_table.find(1, -999)
+    hash_table.remove(1)
+    result3 = hash_table.find(1, -999)
+    results = tf.stack((result1, result2, result3))
+    return results  # expect [-999, 100, -999]
+```
+
+The following test explicitly creates a `tf.function` with `_use_table`.
+Ref-counting causes the C++ resource object created in `_use_table` to be
+destroyed when this `tf.function` returns inside `self.evaluate(results)`. By
+explicitly creating the `tf.function` instead of relying on decorators like
+`@test_util.with_eager_op_as_function` and
+`@test_util.run_in_graph_and_eager_modes` (such as in
+`multiplex_1`),
+you have an explicit place in the test corresponding to where the C++ resource's
+destructor is called.
+
+This test also shows a test that is parameterized for different data types; it
+is actually two tests, one with `tf.int32` input / `float` output and the other
+with `tf.int64` input / `tf.int32` output.
+
+```
+  def test_find_insert_find_tf_function(self, key_dtype, value_dtype):
+    results = def_function.function(
+        lambda: self._use_table(key_dtype, value_dtype))
+    self.assertAllClose(self.evaluate(results), [-999.0, 100.0, -999.0])
+```
+
+Reuse the `BUILD` file to add build
+rules for the Python API wrapper and the op test.
+
+```
+tf_custom_op_library(
+    name = "simple_hash_table_kernel.so",
+    srcs = [
+        "simple_hash_table_kernel.cc",
+        "simple_hash_table_op.cc",
+    ],
+    deps = [
+        "//third_party/absl/container:flat_hash_map",
+        "//third_party/tensorflow/core/lib/gtl:map_util",
+        "//third_party/tensorflow/core/platform:strcat",
+    ],
+)
+
+py_strict_library(
+    name = "simple_hash_table_op",
+    srcs = ["simple_hash_table_op.py"],
+    data = ["simple_hash_table_kernel.so"],
+    srcs_version = "PY3",
+    deps = [
+        "//third_party/py/tensorflow",
+    ],
+)
+
+py_strict_library(
+    name = "simple_hash_table",
+    srcs = ["simple_hash_table.py"],
+    srcs_version = "PY3",
+    deps = [
+        ":simple_hash_table_op",
+        "//third_party/py/tensorflow",
+    ],
+)
+
+tf_py_test(
+    name = "simple_hash_table_test",
+    size = "medium",  # This test blocks because it writes and reads a file,
+    timeout = "short",  # but it still runs quickly.
+    srcs = ["simple_hash_table_test.py"],
+    python_version = "PY3",
+    srcs_version = "PY3",
+    tags = [
+        "no_mac",  # TODO(b/216321151): Re-enable this test.
+    ],
+    deps = [
+        ":simple_hash_table",
+        "//third_party/py/numpy",
+        "//third_party/py/tensorflow",
+        "//third_party/tensorflow/python/framework:errors",
+        "//third_party/tensorflow/python/framework:test_lib",
+    ],
+)
+```
+
+Test the op by running:
+
+<!-- test_snippets_in_readme skip -->
+```shell
+$ bazel test //third_party/tensorflow/google/g3doc/example/simple_hash_table:simple_hash_table_test
+```
+
+### Use the op
+
+Use the op by importing and calling it as follows:
+
+<!-- test_snippets_in_readme skip -->
+```python
+import tensorflow as tf
+
+from tensorflow.examples.custom_ops_doc.simple_hash_table import simple_hash_table
+
+hash_table = simple_hash_table.SimpleHashTable(tf.int32, float, -999.0)
+result1 = hash_table.find(1, -999.0)  # -999.0
+hash_table.insert(1, 100.0)
+result2 = hash_table.find(1, -999.0)  # 100.0
+```
+
+Here, `simple_hash_table` is the name of the Python wrapper that was created
+in this example.
+
+### Summary
+
+In this example, you learned how to implement a simple hash table data structure
+using stateful custom ops.
+
+The table below summarizes the build rules and targets for building and testing
+the `simple_hash_table` op.
+
+Op components                           | Build rule             | Build target               | Source
+--------------------------------------- | ---------------------- | -------------------------- | ------
+Kernels (C++)                           | `tf_custom_op_library` | `simple_hash_table_kernel` | `simple_hash_table_kernel.cc`, `simple_hash_table_op.cc`
+Wrapper (automatically generated)       | N/A.                   | `gen_simple_hash_table_op` | N/A
+Wrapper (with public API and docstring) | `py_strict_library`    | `simple_hash_table_op`, `simple_hash_table`     | `simple_hash_table_op.py`, `simple_hash_table.py`
+Tests                                   | `tf_py_test`           | `simple_hash_table_test`   | `simple_hash_table_test.py`
+<!-- LINT.ThenChange(simple_hash_table.md) -->
+

tensorflow/examples/custom_ops_doc/simple_hash_table/simple_hash_table.py

@@ -0,0 +1,243 @@
+# Copyright 2021 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.
+# ==============================================================================
+"""A wrapper for gen_simple_hash_table_op.py.
+
+This defines a public API and provides a docstring for the C++ Op defined by
+simple_hash_table_kernel.cc
+"""
+
+import tensorflow as tf
+from tensorflow.examples.custom_ops_doc.simple_hash_table.simple_hash_table_op import gen_simple_hash_table_op
+
+
+class SimpleHashTable(tf.saved_model.experimental.TrackableResource):
+  """A simple mutable hash table implementation.
+
+  Implement a simple hash table as a Resource using ref-counting.
+  This demonstrates a Stateful Op for a general Create/Read/Update/Delete
+  (CRUD) style use case.  To instead make an op for a specific lookup table
+  case, it is preferable to follow the implementation style of
+  TensorFlow's internal ops, e.g. use LookupInterface.
+
+  Data can be inserted by calling the `insert` method and removed by calling
+  the `remove` method. It does not support initialization via the init method.
+
+  The `import` and `export` methods allow loading and restoring all of
+  the key, value pairs. These methods (or their corresponding kernels)
+  are intended to be used for supporting SavedModel.
+
+  Example usage:
+    hash_table = simple_hash_table_op.SimpleHashTable(key_dtype, value_dtype,
+                                                      111)
+    result1 = hash_table.find(1, -999)  # -999
+    hash_table.insert(1, 100)
+    result2 = hash_table.find(1, -999)  # 100
+    hash_table.remove(1)
+    result3 = hash_table.find(1, -999)  # -999
+  """
+
+  def __init__(self,
+               key_dtype,
+               value_dtype,
+               default_value,
+               name="SimpleHashTable"):
+    """Creates an empty `SimpleHashTable` object.
+
+    Creates a table, the type of its keys and values are specified by key_dtype
+    and value_dtype, respectively.
+
+    Args:
+      key_dtype: the type of the key tensors.
+      value_dtype: the type of the value tensors.
+      default_value: The value to use if a key is missing in the table.
+      name: A name for the operation (optional).
+
+    Returns:
+      A `SimpleHashTable` object.
+    """
+    super(SimpleHashTable, self).__init__()
+    self._default_value = tf.convert_to_tensor(default_value, dtype=value_dtype)
+    self._value_shape = self._default_value.get_shape()
+    self._key_dtype = key_dtype
+    self._value_dtype = value_dtype
+    self._name = name
+    self._resource_handle = self._create_resource()
+    # Methods that use the Resource get its handle using the
+    # public self.resource_handle property (defined by TrackableResource).
+    # This property calls self._create_resource() the first time
+    # if private self._resource_handle is not preemptively initialized.
+
+  def _create_resource(self):
+    """Create the resource tensor handle.
+
+    `_create_resource` is an override of a method in base class
+    `TrackableResource` that is required for SavedModel support. It can be
+    called by the `resource_handle` property defined by `TrackableResource`.
+
+    Returns:
+      A tensor handle to the lookup table.
+    """
+    assert self._default_value.get_shape().ndims == 0
+    table_ref = gen_simple_hash_table_op.examples_simple_hash_table_create(
+        key_dtype=self._key_dtype,
+        value_dtype=self._value_dtype,
+        name=self._name)
+    return table_ref
+
+  def _serialize_to_tensors(self):
+    """Implements checkpointing protocols for `Trackable`."""
+    tensors = self.export()
+    return {"table-keys": tensors[0], "table-values": tensors[1]}
+
+  def _restore_from_tensors(self, restored_tensors):
+    """Implements checkpointing protocols for `Trackable`."""
+    return gen_simple_hash_table_op.examples_simple_hash_table_import(
+        self.resource_handle, restored_tensors["table-keys"],
+        restored_tensors["table-values"])
+
+  @property
+  def key_dtype(self):
+    """The table key dtype."""
+    return self._key_dtype
+
+  @property
+  def value_dtype(self):
+    """The table value dtype."""
+    return self._value_dtype
+
+  def find(self, key, dynamic_default_value=None, name=None):
+    """Looks up `key` in a table, outputs the corresponding value.
+
+    The `default_value` is used if key not present in the table.
+
+    Args:
+      key: Key to look up. Must match the table's key_dtype.
+      dynamic_default_value: The value to use if the key is missing in the
+        table. If None (by default), the `table.default_value` will be used.
+      name: A name for the operation (optional).
+
+    Returns:
+      A tensor containing the value in the same shape as `key` using the
+        table's value type.
+
+    Raises:
+      TypeError: when `key` do not match the table data types.
+    """
+    with tf.name_scope(name or "%s_lookup_table_find" % self._name):
+      key = tf.convert_to_tensor(key, dtype=self._key_dtype, name="key")
+      if dynamic_default_value is not None:
+        dynamic_default_value = tf.convert_to_tensor(
+            dynamic_default_value,
+            dtype=self._value_dtype,
+            name="default_value")
+      value = gen_simple_hash_table_op.examples_simple_hash_table_find(
+          self.resource_handle, key, dynamic_default_value
+          if dynamic_default_value is not None else self._default_value)
+    return value
+
+  def insert(self, key, value, name=None):
+    """Associates `key` with `value`.
+
+    Args:
+      key: Scalar key to insert.
+      value: Scalar value to be associated with key.
+      name: A name for the operation (optional).
+
+    Returns:
+      The created Operation.
+
+    Raises:
+      TypeError: when `key` or `value` doesn't match the table data
+        types.
+    """
+    with tf.name_scope(name or "%s_lookup_table_insert" % self._name):
+      key = tf.convert_to_tensor(key, self._key_dtype, name="key")
+      value = tf.convert_to_tensor(value, self._value_dtype, name="value")
+      # pylint: disable=protected-access
+      op = gen_simple_hash_table_op.examples_simple_hash_table_insert(
+          self.resource_handle, key, value)
+      return op
+
+  def remove(self, key, name=None):
+    """Remove `key`.
+
+    Args:
+      key: Scalar key to remove.
+      name: A name for the operation (optional).
+
+    Returns:
+      The created Operation.
+
+    Raises:
+      TypeError: when `key` doesn't match the table data type.
+    """
+    with tf.name_scope(name or "%s_lookup_table_remove" % self._name):
+      key = tf.convert_to_tensor(key, self._key_dtype, name="key")
+
+      # For remove, just the key is used by the kernel; no value is used.
+      # But the kernel is specific to key_dtype and value_dtype
+      # (i.e. it uses a <key_dtype, value_dtype> template).
+      # So value_dtype is passed in explicitly. (While
+      # key_dtype is specified implicitly by the dtype of key.)
+
+      # pylint: disable=protected-access
+      op = gen_simple_hash_table_op.examples_simple_hash_table_remove(
+          self.resource_handle, key, value_dtype=self._value_dtype)
+      return op
+
+  def export(self, name=None):
+    """Export all `key` and `value` pairs.
+
+    Args:
+      name: A name for the operation (optional).
+
+    Returns:
+      A tuple of two tensors, the first with the `keys` and the second with
+      the `values`.
+    """
+    with tf.name_scope(name or "%s_lookup_table_export" % self._name):
+      # pylint: disable=protected-access
+      keys, values = gen_simple_hash_table_op.examples_simple_hash_table_export(
+          self.resource_handle,
+          key_dtype=self._key_dtype,
+          value_dtype=self._value_dtype)
+      return keys, values
+
+  def do_import(self, keys, values, name=None):
+    """Import all `key` and `value` pairs.
+
+    (Note that "import" is a python reserved word, so it cannot be the name of
+    a method.)
+
+    Args:
+      keys: Tensor of all keys.
+      values: Tensor of all values.
+      name: A name for the operation (optional).
+
+    Returns:
+      A tuple of two tensors, the first with the `keys` and the second with
+      the `values`.
+    """
+    with tf.name_scope(name or "%s_lookup_table_import" % self._name):
+      # pylint: disable=protected-access
+      op = gen_simple_hash_table_op.examples_simple_hash_table_import(
+          self.resource_handle, keys, values)
+      return op
+
+
+tf.no_gradient("Examples>SimpleHashTableCreate")
+tf.no_gradient("Examples>SimpleHashTableFind")
+tf.no_gradient("Examples>SimpleHashTableInsert")
+tf.no_gradient("Examples>SimpleHashTableRemove")

tensorflow/examples/custom_ops_doc/simple_hash_table/simple_hash_table_kernel.cc

@@ -0,0 +1,354 @@
+/* Copyright 2021 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 <cstddef>
+#include <cstdint>
+#include <string>
+
+#include "absl/container/flat_hash_map.h"
+#include "absl/strings/str_cat.h"
+#include "xla/tsl/platform/errors.h"
+#include "xla/tsl/platform/statusor.h"
+#include "tensorflow/core/framework/allocator.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/op_requires.h"
+#include "tensorflow/core/framework/resource_base.h"
+#include "tensorflow/core/framework/resource_handle.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/lib/gtl/map_util.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/mutex.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/platform/strcat.h"
+#include "tensorflow/core/platform/tstring.h"
+#include "tensorflow/core/platform/types.h"
+#include "tsl/platform/thread_annotations.h"
+
+// Please use the appropriate namespace for your project
+namespace tensorflow {
+namespace custom_op_examples {
+
+// Implement a simple hash table as a Resource using ref-counting.
+// This demonstrates a Stateful Op for a general Create/Read/Update/Delete
+// (CRUD) style use case.  To instead make an op for a specific lookup table
+// case, it is preferable to follow the implementation style of
+// the kernels for TensorFlow's tf.lookup internal ops which use
+// LookupInterface.
+template <class K, class V>
+class SimpleHashTableResource : public ::tensorflow::ResourceBase {
+ public:
+  Status Insert(const Tensor& key, const Tensor& value) {
+    const K key_val = key.flat<K>()(0);
+    const V value_val = value.flat<V>()(0);
+
+    mutex_lock l(mu_);
+    table_[key_val] = value_val;
+    return OkStatus();
+  }
+
+  Status Find(const Tensor& key, Tensor* value, const Tensor& default_value) {
+    // Note that tf_shared_lock could be used instead of mutex_lock
+    // in ops that do not not modify data protected by a mutex, but
+    // go/totw/197 recommends using exclusive lock instead of a shared
+    // lock when the lock is not going to be held for a significant amount
+    // of time.
+    mutex_lock l(mu_);
+
+    const V default_val = default_value.flat<V>()(0);
+    const K key_val = key.flat<K>()(0);
+    auto value_val = value->flat<V>();
+    value_val(0) = gtl::FindWithDefault(table_, key_val, default_val);
+    return OkStatus();
+  }
+
+  Status Remove(const Tensor& key) {
+    mutex_lock l(mu_);
+
+    const K key_val = key.flat<K>()(0);
+    if (table_.erase(key_val) != 1) {
+      return errors::NotFound("Key for remove not found: ", key_val);
+    }
+    return OkStatus();
+  }
+
+  // Save all key, value pairs to tensor outputs to support SavedModel
+  Status Export(OpKernelContext* ctx) {
+    mutex_lock l(mu_);
+    int64_t size = table_.size();
+    Tensor* keys;
+    Tensor* values;
+    TF_RETURN_IF_ERROR(
+        ctx->allocate_output("keys", TensorShape({size}), &keys));
+    TF_RETURN_IF_ERROR(
+        ctx->allocate_output("values", TensorShape({size}), &values));
+    auto keys_data = keys->flat<K>();
+    auto values_data = values->flat<V>();
+    int64_t i = 0;
+    for (auto it = table_.begin(); it != table_.end(); ++it, ++i) {
+      keys_data(i) = it->first;
+      values_data(i) = it->second;
+    }
+    return OkStatus();
+  }
+
+  // Load all key, value pairs from tensor inputs to support SavedModel
+  Status Import(const Tensor& keys, const Tensor& values) {
+    const auto key_values = keys.flat<K>();
+    const auto value_values = values.flat<V>();
+
+    mutex_lock l(mu_);
+    table_.clear();
+    for (int64_t i = 0; i < key_values.size(); ++i) {
+      gtl::InsertOrUpdate(&table_, key_values(i), value_values(i));
+    }
+    return OkStatus();
+  }
+
+  // Create a debug string with the content of the map if this is small,
+  // or some example data if this is large, handling both the cases where the
+  // hash table has many entries and where the entries are long strings.
+  std::string DebugString() const override { return DebugString(3); }
+  std::string DebugString(int num_pairs) const {
+    std::string rval = "SimpleHashTable {";
+    size_t count = 0;
+    const size_t max_kv_str_len = 100;
+    mutex_lock l(mu_);
+    for (const auto& pair : table_) {
+      if (count >= num_pairs) {
+        absl::StrAppend(&rval, "...");
+        break;
+      }
+      std::string kv_str = absl::StrCat(pair.first, ": ", pair.second);
+      absl::StrAppend(&rval, kv_str.substr(0, max_kv_str_len));
+      if (kv_str.length() > max_kv_str_len) absl::StrAppend(&rval, " ...");
+      absl::StrAppend(&rval, ", ");
+      count += 1;
+    }
+    absl::StrAppend(&rval, "}");
+    return rval;
+  }
+
+ private:
+  mutable mutex mu_;
+  absl::flat_hash_map<K, V> table_ TF_GUARDED_BY(mu_);
+};
+
+template <class K, class V>
+class SimpleHashTableCreateOpKernel : public OpKernel {
+ public:
+  explicit SimpleHashTableCreateOpKernel(OpKernelConstruction* ctx)
+      : OpKernel(ctx) {}
+
+  void Compute(OpKernelContext* ctx) override {
+    Tensor handle_tensor;
+    AllocatorAttributes attr;
+    OP_REQUIRES_OK(ctx, ctx->allocate_temp(DT_RESOURCE, TensorShape({}),
+                                           &handle_tensor, attr));
+    handle_tensor.scalar<ResourceHandle>()() =
+        ResourceHandle::MakeRefCountingHandle(
+            new SimpleHashTableResource<K, V>(), ctx->device()->name(),
+            /*dtypes_and_shapes=*/{}, ctx->stack_trace());
+    ctx->set_output(0, handle_tensor);
+  }
+
+ private:
+  // Just to be safe, avoid accidentally copying the kernel.
+  SimpleHashTableCreateOpKernel(const SimpleHashTableCreateOpKernel&) = delete;
+  void operator=(const SimpleHashTableCreateOpKernel&) = delete;
+};
+
+// GetResource retrieves a Resource using a handle from the first
+// input in "ctx" and saves it in "resource" without increasing
+// the reference count for that resource.
+template <class K, class V>
+Status GetResource(OpKernelContext* ctx,
+                   SimpleHashTableResource<K, V>** resource) {
+  const Tensor& handle_tensor = ctx->input(0);
+  const ResourceHandle& handle = handle_tensor.scalar<ResourceHandle>()();
+  typedef SimpleHashTableResource<K, V> resource_type;
+  TF_ASSIGN_OR_RETURN(*resource, handle.GetResource<resource_type>());
+  return OkStatus();
+}
+
+template <class K, class V>
+class SimpleHashTableFindOpKernel : public OpKernel {
+ public:
+  explicit SimpleHashTableFindOpKernel(OpKernelConstruction* ctx)
+      : OpKernel(ctx) {}
+
+  void Compute(OpKernelContext* ctx) override {
+    DataTypeVector expected_inputs = {DT_RESOURCE, DataTypeToEnum<K>::v(),
+                                      DataTypeToEnum<V>::v()};
+    DataTypeVector expected_outputs = {DataTypeToEnum<V>::v()};
+    OP_REQUIRES_OK(ctx, ctx->MatchSignature(expected_inputs, expected_outputs));
+    SimpleHashTableResource<K, V>* resource;
+    OP_REQUIRES_OK(ctx, GetResource(ctx, &resource));
+    // Note that ctx->input(0) is the Resource handle
+    const Tensor& key = ctx->input(1);
+    const Tensor& default_value = ctx->input(2);
+    TensorShape output_shape = default_value.shape();
+    Tensor* out;
+    OP_REQUIRES_OK(ctx, ctx->allocate_output("value", output_shape, &out));
+    OP_REQUIRES_OK(ctx, resource->Find(key, out, default_value));
+  }
+};
+
+template <class K, class V>
+class SimpleHashTableInsertOpKernel : public OpKernel {
+ public:
+  explicit SimpleHashTableInsertOpKernel(OpKernelConstruction* ctx)
+      : OpKernel(ctx) {}
+
+  void Compute(OpKernelContext* ctx) override {
+    DataTypeVector expected_inputs = {DT_RESOURCE, DataTypeToEnum<K>::v(),
+                                      DataTypeToEnum<V>::v()};
+    OP_REQUIRES_OK(ctx, ctx->MatchSignature(expected_inputs, {}));
+    SimpleHashTableResource<K, V>* resource;
+    OP_REQUIRES_OK(ctx, GetResource(ctx, &resource));
+    // Note that ctx->input(0) is the Resource handle
+    const Tensor& key = ctx->input(1);
+    const Tensor& value = ctx->input(2);
+    OP_REQUIRES_OK(ctx, resource->Insert(key, value));
+  }
+};
+
+template <class K, class V>
+class SimpleHashTableRemoveOpKernel : public OpKernel {
+ public:
+  explicit SimpleHashTableRemoveOpKernel(OpKernelConstruction* ctx)
+      : OpKernel(ctx) {}
+
+  void Compute(OpKernelContext* ctx) override {
+    DataTypeVector expected_inputs = {DT_RESOURCE, DataTypeToEnum<K>::v()};
+    OP_REQUIRES_OK(ctx, ctx->MatchSignature(expected_inputs, {}));
+    SimpleHashTableResource<K, V>* resource;
+    OP_REQUIRES_OK(ctx, GetResource(ctx, &resource));
+    // Note that ctx->input(0) is the Resource handle
+    const Tensor& key = ctx->input(1);
+    OP_REQUIRES_OK(ctx, resource->Remove(key));
+  }
+};
+
+template <class K, class V>
+class SimpleHashTableExportOpKernel : public OpKernel {
+ public:
+  explicit SimpleHashTableExportOpKernel(OpKernelConstruction* ctx)
+      : OpKernel(ctx) {}
+
+  void Compute(OpKernelContext* ctx) override {
+    DataTypeVector expected_inputs = {DT_RESOURCE};
+    DataTypeVector expected_outputs = {DataTypeToEnum<K>::v(),
+                                       DataTypeToEnum<V>::v()};
+    OP_REQUIRES_OK(ctx, ctx->MatchSignature(expected_inputs, expected_outputs));
+    SimpleHashTableResource<K, V>* resource;
+    OP_REQUIRES_OK(ctx, GetResource(ctx, &resource));
+    OP_REQUIRES_OK(ctx, resource->Export(ctx));
+  }
+};
+
+template <class K, class V>
+class SimpleHashTableImportOpKernel : public OpKernel {
+ public:
+  explicit SimpleHashTableImportOpKernel(OpKernelConstruction* ctx)
+      : OpKernel(ctx) {}
+
+  void Compute(OpKernelContext* ctx) override {
+    SimpleHashTableResource<K, V>* resource;
+    OP_REQUIRES_OK(ctx, GetResource(ctx, &resource));
+    DataTypeVector expected_inputs = {DT_RESOURCE, DataTypeToEnum<K>::v(),
+                                      DataTypeToEnum<V>::v()};
+    OP_REQUIRES_OK(ctx, ctx->MatchSignature(expected_inputs, {}));
+    const Tensor& keys = ctx->input(1);
+    const Tensor& values = ctx->input(2);
+    OP_REQUIRES(
+        ctx, keys.shape() == values.shape(),
+        errors::InvalidArgument("Shapes of keys and values are not the same: ",
+                                keys.shape().DebugString(), " for keys, ",
+                                values.shape().DebugString(), "for values."));
+
+    int memory_used_before = 0;
+    if (ctx->track_allocations()) {
+      memory_used_before = resource->MemoryUsed();
+    }
+    OP_REQUIRES_OK(ctx, resource->Import(keys, values));
+    if (ctx->track_allocations()) {
+      ctx->record_persistent_memory_allocation(resource->MemoryUsed() -
+                                               memory_used_before);
+    }
+  }
+};
+
+// The "Name" used by REGISTER_KERNEL_BUILDER is defined by REGISTER_OP,
+// see simple_hash_table_op.cc.
+#define REGISTER_KERNEL(key_dtype, value_dtype)               \
+  REGISTER_KERNEL_BUILDER(                                    \
+      Name("Examples>SimpleHashTableCreate")                  \
+          .Device(DEVICE_CPU)                                 \
+          .TypeConstraint<key_dtype>("key_dtype")             \
+          .TypeConstraint<value_dtype>("value_dtype"),        \
+      SimpleHashTableCreateOpKernel<key_dtype, value_dtype>); \
+  REGISTER_KERNEL_BUILDER(                                    \
+      Name("Examples>SimpleHashTableFind")                    \
+          .Device(DEVICE_CPU)                                 \
+          .TypeConstraint<key_dtype>("key_dtype")             \
+          .TypeConstraint<value_dtype>("value_dtype"),        \
+      SimpleHashTableFindOpKernel<key_dtype, value_dtype>);   \
+  REGISTER_KERNEL_BUILDER(                                    \
+      Name("Examples>SimpleHashTableInsert")                  \
+          .Device(DEVICE_CPU)                                 \
+          .TypeConstraint<key_dtype>("key_dtype")             \
+          .TypeConstraint<value_dtype>("value_dtype"),        \
+      SimpleHashTableInsertOpKernel<key_dtype, value_dtype>)  \
+  REGISTER_KERNEL_BUILDER(                                    \
+      Name("Examples>SimpleHashTableRemove")                  \
+          .Device(DEVICE_CPU)                                 \
+          .TypeConstraint<key_dtype>("key_dtype")             \
+          .TypeConstraint<value_dtype>("value_dtype"),        \
+      SimpleHashTableRemoveOpKernel<key_dtype, value_dtype>)  \
+  REGISTER_KERNEL_BUILDER(                                    \
+      Name("Examples>SimpleHashTableExport")                  \
+          .Device(DEVICE_CPU)                                 \
+          .TypeConstraint<key_dtype>("key_dtype")             \
+          .TypeConstraint<value_dtype>("value_dtype"),        \
+      SimpleHashTableExportOpKernel<key_dtype, value_dtype>)  \
+  REGISTER_KERNEL_BUILDER(                                    \
+      Name("Examples>SimpleHashTableImport")                  \
+          .Device(DEVICE_CPU)                                 \
+          .TypeConstraint<key_dtype>("key_dtype")             \
+          .TypeConstraint<value_dtype>("value_dtype"),        \
+      SimpleHashTableImportOpKernel<key_dtype, value_dtype>);
+
+REGISTER_KERNEL(int32, double);
+REGISTER_KERNEL(int32, float);
+REGISTER_KERNEL(int32, int32);
+REGISTER_KERNEL(int32, tstring);
+REGISTER_KERNEL(int64_t, double);
+REGISTER_KERNEL(int64_t, float);
+REGISTER_KERNEL(int64_t, int32);
+REGISTER_KERNEL(int64_t, int64_t);
+REGISTER_KERNEL(int64_t, tstring);
+REGISTER_KERNEL(tstring, bool);
+REGISTER_KERNEL(tstring, double);
+REGISTER_KERNEL(tstring, float);
+REGISTER_KERNEL(tstring, int32);
+REGISTER_KERNEL(tstring, int64_t);
+REGISTER_KERNEL(tstring, tstring);
+#undef REGISTER_KERNEL
+
+}  // namespace custom_op_examples
+}  // namespace tensorflow

tensorflow/examples/custom_ops_doc/simple_hash_table/simple_hash_table_op.cc

@@ -0,0 +1,177 @@
+/* Copyright 2021 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 "xla/tsl/platform/errors.h"
+#include "tensorflow/core/framework/op.h"
+#include "tensorflow/core/framework/shape_inference.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/platform/types.h"
+
+// Please use the appropriate namespace for your project
+namespace tensorflow {
+namespace custom_op_examples {
+
+using ::tensorflow::shape_inference::DimensionHandle;
+using ::tensorflow::shape_inference::InferenceContext;
+using ::tensorflow::shape_inference::ShapeAndType;
+using ::tensorflow::shape_inference::ShapeHandle;
+
+Status ScalarOutput(InferenceContext* c) {
+  c->set_output(0, c->Scalar());
+  return OkStatus();
+}
+
+Status TwoScalarInputs(InferenceContext* c) {
+  ShapeHandle handle;
+  TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 0, &handle));
+  TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &handle));
+  return OkStatus();
+}
+
+Status TwoScalarInputsScalarOutput(InferenceContext* c) {
+  ShapeHandle handle;
+  TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 0, &handle));
+  TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &handle));
+  return ScalarOutput(c);
+}
+
+Status ThreeScalarInputs(InferenceContext* c) {
+  ShapeHandle handle;
+  TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 0, &handle));
+  TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &handle));
+  TF_RETURN_IF_ERROR(c->WithRank(c->input(2), 0, &handle));
+  return OkStatus();
+}
+
+Status ThreeScalarInputsScalarOutput(InferenceContext* c) {
+  ShapeHandle handle;
+  TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 0, &handle));
+  TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &handle));
+  TF_RETURN_IF_ERROR(c->WithRank(c->input(2), 0, &handle));
+  return ScalarOutput(c);
+}
+
+Status ValidateTableType(InferenceContext* c,
+                         const ShapeAndType& key_shape_and_type,
+                         const string& key_dtype_attr,
+                         const ShapeAndType& value_shape_and_type,
+                         const string& value_dtype_attr) {
+  DataType key_dtype;
+  TF_RETURN_IF_ERROR(c->GetAttr(key_dtype_attr, &key_dtype));
+  if (key_shape_and_type.dtype != key_dtype) {
+    return errors::InvalidArgument(
+        "Trying to read value with wrong dtype. "
+        "Expected ",
+        DataTypeString(key_shape_and_type.dtype), " got ",
+        DataTypeString(key_dtype));
+  }
+  DataType value_dtype;
+  TF_RETURN_IF_ERROR(c->GetAttr(value_dtype_attr, &value_dtype));
+  if (value_shape_and_type.dtype != value_dtype) {
+    return errors::InvalidArgument(
+        "Trying to read value with wrong dtype. "
+        "Expected ",
+        DataTypeString(value_shape_and_type.dtype), " got ",
+        DataTypeString(value_dtype));
+  }
+  return OkStatus();
+}
+
+Status ExportShapeFunction(InferenceContext* c) {
+  ShapeHandle handle;
+  TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 0, &handle));
+  auto* handle_data = c->input_handle_shapes_and_types(0);
+  if (handle_data != nullptr && handle_data->size() == 2) {
+    const ShapeAndType& key_shape_and_type = (*handle_data)[0];
+    const ShapeAndType& value_shape_and_type = (*handle_data)[1];
+    TF_RETURN_IF_ERROR(ValidateTableType(c, key_shape_and_type,
+                                         /*key_dtype_attr*/ "key_dtype",
+                                         value_shape_and_type,
+                                         /*value_dtype_attr*/ "value_dtype"));
+  }
+  // Different lookup tables have different output shapes.
+  c->set_output(0, c->UnknownShape());
+  c->set_output(1, c->UnknownShape());
+  return OkStatus();
+}
+
+Status ImportShapeFunction(InferenceContext* c) {
+  ShapeHandle handle;
+  TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 0, &handle));
+
+  ShapeHandle keys;
+  TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 1, &keys));
+  DimensionHandle unused;
+  TF_RETURN_IF_ERROR(
+      c->Merge(c->Dim(keys, 0), c->Dim(c->input(2), 0), &unused));
+  return OkStatus();
+}
+
+// Note that if an op has any Input or Output of type "resource", it
+// is automatically marked as stateful so there is no need to explicitly
+// use "SetIsStateful()".
+// (See FinalizeInputOrOutput in core/framework/op_def_builder.cc.)
+
+REGISTER_OP("Examples>SimpleHashTableCreate")
+    .Output("output: resource")
+    .Attr("key_dtype: type")
+    .Attr("value_dtype: type")
+    .SetShapeFn(ScalarOutput);
+
+REGISTER_OP("Examples>SimpleHashTableFind")
+    .Input("resource_handle: resource")
+    .Input("key: key_dtype")
+    .Input("default_value: value_dtype")
+    .Output("value: value_dtype")
+    .Attr("key_dtype: type")
+    .Attr("value_dtype: type")
+    .SetShapeFn(ThreeScalarInputsScalarOutput);
+
+REGISTER_OP("Examples>SimpleHashTableInsert")
+    .Input("resource_handle: resource")
+    .Input("key: key_dtype")
+    .Input("value: value_dtype")
+    .Attr("key_dtype: type")
+    .Attr("value_dtype: type")
+    .SetShapeFn(ThreeScalarInputs);
+
+REGISTER_OP("Examples>SimpleHashTableRemove")
+    .Input("resource_handle: resource")
+    .Input("key: key_dtype")
+    .Attr("key_dtype: type")
+    .Attr("value_dtype: type")
+    .SetShapeFn(TwoScalarInputs);
+
+REGISTER_OP("Examples>SimpleHashTableExport")
+    .Input("table_handle: resource")
+    .Output("keys: key_dtype")
+    .Output("values: value_dtype")
+    .Attr("key_dtype: type")
+    .Attr("value_dtype: type")
+    .SetShapeFn(ExportShapeFunction);
+
+REGISTER_OP("Examples>SimpleHashTableImport")
+    .Input("table_handle: resource")
+    .Input("keys: key_dtype")
+    .Input("values: value_dtype")
+    .Attr("key_dtype: type")
+    .Attr("value_dtype: type")
+    .SetShapeFn(ImportShapeFunction);
+
+}  // namespace custom_op_examples
+}  // namespace tensorflow

tensorflow/examples/custom_ops_doc/simple_hash_table/simple_hash_table_op.py

@@ -0,0 +1,21 @@
+# Copyright 2021 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.
+# ==============================================================================
+"""Load the simple_hash_table_op kernel."""
+
+import tensorflow as tf
+from tensorflow.python.platform import resource_loader
+
+gen_simple_hash_table_op = tf.load_op_library(
+    resource_loader.get_path_to_datafile("simple_hash_table_kernel.so"))

tensorflow/examples/custom_ops_doc/simple_hash_table/simple_hash_table_test.py

@@ -0,0 +1,125 @@
+# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for simple_hash_table."""
+
+import os.path
+import tempfile
+
+from absl.testing import parameterized
+import tensorflow as tf
+
+from tensorflow.examples.custom_ops_doc.simple_hash_table import simple_hash_table
+from tensorflow.python.eager import def_function
+# This pylint disable is only needed for internal google users
+from tensorflow.python.framework import test_util  # pylint: disable=g-direct-tensorflow-import
+
+
+class SimpleHashTableTest(tf.test.TestCase, parameterized.TestCase):
+
+  # Helper function using "create, find, insert, find, remove, find
+  def _use_table(self, key_dtype, value_dtype):
+    hash_table = simple_hash_table.SimpleHashTable(key_dtype, value_dtype, 111)
+    result1 = hash_table.find(1, -999)
+    hash_table.insert(1, 100)
+    result2 = hash_table.find(1, -999)
+    hash_table.remove(1)
+    result3 = hash_table.find(1, -999)
+    results = tf.stack((result1, result2, result3))
+    return results  # expect [-999, 100, -999]
+
+  # Test of "create, find, insert, find" in eager mode.
+  @parameterized.named_parameters(('int32_float', tf.int32, float),
+                                  ('int64_int32', tf.int64, tf.int32))
+  def test_find_insert_find_eager(self, key_dtype, value_dtype):
+    results = self._use_table(key_dtype, value_dtype)
+    self.assertAllClose(results, [-999, 100, -999])
+
+  # Test of "create, find, insert, find" in a tf.function. Note that the
+  # creation and use of the ref-counted resource occurs inside a single
+  # self.evaluate.
+  @parameterized.named_parameters(('int32_float', tf.int32, float),
+                                  ('int64_int32', tf.int64, tf.int32))
+  def test_find_insert_find_tf_function(self, key_dtype, value_dtype):
+    results = def_function.function(
+        lambda: self._use_table(key_dtype, value_dtype))
+    self.assertAllClose(self.evaluate(results), [-999.0, 100.0, -999.0])
+
+  # strings for key and value
+  def test_find_insert_find_strings_eager(self):
+    default = 'Default'
+    foo = 'Foo'
+    bar = 'Bar'
+    hash_table = simple_hash_table.SimpleHashTable(tf.string, tf.string,
+                                                   default)
+    result1 = hash_table.find(foo, default)
+    self.assertEqual(result1, default)
+    hash_table.insert(foo, bar)
+    result2 = hash_table.find(foo, default)
+    self.assertEqual(result2, bar)
+
+  def test_export(self):
+    table = simple_hash_table.SimpleHashTable(
+        tf.int64, tf.int64, default_value=-1)
+    table.insert(1, 100)
+    table.insert(2, 200)
+    table.insert(3, 300)
+    keys, values = self.evaluate(table.export())
+    self.assertAllEqual(sorted(keys), [1, 2, 3])
+    self.assertAllEqual(sorted(values), [100, 200, 300])
+
+  def test_import(self):
+    table = simple_hash_table.SimpleHashTable(
+        tf.int64, tf.int64, default_value=-1)
+    keys = tf.constant([1, 2, 3], dtype=tf.int64)
+    values = tf.constant([100, 200, 300], dtype=tf.int64)
+    table.do_import(keys, values)
+    self.assertEqual(table.find(1), 100)
+    self.assertEqual(table.find(2), 200)
+    self.assertEqual(table.find(3), 300)
+    self.assertEqual(table.find(9), -1)
+
+  @test_util.run_v2_only
+  def testSavedModelSaveRestore(self):
+    save_dir = os.path.join(self.get_temp_dir(), 'save_restore')
+    save_path = os.path.join(tempfile.mkdtemp(prefix=save_dir), 'hash')
+
+    # TODO(b/203097231) is there an alternative that is not __internal__?
+    root = tf.__internal__.tracking.AutoTrackable()
+
+    default_value = -1
+    root.table = simple_hash_table.SimpleHashTable(
+        tf.int64, tf.int64, default_value=default_value)
+
+    @def_function.function(input_signature=[tf.TensorSpec((), tf.int64)])
+    def lookup(key):
+      return root.table.find(key)
+
+    root.lookup = lookup
+
+    root.table.insert(1, 100)
+    root.table.insert(2, 200)
+    root.table.insert(3, 300)
+    self.assertEqual(root.lookup(2), 200)
+    self.assertAllEqual(3, len(self.evaluate(root.table.export()[0])))
+    tf.saved_model.save(root, save_path)
+
+    del root
+    loaded = tf.saved_model.load(save_path)
+    self.assertEqual(loaded.lookup(2), 200)
+    self.assertEqual(loaded.lookup(10), -1)
+
+
+if __name__ == '__main__':
+  tf.test.main()

tensorflow/examples/custom_ops_doc/sleep/BUILD

@@ -0,0 +1,58 @@
+# (non-blocking) op using AsyncOpKernel
+
+load("//tensorflow:strict.default.bzl", "py_strict_binary", "py_strict_library")
+load("//tensorflow:tensorflow.bzl", "tf_custom_op_library")
+load("//tensorflow:tensorflow.default.bzl", "tf_py_test")
+
+licenses(["notice"])
+
+tf_custom_op_library(
+    name = "sleep_kernel.so",
+    srcs = [
+        "sleep_kernel.cc",
+        "sleep_op.cc",
+    ],
+    deps = [
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/time",
+    ],
+)
+
+py_strict_library(
+    name = "sleep_op",
+    srcs = ["sleep_op.py"],
+    data = ["sleep_kernel.so"],
+    srcs_version = "PY3",
+)
+
+py_strict_binary(
+    name = "sleep_bin",
+    srcs = ["sleep_bin.py"],
+    srcs_version = "PY3",
+    deps = [
+        ":sleep_op",
+        "//tensorflow:tensorflow_py",
+        "//third_party/py/numpy",
+        "@absl_py//absl:app",
+    ],
+)
+
+tf_py_test(
+    name = "sleep_test",
+    size = "medium",  # This test blocks using sleep,
+    timeout = "short",  # but it still runs quickly.
+    srcs = ["sleep_test.py"],
+    python_version = "PY3",
+    srcs_version = "PY3",
+    tags = [
+        "no_mac",  # TODO(b/216321151): Re-enable this test.
+    ],
+    deps = [
+        ":sleep_op",
+        "//tensorflow:tensorflow_py",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:test_lib",
+        "//third_party/py/numpy",
+    ],
+)

tensorflow/examples/custom_ops_doc/sleep/README.md

@@ -0,0 +1,640 @@
+<!-- LINT.IfChange -->
+# Create an asynchronous sleep op
+
+This guide provides an end-to-end example for an asynchronous custom op. The
+example implements an asynchronous sleep op, and contrasts the implementation
+with a synchronous sleep op.
+
+Asynchronous ops allow other ops to execute while the asynchronous op waits.
+Unlike synchronous ops, an asynchronous op does not block other ops as it waits.
+
+The content on this page assumes familiarity with the high-level process for
+adding custom ops to TensorFlow. For additional context,
+read the
+[OSS guide on creating custom ops](https://www.tensorflow.org/guide/create_op).
+
+## Background information on asynchronous ops
+
+Asynchronous ops are recommended for cases where the number of ops that may be
+waiting at a given time is significantly larger than the desired number of
+threads. In general, you should use asynchronous ops for most cases where an op
+waits, though for specific cases (e.g. short waits) the overhead and complexity
+of an asynchronous op may outweigh the benefits.
+
+You can also use event and queuing based techniques, rather than using threads.
+For cases where the op interfaces to something that already uses a callback, an
+asynchronous op implementation is more straightforward than the alternatives.
+
+Asynchronous ops are derived from
+[`AsyncOpKernel`](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/framework/op_kernel.h#L225)
+and use the `ComputeAsync()` method to override the default `Compute()` method
+used by synchronous ops
+([`OpKernel`](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/framework/op_kernel.h#L104)).
+
+An op can delegate the task of producing a result to another function, method,
+or closure using `ComputeAsync()`. The op can then return before all the work
+has completed. For example, the op can schedule a closure to run in another
+thread and then return immediately.
+
+### The `done` callback function
+
+The `ComputeAsync()` method contains a `done` parameter, which is passed to
+other functions as a callback function. Once the other function sets the output
+or completes the work, it calls `done` to notify that the op has finished.
+
+The `done` function must be called exactly once in every execution path. Any
+paths in `ComputeAsync()` that return early, whether due to error handling or
+cases where results are produced quickly, must call `done`. Similarly, if the
+function that receives the `done` callback has any execution paths that return
+early, the function must call `done` in these paths in addition to calling
+`done` before it returns normally.
+
+Once `done` is triggered, any ops that depend on the output(s) of this op can
+execute.
+
+<!-- test_snippets_in_readme skip -->
+```c++
+class MyAsyncOp : public AsyncOpKernel {
+ // …
+  void ComputeAsync(OpKernelContext* ctx, DoneCallback done) override {
+    // …
+    OP_REQUIRES_OK_ASYNC(ctx, ErrorChecking(), done);
+    // …
+    // ScheduleOrRequest is pseudocode for scheduling or requesting work elsewhere
+    ScheduleOrRequest(delegate, done /* , … */ );
+    // return quickly without waiting for the result
+  }
+
+}
+void delegate(DoneCallback done /* , … */ ) {
+  // set output of op and/or do work of op
+  // …
+  done();
+}
+```
+
+#### Alternatives for parallelizing computation
+
+While this example implements an op that waits, there are other cases where you
+may seek an alternative to executing computationally-intensive ops in parallel.
+For example:
+
+*   Ops in `tf.function` graphs can be executed in parallel by multiple threads
+    from the inter-op thread pool. You can use synchronous ops when there are at
+    least as many threads as ops that can run at the same time. Even if there
+    are fewer threads than ops, this may still be a sufficient implementation.
+*   `tf.data` can be used for cases that create an input pipe.
+    `tf.data.dataset.map()` can be used with a synchronous custom op, where the
+    `num_parallel_calls` parameter is used to run that op in parallel with a
+    thread pool provided by `tf.data`. For more information, see
+    [Better performance with the tf.data API](https://www.tensorflow.org/guide/data_performance)
+    and
+    [Preprocessing data](https://www.tensorflow.org/guide/data#preprocessing_data).
+*   [tf.map_fn](https://www.tensorflow.org/api_docs/python/tf/map_fn) and
+    [tf.while_loop](https://www.tensorflow.org/api_docs/python/tf/while_loop)
+    can be used with a synchronous custom op and the `parallel_iterations`
+    parameter to run that op in parallel using multiple threads.
+
+## Creating an asynchronous sleep op
+
+This example demonstrates how you can create an asynchronous sleep op,
+`sleep_op`, that waits for a specified amount of time while letting other ops
+run. While this example does not wait for any particular purpose, it illustrates
+the pattern for ops that delay or poll while waiting for something, such as
+another op's result.
+
+For contrast, the example also includes a synchronous version of the same sleep
+op.
+
+The following example waits for one second with the synchronous and asynchronous
+versions:
+
+<!-- test_snippets_in_readme skip -->
+```python
+sleep_op.SyncSleep(1.0)
+# tf.Tensor(1.0, shape=(), dtype=float32)
+sleep_op.AsyncSleep(1.0)
+# tf.Tensor(0.999892, shape=(), dtype=float32)
+```
+
+The synchronous version (`sleep_op.SyncSleep`) simply blocks time by calling
+`absl::SleepFor` from
+[clock.h](https://github.com/abseil/abseil-cpp/blob/8c6e53ef3adb1227fffa442c50349dab134a54bc/absl/time/clock.h)
+(similar to `sleep`) for the amount of time specified in the input (1 second),
+and returns that delay value in the output.
+
+The asynchronous version (`sleep_op.AsyncSleep`) uses the input (1 second) and
+the current time to compute the wake-up time, which is the point in time when
+the function stops waiting. The op schedules a function that receives this
+wake-up time to run in a threadpool, and returns immediately.
+
+The function either begins running immediately after it is scheduled or after
+blocking for some time. If the current time is before the wake-up time, the
+function computes the difference and blocks time using `absl::SleepFor`.
+Otherwise, if the current time is equal to or after the wake-up time, it does
+not call `absl::SleepFor`. The function sets the output to the time specified
+for sleep, or 0 if it does not call `absl::SleepFor`, and calls the `done`
+callback.
+
+In the example above, the asynchronous function called sleep for 0.999892
+seconds after being blocked for 0.000108 seconds waiting for the function to
+run. This waits for a total time of 1 second, as specified in the input.
+
+This example contains C++ and Python code snippets to illustrate the code flow.
+These snippets may be missing namespace declarations, imports, and test cases.
+
+### Step 1 - Define the op interface
+
+Define the op interface and register it using the `REGISTER_OP` macro.
+
+```
+REGISTER_OP("Examples>AsyncSleep")
+    .Input("delay: float")
+    .Output("output: float")
+    .SetShapeFn([](tensorflow::shape_inference::InferenceContext* c) {
+      tensorflow::shape_inference::ShapeHandle handle;
+      TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 0, &handle));
+      return (ScalarOutput(c));
+    })
+    .Doc(R"doc(
+Pause for `delay` seconds (which need not be an integer).
+
+This is an asynchronous (non-blocking) version of sleep. It is intended to
+be an example of how to implement ops that do I/O or that block on other ops.
+
+delay: tf.Tensor which is a scalar of type float.
+
+Returns the time spent in blocking sleep (which may be less that `delay` or
+zero if other ops run while this is waiting asynchronously).
+)doc");
+```
+```
+REGISTER_OP("Examples>SyncSleep")
+    .Input("delay: float")
+    .Output("output: float")
+    .SetShapeFn([](tensorflow::shape_inference::InferenceContext* c) {
+      tensorflow::shape_inference::ShapeHandle handle;
+      TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 0, &handle));
+      return (ScalarOutput(c));
+    })
+    .Doc(R"doc(
+Pause for `delay` seconds (which need not be an integer).
+
+This is a synchronous (blocking) version of sleep. It's purpose is
+to be contrasted with Examples>AsyncSleep.
+
+delay: tf.Tensor which is a scalar of type float.
+
+Returns `delay`.
+)doc");
+```
+
+The op registers two examples, an asynchronous op (`Examples>AsyncSleep`) and a
+synchronous op (`Examples>SyncSleep`). Both examples accept `delay` as an input,
+which is a scalar of type float. The `AsyncSleep` example returns the time spent
+in blocking sleep (which may be less than `delay` if other ops run while the op
+is waiting asynchronously), while the `SyncSleep` example simply returns the
+value of `delay`.
+
+### Step 2 - Register the op implementation (kernel)
+
+The C++ kernel in
+`sleep_kernel.cc`
+implements both a synchronous (`SyncSleepOp`) and asynchronous (`AsyncSleepOp`)
+sleep op.
+
+<!-- test_snippets_in_readme skip -->
+```c++
+REGISTER_KERNEL_BUILDER(
+    Name("Examples>AsyncSleep").Device(::tensorflow::DEVICE_CPU), AsyncSleepOp)
+REGISTER_KERNEL_BUILDER(
+    Name("Examples>SyncSleep").Device(::tensorflow::DEVICE_CPU), SyncSleepOp)
+```
+
+### Step 3 - Implement the op kernel(s)
+
+In the `sleep_kernel.cc` op kernel, create two classes: one derived from
+`OpKernel` and another derived from `AsyncOpKernel`.
+
+The `OpKernel` class will be familiar if you have followed the other custom op
+examples in this series. It implements a `Compute()` method, which is used for
+the synchronous sleep op. The new `AsyncOpKernel` class uses `ComputeAsync()`,
+which is used for the asynchronous sleep op.
+
+```
+class AsyncSleepOp : public AsyncOpKernel {
+ public:
+  explicit AsyncSleepOp(OpKernelConstruction* ctx) : AsyncOpKernel(ctx) {}
+  AsyncSleepOp(const AsyncSleepOp& other) = delete;
+  AsyncSleepOp& operator=(const AsyncSleepOp& other) = delete;
+  ~AsyncSleepOp() override = default;
+
+  // Implementations of ComputeAsync() must ensure that `done` is (eventually)
+  // called exactly once to signal the completion of the computation. The
+  // implementation of ComputeAsync() must not block on the execution of another
+  // OpKernel. `done` may be called by the current thread, or by another thread.
+  // `context` is guaranteed to stay alive until the `done` callback starts.
+  // For example, use OP_REQUIRES_ASYNC which takes the `done` parameter
+  // as an input and calls `done` for the case of exiting early with an error
+  // (instead of OP_REQUIRES).
+  //
+  // Since it is possible that the unblocking kernel may never run (due to an
+  // error or cancellation), in most cases the AsyncOpKernel should implement
+  // cancellation support via `context->cancellation_manager()`.
+  // TODO (schwartzedward): should this use cancellation support?
+  //
+  // WARNING: As soon as the `done` callback starts, `context` and `this` may be
+  // deleted. No code depending on these objects should execute after the call
+  // to `done`.
+  void ComputeAsync(OpKernelContext* ctx, DoneCallback done) override {
+    const auto& delay_tensor = ctx->input(0);
+    OP_REQUIRES_ASYNC(
+        ctx, ::tensorflow::TensorShapeUtils::IsScalar(delay_tensor.shape()),
+        InvalidArgument("Input `delay` must be a scalar."),
+        done);  // Important: call `done` in every execution path
+    const float delay = delay_tensor.flat<float>()(0);
+    OP_REQUIRES_ASYNC(ctx, delay >= 0.0,
+                      InvalidArgument("Input `delay` must be non-negative."),
+                      done);  // Important: call `done` in every execution path
+    auto thread_pool = ctx->device()->tensorflow_cpu_worker_threads()->workers;
+    OP_REQUIRES_ASYNC(ctx, thread_pool != nullptr,
+                      Internal("No thread_pool found."),
+                      done);  // Important: call `done` in every execution path
+
+    Tensor* output_tensor = nullptr;
+    OP_REQUIRES_OK_ASYNC(
+        ctx, ctx->allocate_output(0, delay_tensor.shape(), &output_tensor),
+        done);  // Important: call `done` in every execution path
+
+    absl::Time now = absl::Now();
+    absl::Time when = now + absl::Seconds(delay);
+    VLOG(1) << "BEFORE ASYNC SLEEP " << ctx->op_kernel().name() << " now "
+            << now << " when " << when;
+    thread_pool->Schedule([this, output_tensor, when, done] {
+      this->sleeper(output_tensor, when, done);
+    });
+    // Note that `done` is normally called by sleeper(), it is not normally
+    // called by this function.
+  }
+
+ private:
+  void sleeper(Tensor* output_tensor, absl::Time when, DoneCallback done) {
+    absl::Time now = absl::Now();
+    int64_t delay_us = 0;
+    if (now < when) {
+      delay_us = absl::ToInt64Microseconds(when - now);
+      VLOG(1) << "MIDDLE ASYNC SLEEP " << delay_us;
+      absl::SleepFor(when - now);
+      VLOG(1) << "AFTER ASYNC SLEEP " << delay_us;
+    } else {
+      VLOG(1) << "MIDDLE/AFTER ASYNC SKIP SLEEP";
+    }
+    auto output = output_tensor->template flat<float>();
+    output(0) = static_cast<float>(delay_us) / 1000000.0;
+    done();  // Important: call `done` in every execution path
+  }
+};
+```
+
+The implementation overrides `ComputeAsync`, which has a `done` callback
+argument. For error checking, it uses `OP_REQUIRES_ASYNC` and
+`OP_REQUIRES_OK_ASYNC`, which calls `done` if returning with an error.
+
+Note: For asynchronous ops, always use `OP_REQUIRES_ASYNC` and
+`OP_REQUIRES_OK_ASYNC`. Asynchronous ops cannot use `OP_REQUIRES` or
+`OP_REQUIRES_OK`.
+
+The wake-up time for the op is computed using `absl::Now` from
+[clock.h](https://github.com/abseil/abseil-cpp/blob/8c6e53ef3adb1227fffa442c50349dab134a54bc/absl/time/clock.h).
+The sleeper helper method receives the wake-up time and `done` callback, and
+runs in the `tensorflow_cpu_worker_threads` thread pool.
+
+For normal operations, `done` is called by the sleeper method, not by
+`ComputeAsync`.
+
+```
+  void ComputeAsync(OpKernelContext* ctx, DoneCallback done) override {
+    const auto& delay_tensor = ctx->input(0);
+    OP_REQUIRES_ASYNC(
+        ctx, ::tensorflow::TensorShapeUtils::IsScalar(delay_tensor.shape()),
+        InvalidArgument("Input `delay` must be a scalar."),
+        done);  // Important: call `done` in every execution path
+    const float delay = delay_tensor.flat<float>()(0);
+    OP_REQUIRES_ASYNC(ctx, delay >= 0.0,
+                      InvalidArgument("Input `delay` must be non-negative."),
+                      done);  // Important: call `done` in every execution path
+    auto thread_pool = ctx->device()->tensorflow_cpu_worker_threads()->workers;
+    OP_REQUIRES_ASYNC(ctx, thread_pool != nullptr,
+                      Internal("No thread_pool found."),
+                      done);  // Important: call `done` in every execution path
+
+    Tensor* output_tensor = nullptr;
+    OP_REQUIRES_OK_ASYNC(
+        ctx, ctx->allocate_output(0, delay_tensor.shape(), &output_tensor),
+        done);  // Important: call `done` in every execution path
+
+    absl::Time now = absl::Now();
+    absl::Time when = now + absl::Seconds(delay);
+    VLOG(1) << "BEFORE ASYNC SLEEP " << ctx->op_kernel().name() << " now "
+            << now << " when " << when;
+    thread_pool->Schedule([this, output_tensor, when, done] {
+      this->sleeper(output_tensor, when, done);
+    });
+    // Note that `done` is normally called by sleeper(), it is not normally
+    // called by this function.
+  }
+```
+
+If the current time is before the specified wake-up time, the `sleeper` helper
+method sleeps for the difference between the current time and wake-up time. It
+sets the output and calls the `done` callback, which notifies the op that the
+output is set.
+
+Calling `done` also invalidates `ctx`, so its contents cannot be used after
+calling `done`.
+
+```
+  void sleeper(Tensor* output_tensor, absl::Time when, DoneCallback done) {
+    absl::Time now = absl::Now();
+    int64_t delay_us = 0;
+    if (now < when) {
+      delay_us = absl::ToInt64Microseconds(when - now);
+      VLOG(1) << "MIDDLE ASYNC SLEEP " << delay_us;
+      absl::SleepFor(when - now);
+      VLOG(1) << "AFTER ASYNC SLEEP " << delay_us;
+    } else {
+      VLOG(1) << "MIDDLE/AFTER ASYNC SKIP SLEEP";
+    }
+    auto output = output_tensor->template flat<float>();
+    output(0) = static_cast<float>(delay_us) / 1000000.0;
+    done();  // Important: call `done` in every execution path
+  }
+```
+
+#### Compile the op
+
+Compile the C++ op to create a kernel library and Python wrapper that enables
+you to use the op with TensorFlow.
+
+The `BUILD` file declares the dependencies and the output build targets. Refer to
+[building for OSS](https://www.tensorflow.org/guide/create_op#build_the_op_library).
+
+You will be reusing the `BUILD` file later in this example.
+
+```
+tf_custom_op_library(
+    name = "sleep_kernel.so",
+    srcs = [
+        "sleep_kernel.cc",
+        "sleep_op.cc",
+    ],
+    deps = [
+        "//third_party/absl/log",
+        "//third_party/absl/status",
+        "//third_party/absl/time",
+    ],
+)
+
+py_strict_library(
+    name = "sleep_op",
+    srcs = ["sleep_op.py"],
+    data = ["sleep_kernel.so"],
+    srcs_version = "PY3",
+)
+
+py_strict_binary(
+    name = "sleep_bin",
+    srcs = ["sleep_bin.py"],
+    srcs_version = "PY3",
+    deps = [
+        ":sleep_op",
+        "//third_party/py/numpy",
+        "//third_party/py/tensorflow",
+        "@absl_py//absl:app",
+    ],
+)
+```
+
+### Step 4 - Create the Python wrapper
+
+To create the Python wrapper, import and implement a function that serves as the
+op's public API and provides a docstring.
+
+```
+def AsyncSleep(delay, name=None):
+  """Pause for `delay` seconds (which need not be an integer).
+
+  This is an asynchronous (non-blocking) version of a sleep op. It includes
+  any time spent being blocked by another thread in `delay`. If it is blocked
+  for a fraction of the time specified by `delay`, it only calls `sleep`
+  (actually `usleep`) only for the remainder. If it is blocked for the full
+  time specified by `delay` or more, it returns without explicitly calling
+  `sleep`.
+
+  Args:
+    delay: tf.Tensor which is a scalar of type float.
+    name: An optional name for the op.
+
+  Returns:
+    The `delay` value.
+  """
+  return gen_sleep_op.examples_async_sleep(delay=delay, name=name)
+```
+
+```
+def SyncSleep(delay, name=None):
+  """Pause for `delay` seconds (which need not be an integer).
+
+  This is a synchronous (blocking) version of a sleep op. It's purpose is
+  to be contrasted with Examples>AsyncSleep.
+
+  Args:
+    delay: tf.Tensor which is a scalar of type float.
+    name: An optional name for the op.
+
+  Returns:
+    The `delay` value.
+  """
+  return gen_sleep_op.examples_sync_sleep(delay=delay, name=name)
+```
+
+### Step 5 - Test the op
+
+Create op tests using classes derived from
+[`tf.test.TestCase`](https://www.tensorflow.org/api_docs/python/tf/test/TestCase).
+
+When writing tests to ensure that the op works correctly in both graph and eager
+executions, it is important to note that errors in the op code may be detected
+in two distinct phases of code execution depending on how it is executed (eager
+or graph). Errors may be detected early by the shape function or a
+bit later from the logic in the `Compute` and `ComputeAsync` methods. This may
+lead to differing error types and/or messages.
+
+```
+class SleepTest(tf.test.TestCase):
+
+  def _check_sleep(self, op):
+    """Check that one sleep op works in isolation.
+
+    See sleep_bin.py for an example of how the synchronous and asynchronous
+    sleep ops differ in behavior.
+
+    Args:
+      op: The sleep op, either sleep_op.SyncSleep or sleep_op.AsyncSleep.
+    """
+    delay = 0.3  # delay in seconds
+    start_t = time.time()
+    func = tf.function(lambda: op(delay))
+    results = self.evaluate(func())
+    end_t = time.time()
+    delta_t = end_t - start_t
+    self.assertEqual(results.shape, tuple())
+    self.assertGreater(delta_t, 0.9 * delay)
+
+  def test_sync_sleep(self):
+    self._check_sleep(sleep_op.SyncSleep)
+
+  def test_async_sleep(self):
+    self._check_sleep(sleep_op.AsyncSleep)
+
+  def test_async_sleep_error(self):
+    # It is import that ComputeAsync() calls its done() callback if it returns
+    # early due to an error.
+    func = tf.function(lambda: sleep_op.AsyncSleep(-1.0))
+    with self.assertRaisesRegex(errors_impl.InvalidArgumentError,
+                                'Input `delay` must be non-negative.'):
+      self.evaluate(func())
+```
+
+The tests in `sleep_test.py` only tests a single sleep op in isolation, so they
+do not demonstrate the difference in behavior between the asynchronous and
+synchronous versions. Ideally, tests should cover every codepath to confirm that
+`done` is always called.
+
+You can test the op and create a build that shows the difference between
+synchronous and asynchronous behavior. After running the `sleep_bin.py` binary,
+the output will look something like this:
+
+<!-- test_snippets_in_readme skip -->
+```
+Using synchronous sleep op with each of 50 ops
+sleeping for about 1.00 seconds, so total time is about 1.00 * ceil(50 /
+NUMBER_OF_THREADS). 16 is a typical number of threads, which would be 4.00
+seconds. The actual time will be a little greater.
+
+Total time = 4.555 seconds using <function SyncSleep at 0x7f19a401b4d0> Returned
+values from the ops:
+
+[1. 1.0001 1.0002 1.0003 1.0004 1.0005 1.0006 1.0007 1.0008 1.0009 1.001 1.0011
+1.0012 1.0013 1.0014 1.0015 1.0016 1.0017 1.0018 1.0019 1.002 1.0021 1.0022
+1.0023 1.0024 1.0025 1.0026 1.0027 1.0028 1.0029 1.003 1.0031 1.0032 1.0033
+1.0034 1.0035 1.0036 1.0037 1.0038 1.0039 1.004 1.0041 1.0042 1.0043 1.0044
+1.0045 1.0046 1.0047 1.0048 1.0049]
+
+Using asynchronous sleep op with each of 50 ops sleeping only as much as
+necessary so they finish after at least 1.00 seconds. Time that an op spends
+blocked waiting to finish counts as all or part of its delay. The returned
+values show how long each ops sleeps or 0 if the op does not need to sleep. The
+expected total time will be a little greater than the requested delay of 1.00
+seconds.
+
+Total time = 1.450 seconds using <function AsyncSleep at 0x7f19a401b170>
+Returned values from the ops: [0. 0. 0. 0. 0. 1.0001 1.0005 0. 0. 1.0008 0. 0.
+1.0011 1.0013 0. 0. 0.0004 0.0012 1.0018 0.0001 1.002 0.001 0. 0.0006 0. 0.
+0.0005 0. 0.001 0. 0. 1.003 1.0031 0.0003 0.0005 0.0033 1.0032 0.0004 1.0037
+1.0038 0.0014 0.0023 0.0008 0.0024 0.0014 1.0044 1.0045 1.0046 1.0047 0.004 ]
+```
+
+The returned values from the synchronous ops are the same as the input delay
+values, which are all different but close to 1 second. The returned values from
+the asynchronous ops are the time blocked using `absl::SleepFor`, or 0 if the op
+did not sleep.
+
+In this example, 16 of the 50 ops slept for approximately the entire requested
+time, 17 of the 50 ops slept for a fraction of the time, and 17 did not sleep at
+all.
+
+### Use the op
+
+Use sleep_bin.py to explore the differences between synchronous and asynchronous
+behavior.
+
+The file creates a
+[`tf.stack`](https://www.tensorflow.org/api_docs/python/tf/stack) of 50 sleep
+custom ops that can potentially run in parallel. Each op receives a slightly
+different input argument so that the ops are not combined as common
+subexpressions.
+
+Note: Multiple sleep ops with different inputs can be used in one `tf.function`
+without being combined by optimization.
+
+```
+def stack50(op, delay):
+  """Create a tf.stack of 50 sleep ops.
+
+  Args:
+    op: The sleep op, either sleep_op.SyncSleep or sleep_op.AsyncSleep.
+    delay: Each op should finish at least float `delay` seconds after it starts.
+  """
+  n = 50
+  delays = delay + tf.range(0, n, dtype=float) / 10000.0
+  start_t = time.time()
+  func = tf.function(lambda: tf.stack([op(delays[i]) for i in range(n)]))
+  r_numpy = func().numpy()
+  end_t = time.time()
+  print('')
+  print('Total time = %5.3f seconds using %s' % (end_t - start_t, str(op)))
+  print('Returned values from the ops:')
+  np.set_printoptions(precision=4, suppress=True)
+  print(r_numpy)
+  sys.stdout.flush()
+```
+
+When `stack50` is called with the synchronous op, each of the 50 ops sleeps
+independently. The number of threads used to execute ops is configurable. With
+the default of 16 threads, the total time is a little greater than 4 seconds
+(ceil(50/16) = 4, with a delay of 1 second).
+
+<!-- test_snippets_in_readme skip -->
+```python
+  delay_seconds = 1.0
+  stack50(sleep_op.SyncSleep, delay_seconds)
+```
+
+When `stack50` is called with the asynchronous op, each op considers any time it
+spends blocked waiting to be scheduled in a thread pool as part of its delay,
+and only sleeps for the remaining time. For a delay of 1 second, the total
+time is a little greater than 1 second (regardless of how many ops are sleeping
+or how many threads are in the thread pool).
+
+<!-- test_snippets_in_readme skip -->
+```python
+  delay_seconds = 1.0
+  stack50(sleep_op.AsyncSleep, delay_seconds)
+```
+
+A more appropriate solution for actual delay or polling use cases is to
+use queuing or scheduling, where one or more asynchronous ops manage multiple
+delays or polls.
+
+### Summary
+
+In this example, you learned how to implement a synchronous and asynchronous
+custom op for GPU. Using a helper method, you implemented an asynchronous sleep
+op.
+
+The table below summarizes the build rules and targets for building and testing
+the `sleep` op.
+
+| Op components  | Build rule             | Build target   | Source            |
+| -------------- | ---------------------- | -------------- | ----------------- |
+| Kernels (C++)  | `tf_custom_op_library` | `sleep_kernel` | `sleep_kernel.cc` |
+| Wrapper        | N/A                    | `gen_sleep_op` | N/A               |
+: (automatically :                        :                :                   :
+: generated)     :                        :                :                   :
+| Wrapper (with  | `py_strict_library`    | `sleep_op`     | `sleep_op.py`     |
+: public API and :                        :                :                   :
+: docstring)     :                        :                :                   :
+| Tests          | `tf_py_test`           | `sleep_test`   | `sleep_test.py`   |
+| Example        | `py_strict_binary`     | `sleep_bin`    | `sleep_bin.py`    |
+<!-- LINT.ThenChange(sleep.md) -->

tensorflow/examples/custom_ops_doc/sleep/sleep_bin.py

@@ -0,0 +1,72 @@
+# Copyright 2021 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.
+# ==============================================================================
+"""Simple binary for sleep."""
+
+import math
+import sys
+import time
+
+from absl import app
+import numpy as np
+import tensorflow as tf
+
+from tensorflow.examples.custom_ops_doc.sleep import sleep_op
+
+
+def stack50(op, delay):
+  """Create a tf.stack of 50 sleep ops.
+
+  Args:
+    op: The sleep op, either sleep_op.SyncSleep or sleep_op.AsyncSleep.
+    delay: Each op should finish at least float `delay` seconds after it starts.
+  """
+  n = 50
+  delays = delay + tf.range(0, n, dtype=float) / 10000.0
+  start_t = time.time()
+  func = tf.function(lambda: tf.stack([op(delays[i]) for i in range(n)]))
+  r_numpy = func().numpy()
+  end_t = time.time()
+  print('')
+  print('Total time = %5.3f seconds using %s' % (end_t - start_t, str(op)))
+  print('Returned values from the ops:')
+  np.set_printoptions(precision=4, suppress=True)
+  print(r_numpy)
+  sys.stdout.flush()
+
+
+def main(argv):
+  del argv  # not used
+  delay_seconds = 1.0
+  print("""
+Using synchronous sleep op with each of 50 ops sleeping for about %0.2f seconds,
+so total time is about %0.2f * ceil(50 / NUMBER_OF_THREADS). 16 is a typical
+number of threads, which would be %0.2f seconds. The actual time will be
+a little greater.
+""" % (delay_seconds, delay_seconds, delay_seconds * math.ceil(50.0 / 16.0)))
+  stack50(sleep_op.SyncSleep, delay_seconds)
+
+  print("""
+Using asynchronous sleep op with each of 50 ops sleeping only as much as
+necessary so they finish after at least %0.2f seconds. Time that
+an op spends blocked waiting to finish counts as all or part of its delay.
+The returned values show how long each ops sleeps or 0 if the op does not
+need to sleep. The expected total time will be a little greater than
+the requested delay of %0.2f seconds.
+""" % (delay_seconds, delay_seconds))
+  stack50(sleep_op.AsyncSleep, delay_seconds)
+
+
+if __name__ == '__main__':
+  app.run(main)

tensorflow/examples/custom_ops_doc/sleep/sleep_kernel.cc

@@ -0,0 +1,139 @@
+/* Copyright 2021 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 "absl/log/log.h"
+#include "absl/time/clock.h"
+#include "absl/time/time.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/register_types.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/platform/threadpool.h"
+
+// Please use the appropriate namespace for your project
+namespace tensorflow {
+namespace custom_op_examples {
+
+using ::tensorflow::errors::Internal;
+using ::tensorflow::errors::InvalidArgument;
+
+class AsyncSleepOp : public AsyncOpKernel {
+ public:
+  explicit AsyncSleepOp(OpKernelConstruction* ctx) : AsyncOpKernel(ctx) {}
+  AsyncSleepOp(const AsyncSleepOp& other) = delete;
+  AsyncSleepOp& operator=(const AsyncSleepOp& other) = delete;
+  ~AsyncSleepOp() override = default;
+
+  // Implementations of ComputeAsync() must ensure that `done` is (eventually)
+  // called exactly once to signal the completion of the computation. The
+  // implementation of ComputeAsync() must not block on the execution of another
+  // OpKernel. `done` may be called by the current thread, or by another thread.
+  // `context` is guaranteed to stay alive until the `done` callback starts.
+  // For example, use OP_REQUIRES_ASYNC which takes the `done` parameter
+  // as an input and calls `done` for the case of exiting early with an error
+  // (instead of OP_REQUIRES).
+  //
+  // Since it is possible that the unblocking kernel may never run (due to an
+  // error or cancellation), in most cases the AsyncOpKernel should implement
+  // cancellation support via `context->cancellation_manager()`.
+  // TODO (schwartzedward): should this use cancellation support?
+  //
+  // WARNING: As soon as the `done` callback starts, `context` and `this` may be
+  // deleted. No code depending on these objects should execute after the call
+  // to `done`.
+  void ComputeAsync(OpKernelContext* ctx, DoneCallback done) override {
+    const auto& delay_tensor = ctx->input(0);
+    OP_REQUIRES_ASYNC(
+        ctx, ::tensorflow::TensorShapeUtils::IsScalar(delay_tensor.shape()),
+        InvalidArgument("Input `delay` must be a scalar."),
+        done);  // Important: call `done` in every execution path
+    const float delay = delay_tensor.flat<float>()(0);
+    OP_REQUIRES_ASYNC(ctx, delay >= 0.0,
+                      InvalidArgument("Input `delay` must be non-negative."),
+                      done);  // Important: call `done` in every execution path
+    auto thread_pool = ctx->device()->tensorflow_cpu_worker_threads()->workers;
+    OP_REQUIRES_ASYNC(ctx, thread_pool != nullptr,
+                      Internal("No thread_pool found."),
+                      done);  // Important: call `done` in every execution path
+
+    Tensor* output_tensor = nullptr;
+    OP_REQUIRES_OK_ASYNC(
+        ctx, ctx->allocate_output(0, delay_tensor.shape(), &output_tensor),
+        done);  // Important: call `done` in every execution path
+
+    absl::Time now = absl::Now();
+    absl::Time when = now + absl::Seconds(delay);
+    VLOG(1) << "BEFORE ASYNC SLEEP " << ctx->op_kernel().name() << " now "
+            << now << " when " << when;
+    thread_pool->Schedule([this, output_tensor, when, done] {
+      this->sleeper(output_tensor, when, done);
+    });
+    // Note that `done` is normally called by sleeper(), it is not normally
+    // called by this function.
+  }
+
+ private:
+  void sleeper(Tensor* output_tensor, absl::Time when, DoneCallback done) {
+    absl::Time now = absl::Now();
+    int64_t delay_us = 0;
+    if (now < when) {
+      delay_us = absl::ToInt64Microseconds(when - now);
+      VLOG(1) << "MIDDLE ASYNC SLEEP " << delay_us;
+      absl::SleepFor(when - now);
+      VLOG(1) << "AFTER ASYNC SLEEP " << delay_us;
+    } else {
+      VLOG(1) << "MIDDLE/AFTER ASYNC SKIP SLEEP";
+    }
+    auto output = output_tensor->template flat<float>();
+    output(0) = static_cast<float>(delay_us) / 1000000.0;
+    done();  // Important: call `done` in every execution path
+  }
+};
+
+class SyncSleepOp : public OpKernel {
+ public:
+  explicit SyncSleepOp(OpKernelConstruction* ctx) : OpKernel(ctx) {}
+  SyncSleepOp(const SyncSleepOp& other) = delete;
+  SyncSleepOp& operator=(const SyncSleepOp& other) = delete;
+  ~SyncSleepOp() override = default;
+
+  void Compute(OpKernelContext* ctx) override {
+    const auto& delay_tensor = ctx->input(0);
+    OP_REQUIRES(ctx,
+                ::tensorflow::TensorShapeUtils::IsScalar(delay_tensor.shape()),
+                InvalidArgument("Input `delay` must be a scalar."));
+    const float delay = delay_tensor.flat<float>()(0);
+    OP_REQUIRES(ctx, delay >= 0.0,
+                InvalidArgument("Input `delay` must be non-negative."));
+    VLOG(1) << "BEFORE SYNC SLEEP" << ctx->op_kernel().name();
+    absl::SleepFor(absl::Seconds(delay));
+    VLOG(1) << "AFTER SYNC SLEEP" << ctx->op_kernel().name();
+
+    Tensor* output_tensor = nullptr;
+    OP_REQUIRES_OK(
+        ctx, ctx->allocate_output(0, delay_tensor.shape(), &output_tensor));
+    auto output = output_tensor->template flat<float>();
+    output(0) = delay;
+  }
+};
+
+REGISTER_KERNEL_BUILDER(
+    Name("Examples>AsyncSleep").Device(::tensorflow::DEVICE_CPU), AsyncSleepOp)
+REGISTER_KERNEL_BUILDER(
+    Name("Examples>SyncSleep").Device(::tensorflow::DEVICE_CPU), SyncSleepOp)
+
+}  // namespace custom_op_examples
+}  // namespace tensorflow

tensorflow/examples/custom_ops_doc/sleep/sleep_op.cc

@@ -0,0 +1,72 @@
+/* Copyright 2021 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 "absl/status/status.h"
+#include "tensorflow/core/framework/op.h"
+#include "tensorflow/core/framework/shape_inference.h"
+
+// Use a namespace when registering by prepending the
+// package's name to the op’s name and separate with a '>'.
+// This is the recommendation for out-of-tree ops to avoid name collisions in
+// "Best practices for custom operations in TensorFlow"
+// https://github.com/tensorflow/community/blob/master/rfcs/20190726-custom-ops.md
+
+// AsyncSleep returns it's `delay` input so it can be used as a function.
+
+using ::tensorflow::shape_inference::InferenceContext;
+
+absl::Status ScalarOutput(InferenceContext* c) {
+  c->set_output(0, c->Scalar());
+  return absl::OkStatus();
+}
+
+REGISTER_OP("Examples>AsyncSleep")
+    .Input("delay: float")
+    .Output("output: float")
+    .SetShapeFn([](tensorflow::shape_inference::InferenceContext* c) {
+      tensorflow::shape_inference::ShapeHandle handle;
+      TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 0, &handle));
+      return (ScalarOutput(c));
+    })
+    .Doc(R"doc(
+Pause for `delay` seconds (which need not be an integer).
+
+This is an asynchronous (non-blocking) version of sleep. It is intended to
+be an example of how to implement ops that do I/O or that block on other ops.
+
+delay: tf.Tensor which is a scalar of type float.
+
+Returns the time spent in blocking sleep (which may be less that `delay` or
+zero if other ops run while this is waiting asynchronously).
+)doc");
+
+REGISTER_OP("Examples>SyncSleep")
+    .Input("delay: float")
+    .Output("output: float")
+    .SetShapeFn([](tensorflow::shape_inference::InferenceContext* c) {
+      tensorflow::shape_inference::ShapeHandle handle;
+      TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 0, &handle));
+      return (ScalarOutput(c));
+    })
+    .Doc(R"doc(
+Pause for `delay` seconds (which need not be an integer).
+
+This is a synchronous (blocking) version of sleep. It's purpose is
+to be contrasted with Examples>AsyncSleep.
+
+delay: tf.Tensor which is a scalar of type float.
+
+Returns `delay`.
+)doc");

tensorflow/examples/custom_ops_doc/sleep/sleep_op.py

@@ -0,0 +1,69 @@
+# Copyright 2021 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.
+# ==============================================================================
+"""A wrapper for gen_sleep_op.py.
+
+This defines a public API (and provides a docstring for it) for the C++ Op
+defined by sleep_kernel.cc
+"""
+
+import tensorflow as tf
+from tensorflow.python.platform import resource_loader
+
+_sleep_module = tf.load_op_library(
+    resource_loader.get_path_to_datafile("sleep_kernel.so"))
+
+examples_async_sleep = _sleep_module.examples_async_sleep
+examples_sync_sleep = _sleep_module.examples_sync_sleep
+
+
+# In this example, this Python function is a trivial wrapper for the C++ Op:
+# it provides a public API and docstring that are equivalent to the API
+# and documentation of the C++ op. The motivation for it is to be a placeholder
+# that allows a wider variety of non-breaking future changes than are possible
+# with the generated wrapper alone. Having this wrapper is optional.
+def AsyncSleep(delay, name=None):
+  """Pause for `delay` seconds (which need not be an integer).
+
+  This is an asynchronous (non-blocking) version of a sleep op. It includes
+  any time spent being blocked by another thread in `delay`. If it is blocked
+  for a fraction of the time specified by `delay`, it only calls `sleep`
+  (actually `usleep`) only for the remainder. If it is blocked for the full
+  time specified by `delay` or more, it returns without explicitly calling
+  `sleep`.
+
+  Args:
+    delay: tf.Tensor which is a scalar of type float.
+    name: An optional name for the op.
+
+  Returns:
+    The `delay` value.
+  """
+  return examples_async_sleep(delay=delay, name=name)
+
+
+def SyncSleep(delay, name=None):
+  """Pause for `delay` seconds (which need not be an integer).
+
+  This is a synchronous (blocking) version of a sleep op. It's purpose is
+  to be contrasted with Examples>AsyncSleep.
+
+  Args:
+    delay: tf.Tensor which is a scalar of type float.
+    name: An optional name for the op.
+
+  Returns:
+    The `delay` value.
+  """
+  return examples_sync_sleep(delay=delay, name=name)

tensorflow/examples/custom_ops_doc/sleep/sleep_test.py

@@ -0,0 +1,62 @@
+# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for sleep."""
+
+import time
+
+import tensorflow as tf
+
+from tensorflow.examples.custom_ops_doc.sleep import sleep_op
+# This pylint disable is only needed for internal google users
+from tensorflow.python.framework import errors_impl  # pylint: disable=g-direct-tensorflow-import
+
+
+class SleepTest(tf.test.TestCase):
+
+  def _check_sleep(self, op):
+    """Check that one sleep op works in isolation.
+
+    See sleep_bin.py for an example of how the synchronous and asynchronous
+    sleep ops differ in behavior.
+
+    Args:
+      op: The sleep op, either sleep_op.SyncSleep or sleep_op.AsyncSleep.
+    """
+    delay = 0.3  # delay in seconds
+    start_t = time.time()
+    func = tf.function(lambda: op(delay))
+    results = self.evaluate(func())
+    end_t = time.time()
+    delta_t = end_t - start_t
+    self.assertEqual(results.shape, tuple())
+    self.assertGreater(delta_t, 0.9 * delay)
+
+  def test_sync_sleep(self):
+    self._check_sleep(sleep_op.SyncSleep)
+
+  def test_async_sleep(self):
+    self._check_sleep(sleep_op.AsyncSleep)
+
+  def test_async_sleep_error(self):
+    # It is import that ComputeAsync() calls its done() callback if it returns
+    # early due to an error.
+    func = tf.function(lambda: sleep_op.AsyncSleep(-1.0))
+    with self.assertRaisesRegex(errors_impl.InvalidArgumentError,
+                                'Input `delay` must be non-negative.'):
+      self.evaluate(func())
+
+
+if __name__ == '__main__':
+  tf.test.main()

tensorflow/examples/image_retraining/README.md

@@ -0,0 +1,15 @@
+**NOTE: This code has moved to**
+https://github.com/tensorflow/hub/tree/master/examples/image_retraining
+
+retrain.py is an example script that shows how one can adapt a pretrained
+network for other classification problems (including use with TFLite and
+quantization).
+
+As of TensorFlow 1.7, it is recommended to use a pretrained network from
+TensorFlow Hub, using the new version of this example found in the location
+above, as explained in TensorFlow's revised
+[image retraining tutorial](https://www.tensorflow.org/hub/tutorials/tf2_image_retraining).
+
+Older versions of this example (using frozen GraphDefs instead of
+TensorFlow Hub modules) are available in the release branches of
+TensorFlow versions up to and including 1.7.

tensorflow/examples/label_image/BUILD

@@ -0,0 +1,74 @@
+# Description:
+#   TensorFlow C++ inference example for labeling images.
+
+load("@xla//third_party/rules_python/python:py_binary.bzl", "py_binary")
+load("//tensorflow:tensorflow.bzl", "tf_cc_binary")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = ["//tensorflow:internal"],
+    licenses = ["notice"],
+)
+
+exports_files(["data/grace_hopper.jpg"])
+
+tf_cc_binary(
+    name = "label_image",
+    srcs = [
+        "main.cc",
+    ],
+    linkopts = select({
+        "//tensorflow:android": [
+            "-pie",
+            "-landroid",
+            "-ljnigraphics",
+            "-llog",
+            "-lm",
+            "-z defs",
+            "-s",
+            "-Wl,--exclude-libs,ALL",
+        ],
+        "//conditions:default": ["-lm"],
+    }),
+    deps = select({
+        "//tensorflow:android": [
+            "//tensorflow/core:portable_tensorflow_lib",
+            # cc:android_tensorflow_image_op is for including jpeg/gif/png
+            # decoder to enable real-image evaluation on Android
+            "//tensorflow/core/kernels/image:android_tensorflow_image_op",
+        ],
+        "//conditions:default": [
+            "//tensorflow/core:core_cpu",
+            "//tensorflow/core:framework",
+            "//tensorflow/core:framework_internal",
+            "//tensorflow/core:lib",
+            "//tensorflow/core:protos_all_cc",
+            "//tensorflow/core:tensorflow",
+        ],
+    }) + [
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:string_view",
+        "@com_google_absl//absl/types:span",
+        "//tensorflow/cc:ops",
+        "//tensorflow/cc:scope",
+        "@xla//xla/tsl/util:command_line_flags",
+        "@xla//xla/tsl/platform:status",
+        "@xla//xla/tsl/platform:types",
+        # cc:cc_ops is used to include image ops (for label_image)
+        # Jpg, gif, and png related code won't be included
+        "//tensorflow/cc:cc_ops",
+    ],
+)
+
+py_binary(
+    name = "label_image_py",
+    srcs = ["label_image.py"],
+    main = "label_image.py",
+    strict_deps = True,
+    deps = [
+        "//tensorflow:tensorflow_py",
+        "//third_party/py/numpy",
+    ],
+)

tensorflow/examples/label_image/README.md

@@ -0,0 +1,99 @@
+# TensorFlow C++ and Python Image Recognition Demo
+
+This example shows how you can load a pre-trained TensorFlow network and use it
+to recognize objects in images in C++. For Java see the [Java
+README](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/java),
+and for Go see the [godoc
+example](https://godoc.org/github.com/tensorflow/tensorflow/tensorflow/go#ex-package).
+
+## Description
+
+This demo uses a Google Inception model to classify image files that are passed
+in on the command line.
+
+## To build/install/run
+
+The TensorFlow `GraphDef` that contains the model definition and weights is not
+packaged in the repo because of its size. Instead, you must first download the
+file to the `data` directory in the source tree:
+
+```bash
+$ curl -L "https://storage.googleapis.com/download.tensorflow.org/models/inception_v3_2016_08_28_frozen.pb.tar.gz" |
+  tar -C tensorflow/examples/label_image/data -xz
+```
+
+Then, as long as you've managed to build the main TensorFlow framework, you
+should have everything you need to run this example installed already.
+
+Once extracted, see the labels file in the data directory for the possible
+classifications, which are the 1,000 categories used in the Imagenet
+competition.
+
+To build it, run this command:
+
+```bash
+$ bazel build tensorflow/examples/label_image/...
+```
+
+That should build a binary executable that you can then run like this:
+
+```bash
+$ bazel-bin/tensorflow/examples/label_image/label_image
+```
+
+This uses the default example image that ships with the framework, and should
+output something similar to this:
+
+```
+I tensorflow/examples/label_image/main.cc:206] military uniform (653): 0.834306
+I tensorflow/examples/label_image/main.cc:206] mortarboard (668): 0.0218692
+I tensorflow/examples/label_image/main.cc:206] academic gown (401): 0.0103579
+I tensorflow/examples/label_image/main.cc:206] pickelhaube (716): 0.00800814
+I tensorflow/examples/label_image/main.cc:206] bulletproof vest (466): 0.00535088
+```
+
+In this case, we're using the default image of Admiral Grace Hopper, and you can
+see the network correctly spots she's wearing a military uniform, with a high
+score of 0.8.
+
+Next, try it out on your own images by supplying the --image= argument, e.g.
+
+```bash
+$ bazel-bin/tensorflow/examples/label_image/label_image --image=my_image.png
+```
+
+For a more detailed look at this code, you can check out the C++ section of the
+[Inception tutorial](https://github.com/tensorflow/docs/blob/master/site/en/r1/tutorials/images/image_recognition.md).
+
+## Python implementation
+
+label_image.py is a python implementation that provides code corresponding to
+the C++ code here. This gives more intuitive mapping between C++ and Python than
+the Python code mentioned in the
+[Inception tutorial](https://github.com/tensorflow/docs/blob/master/site/en/r1/tutorials/images/image_recognition.md).
+and could be easier to add visualization or debug code.
+
+`bazel-bin/tensorflow/examples/label_image/label_image_py` should be there after
+```bash
+$ bazel build tensorflow/examples/label_image/...
+```
+
+Run
+
+```bash
+$ bazel-bin/tensorflow/examples/label_image/label_image_py
+```
+
+Or, with tensorflow python package installed, you can run it like:
+```bash
+$ python3 tensorflow/examples/label_image/label_image.py
+```
+
+And get result similar to this:
+```
+military uniform 0.834305
+mortarboard 0.0218694
+academic gown 0.0103581
+pickelhaube 0.00800818
+bulletproof vest 0.0053509
+```

tensorflow/examples/label_image/label_image.py

@@ -0,0 +1,131 @@
+# 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.
+# ==============================================================================
+
+import argparse
+
+import numpy as np
+import tensorflow as tf
+tf.compat.v1.disable_eager_execution()
+
+def load_graph(model_file):
+  graph = tf.Graph()
+  graph_def = tf.compat.v1.GraphDef()
+
+  with open(model_file, "rb") as f:
+    graph_def.ParseFromString(f.read())
+  with graph.as_default():
+    tf.import_graph_def(graph_def)
+
+  return graph
+
+
+def read_tensor_from_image_file(file_name,
+                                input_height=299,
+                                input_width=299,
+                                input_mean=0,
+                                input_std=255):
+  input_name = "file_reader"
+  output_name = "normalized"
+  file_reader = tf.io.read_file(file_name, input_name)
+  if file_name.endswith(".png"):
+    image_reader = tf.io.decode_png(file_reader, channels=3, name="png_reader")
+  elif file_name.endswith(".gif"):
+    image_reader = tf.squeeze(tf.io.decode_gif(file_reader, name="gif_reader"))
+  elif file_name.endswith(".bmp"):
+    image_reader = tf.io.decode_bmp(file_reader, name="bmp_reader")
+  else:
+    image_reader = tf.io.decode_jpeg(
+        file_reader, channels=3, name="jpeg_reader")
+  float_caster = tf.cast(image_reader, tf.float32)
+  dims_expander = tf.expand_dims(float_caster, 0)
+  resized = tf.compat.v1.image.resize_bilinear(
+      dims_expander, [input_height, input_width]
+  )
+  normalized = tf.divide(tf.subtract(resized, [input_mean]), [input_std])
+  sess = tf.compat.v1.Session()
+  return sess.run(normalized)
+
+
+def load_labels(label_file):
+  proto_as_ascii_lines = tf.io.gfile.GFile(label_file).readlines()
+  return [l.rstrip() for l in proto_as_ascii_lines]
+
+
+if __name__ == "__main__":
+  file_name = "tensorflow/examples/label_image/data/grace_hopper.jpg"
+  model_file = \
+    "tensorflow/examples/label_image/data/inception_v3_2016_08_28_frozen.pb"
+  label_file = "tensorflow/examples/label_image/data/imagenet_slim_labels.txt"
+  input_height = 299
+  input_width = 299
+  input_mean = 0
+  input_std = 255
+  input_layer = "input"
+  output_layer = "InceptionV3/Predictions/Reshape_1"
+
+  parser = argparse.ArgumentParser()
+  parser.add_argument("--image", help="image to be processed")
+  parser.add_argument("--graph", help="graph/model to be executed")
+  parser.add_argument("--labels", help="name of file containing labels")
+  parser.add_argument("--input_height", type=int, help="input height")
+  parser.add_argument("--input_width", type=int, help="input width")
+  parser.add_argument("--input_mean", type=int, help="input mean")
+  parser.add_argument("--input_std", type=int, help="input std")
+  parser.add_argument("--input_layer", help="name of input layer")
+  parser.add_argument("--output_layer", help="name of output layer")
+  args = parser.parse_args()
+
+  if args.graph:
+    model_file = args.graph
+  if args.image:
+    file_name = args.image
+  if args.labels:
+    label_file = args.labels
+  if args.input_height:
+    input_height = args.input_height
+  if args.input_width:
+    input_width = args.input_width
+  if args.input_mean:
+    input_mean = args.input_mean
+  if args.input_std:
+    input_std = args.input_std
+  if args.input_layer:
+    input_layer = args.input_layer
+  if args.output_layer:
+    output_layer = args.output_layer
+
+  graph = load_graph(model_file)
+  t = read_tensor_from_image_file(
+      file_name,
+      input_height=input_height,
+      input_width=input_width,
+      input_mean=input_mean,
+      input_std=input_std)
+
+  input_name = "import/" + input_layer
+  output_name = "import/" + output_layer
+  input_operation = graph.get_operation_by_name(input_name)
+  output_operation = graph.get_operation_by_name(output_name)
+
+  with tf.compat.v1.Session(graph=graph) as sess:
+    results = sess.run(output_operation.outputs[0], {
+        input_operation.outputs[0]: t
+    })
+  results = np.squeeze(results)
+
+  top_k = results.argsort()[-5:][::-1]
+  labels = load_labels(label_file)
+  for i in top_k:
+    print(labels[i], results[i])

tensorflow/examples/label_image/main.cc

@@ -0,0 +1,402 @@
+/* Copyright 2015 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.
+==============================================================================*/
+
+// A minimal but useful C++ example showing how to load an Imagenet-style object
+// recognition TensorFlow model, prepare input images for it, run them through
+// the graph, and interpret the results.
+//
+// It's designed to have as few dependencies and be as clear as possible, so
+// it's more verbose than it could be in production code. In particular, using
+// auto for the types of a lot of the returned values from TensorFlow calls can
+// remove a lot of boilerplate, but I find the explicit types useful in sample
+// code to make it simple to look up the classes involved.
+//
+// To use it, compile and then run in a working directory with the
+// learning/brain/tutorials/label_image/data/ folder below it, and you should
+// see the top five labels for the example Lena image output. You can then
+// customize it to use your own models or images by changing the file names at
+// the top of the main() function.
+//
+// The googlenet_graph.pb file included by default is created from Inception.
+//
+// Note that, for GIF inputs, to reuse existing code, only single-frame ones
+// are supported.
+
+#include <algorithm>
+#include <cstddef>
+#include <cstdint>
+#include <fstream>
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "absl/strings/match.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+#include "absl/types/span.h"
+#include "tensorflow/cc/framework/ops.h"
+#include "tensorflow/cc/framework/scope.h"
+#include "tensorflow/cc/ops/array_ops.h"
+#include "tensorflow/cc/ops/const_op.h"
+#include "tensorflow/cc/ops/image_ops.h"
+#include "tensorflow/cc/ops/math_ops.h"
+#include "tensorflow/cc/ops/nn_ops.h"
+#include "xla/tsl/platform/env.h"
+#include "xla/tsl/platform/errors.h"
+#include "xla/tsl/platform/status.h"
+#include "xla/tsl/platform/types.h"
+#include "xla/tsl/util/command_line_flags.h"
+#include "tensorflow/core/framework/graph.pb.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/tensor_types.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/lib/core/stringpiece.h"
+#include "tensorflow/core/lib/strings/str_util.h"
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/file_system.h"
+#include "tensorflow/core/platform/init_main.h"
+#include "tensorflow/core/platform/logging.h"
+#include "tensorflow/core/platform/path.h"
+#include "tensorflow/core/platform/tstring.h"
+#include "tensorflow/core/platform/types.h"
+#include "tensorflow/core/public/session.h"
+#include "tensorflow/core/public/session_options.h"
+#include "tensorflow/core/util/command_line_flags.h"
+
+// These are all common classes it's handy to reference with no namespace.
+using tensorflow::Flag;
+using tensorflow::int32;
+using tensorflow::Status;
+using tensorflow::string;
+using tensorflow::Tensor;
+using tensorflow::tstring;
+
+// Takes a file name, and loads a list of labels from it, one per line, and
+// returns a vector of the strings. It pads with empty strings so the length
+// of the result is a multiple of 16, because our model expects that.
+Status ReadLabelsFile(const string& file_name, std::vector<string>* result,
+                      size_t* found_label_count) {
+  std::ifstream file(file_name);
+  if (!file) {
+    return absl::NotFoundError(
+        absl::StrCat("Labels file ", file_name, " not found."));
+  }
+  result->clear();
+  string line;
+  while (std::getline(file, line)) {
+    result->push_back(line);
+  }
+  *found_label_count = result->size();
+  const int padding = 16;
+  while (result->size() % padding) {
+    result->emplace_back();
+  }
+  return absl::OkStatus();
+}
+
+static Status ReadEntireFile(tensorflow::Env* env, const string& filename,
+                             Tensor* output) {
+  uint64_t file_size = 0;
+  TF_RETURN_IF_ERROR(env->GetFileSize(filename, &file_size));
+
+  string contents;
+  contents.resize(file_size);
+
+  std::unique_ptr<tensorflow::RandomAccessFile> file;
+  TF_RETURN_IF_ERROR(env->NewRandomAccessFile(filename, &file));
+
+  absl::string_view data;
+  TF_RETURN_IF_ERROR(
+      file->Read(0, data, absl::MakeSpan(&contents[0], file_size)));
+  if (data.size() != file_size) {
+    return absl::DataLossError(absl::StrCat("Truncated read of '", filename,
+                                            "' expected ", file_size, " got ",
+                                            data.size()));
+  }
+  output->scalar<tstring>()() = tstring(data);
+  return absl::OkStatus();
+}
+
+// Given an image file name, read in the data, try to decode it as an image,
+// resize it to the requested size, and then scale the values as desired.
+Status ReadTensorFromImageFile(const string& file_name, const int input_height,
+                               const int input_width, const float input_mean,
+                               const float input_std,
+                               std::vector<Tensor>* out_tensors) {
+  auto root = tensorflow::Scope::NewRootScope();
+  using namespace ::tensorflow::ops;  // NOLINT(build/namespaces)
+
+  string input_name = "file_reader";
+  string output_name = "normalized";
+
+  // read file_name into a tensor named input
+  Tensor input(tensorflow::DT_STRING, tensorflow::TensorShape());
+  TF_RETURN_IF_ERROR(
+      ReadEntireFile(tensorflow::Env::Default(), file_name, &input));
+
+  // use a placeholder to read input data
+  auto file_reader =
+      Placeholder(root.WithOpName("input"), tensorflow::DataType::DT_STRING);
+
+  std::vector<std::pair<string, tensorflow::Tensor>> inputs = {
+      {"input", input},
+  };
+
+  // Now try to figure out what kind of file it is and decode it.
+  const int wanted_channels = 3;
+  tensorflow::Output image_reader;
+  if (absl::EndsWith(file_name, ".png")) {
+    image_reader = DecodePng(root.WithOpName("png_reader"), file_reader,
+                             DecodePng::Channels(wanted_channels));
+  } else if (absl::EndsWith(file_name, ".gif")) {
+    // gif decoder returns 4-D tensor, remove the first dim
+    image_reader =
+        Squeeze(root.WithOpName("squeeze_first_dim"),
+                DecodeGif(root.WithOpName("gif_reader"), file_reader));
+  } else if (absl::EndsWith(file_name, ".bmp")) {
+    image_reader = DecodeBmp(root.WithOpName("bmp_reader"), file_reader);
+  } else {
+    // Assume if it's neither a PNG nor a GIF then it must be a JPEG.
+    image_reader = DecodeJpeg(root.WithOpName("jpeg_reader"), file_reader,
+                              DecodeJpeg::Channels(wanted_channels));
+  }
+  // Now cast the image data to float so we can do normal math on it.
+  auto float_caster =
+      Cast(root.WithOpName("float_caster"), image_reader, tensorflow::DT_FLOAT);
+  // The convention for image ops in TensorFlow is that all images are expected
+  // to be in batches, so that they're four-dimensional arrays with indices of
+  // [batch, height, width, channel]. Because we only have a single image, we
+  // have to add a batch dimension of 1 to the start with ExpandDims().
+  auto dims_expander = ExpandDims(root, float_caster, 0);
+  // Bilinearly resize the image to fit the required dimensions.
+  auto resized = ResizeBilinear(
+      root, dims_expander,
+      Const(root.WithOpName("size"), {input_height, input_width}));
+  // Subtract the mean and divide by the scale.
+  Div output_op(root.WithOpName(output_name), Sub(root, resized, {input_mean}),
+                {input_std});
+
+  // This runs the GraphDef network definition that we've just constructed, and
+  // returns the results in the output tensor.
+  tensorflow::GraphDef graph;
+  TF_RETURN_IF_ERROR(root.ToGraphDef(&graph));
+
+  std::unique_ptr<tensorflow::Session> session(
+      tensorflow::NewSession(tensorflow::SessionOptions()));
+  TF_RETURN_IF_ERROR(session->Create(graph));
+  TF_RETURN_IF_ERROR(session->Run({inputs}, {output_name}, {}, out_tensors));
+  return absl::OkStatus();
+}
+
+// Reads a model graph definition from disk, and creates a session object you
+// can use to run it.
+Status LoadGraph(const string& graph_file_name,
+                 std::unique_ptr<tensorflow::Session>* session) {
+  tensorflow::GraphDef graph_def;
+  Status load_graph_status =
+      ReadBinaryProto(tensorflow::Env::Default(), graph_file_name, &graph_def);
+  if (!load_graph_status.ok()) {
+    return absl::NotFoundError(absl::StrCat("Failed to load compute graph at '",
+                                            graph_file_name, "'"));
+  }
+  session->reset(tensorflow::NewSession(tensorflow::SessionOptions()));
+  Status session_create_status = (*session)->Create(graph_def);
+  if (!session_create_status.ok()) {
+    return session_create_status;
+  }
+  return absl::OkStatus();
+}
+
+// Analyzes the output of the Inception graph to retrieve the highest scores and
+// their positions in the tensor, which correspond to categories.
+Status GetTopLabels(const std::vector<Tensor>& outputs, int how_many_labels,
+                    Tensor* indices, Tensor* scores) {
+  auto root = tensorflow::Scope::NewRootScope();
+  using namespace ::tensorflow::ops;  // NOLINT(build/namespaces)
+
+  string output_name = "top_k";
+  TopK top_k(root.WithOpName(output_name), outputs[0], how_many_labels);
+  // This runs the GraphDef network definition that we've just constructed, and
+  // returns the results in the output tensors.
+  tensorflow::GraphDef graph;
+  TF_RETURN_IF_ERROR(root.ToGraphDef(&graph));
+
+  std::unique_ptr<tensorflow::Session> session(
+      tensorflow::NewSession(tensorflow::SessionOptions()));
+  TF_RETURN_IF_ERROR(session->Create(graph));
+  // The TopK node returns two outputs, the scores and their original indices,
+  // so we have to append :0 and :1 to specify them both.
+  std::vector<Tensor> out_tensors;
+  TF_RETURN_IF_ERROR(session->Run({}, {output_name + ":0", output_name + ":1"},
+                                  {}, &out_tensors));
+  *scores = out_tensors[0];
+  *indices = out_tensors[1];
+  return absl::OkStatus();
+}
+
+// Given the output of a model run, and the name of a file containing the labels
+// this prints out the top five highest-scoring values.
+Status PrintTopLabels(const std::vector<Tensor>& outputs,
+                      const string& labels_file_name) {
+  std::vector<string> labels;
+  size_t label_count;
+  Status read_labels_status =
+      ReadLabelsFile(labels_file_name, &labels, &label_count);
+  if (!read_labels_status.ok()) {
+    LOG(ERROR) << read_labels_status;
+    return read_labels_status;
+  }
+  const int how_many_labels = std::min(5, static_cast<int>(label_count));
+  Tensor indices;
+  Tensor scores;
+  TF_RETURN_IF_ERROR(GetTopLabels(outputs, how_many_labels, &indices, &scores));
+  tensorflow::TTypes<float>::Flat scores_flat = scores.flat<float>();
+  tensorflow::TTypes<int32>::Flat indices_flat = indices.flat<int32>();
+  for (int pos = 0; pos < how_many_labels; ++pos) {
+    const int label_index = indices_flat(pos);
+    const float score = scores_flat(pos);
+    LOG(INFO) << labels[label_index] << " (" << label_index << "): " << score;
+  }
+  return absl::OkStatus();
+}
+
+// This is a testing function that returns whether the top label index is the
+// one that's expected.
+Status CheckTopLabel(const std::vector<Tensor>& outputs, int expected,
+                     bool* is_expected) {
+  *is_expected = false;
+  Tensor indices;
+  Tensor scores;
+  const int how_many_labels = 1;
+  TF_RETURN_IF_ERROR(GetTopLabels(outputs, how_many_labels, &indices, &scores));
+  tensorflow::TTypes<int32>::Flat indices_flat = indices.flat<int32>();
+  if (indices_flat(0) != expected) {
+    LOG(ERROR) << "Expected label #" << expected << " but got #"
+               << indices_flat(0);
+    *is_expected = false;
+  } else {
+    *is_expected = true;
+  }
+  return absl::OkStatus();
+}
+
+int main(int argc, char* argv[]) {
+  // These are the command-line flags the program can understand.
+  // They define where the graph and input data is located, and what kind of
+  // input the model expects. If you train your own model, or use something
+  // other than inception_v3, then you'll need to update these.
+  string image = "tensorflow/examples/label_image/data/grace_hopper.jpg";
+  string graph =
+      "tensorflow/examples/label_image/data/inception_v3_2016_08_28_frozen.pb";
+  string labels =
+      "tensorflow/examples/label_image/data/imagenet_slim_labels.txt";
+  int32_t input_width = 299;
+  int32_t input_height = 299;
+  float input_mean = 0;
+  float input_std = 255;
+  string input_layer = "input";
+  string output_layer = "InceptionV3/Predictions/Reshape_1";
+  bool self_test = false;
+  string root_dir = "";
+  std::vector<Flag> flag_list = {
+      Flag("image", &image, "image to be processed"),
+      Flag("graph", &graph, "graph to be executed"),
+      Flag("labels", &labels, "name of file containing labels"),
+      Flag("input_width", &input_width, "resize image to this width in pixels"),
+      Flag("input_height", &input_height,
+           "resize image to this height in pixels"),
+      Flag("input_mean", &input_mean, "scale pixel values to this mean"),
+      Flag("input_std", &input_std, "scale pixel values to this std deviation"),
+      Flag("input_layer", &input_layer, "name of input layer"),
+      Flag("output_layer", &output_layer, "name of output layer"),
+      Flag("self_test", &self_test, "run a self test"),
+      Flag("root_dir", &root_dir,
+           "interpret image and graph file names relative to this directory"),
+  };
+  string usage = tensorflow::Flags::Usage(argv[0], flag_list);
+  const bool parse_result = tensorflow::Flags::Parse(&argc, argv, flag_list);
+  if (!parse_result) {
+    LOG(ERROR) << usage;
+    return -1;
+  }
+
+  // We need to call this to set up global state for TensorFlow.
+  tensorflow::port::InitMain(argv[0], &argc, &argv);
+  if (argc > 1) {
+    LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage;
+    return -1;
+  }
+
+  // First we load and initialize the model.
+  std::unique_ptr<tensorflow::Session> session;
+  string graph_path = tensorflow::io::JoinPath(root_dir, graph);
+  Status load_graph_status = LoadGraph(graph_path, &session);
+  if (!load_graph_status.ok()) {
+    LOG(ERROR) << load_graph_status;
+    return -1;
+  }
+
+  // Get the image from disk as a float array of numbers, resized and normalized
+  // to the specifications the main graph expects.
+  std::vector<Tensor> resized_tensors;
+  string image_path = tensorflow::io::JoinPath(root_dir, image);
+  Status read_tensor_status =
+      ReadTensorFromImageFile(image_path, input_height, input_width, input_mean,
+                              input_std, &resized_tensors);
+  if (!read_tensor_status.ok()) {
+    LOG(ERROR) << read_tensor_status;
+    return -1;
+  }
+  const Tensor& resized_tensor = resized_tensors[0];
+
+  // Actually run the image through the model.
+  std::vector<Tensor> outputs;
+  Status run_status = session->Run({{input_layer, resized_tensor}},
+                                   {output_layer}, {}, &outputs);
+  if (!run_status.ok()) {
+    LOG(ERROR) << "Running model failed: " << run_status;
+    return -1;
+  }
+
+  // This is for automated testing to make sure we get the expected result with
+  // the default settings. We know that label 653 (military uniform) should be
+  // the top label for the Admiral Hopper image.
+  if (self_test) {
+    bool expected_matches;
+    Status check_status = CheckTopLabel(outputs, 653, &expected_matches);
+    if (!check_status.ok()) {
+      LOG(ERROR) << "Running check failed: " << check_status;
+      return -1;
+    }
+    if (!expected_matches) {
+      LOG(ERROR) << "Self-test failed!";
+      return -1;
+    }
+  }
+
+  // Do something interesting with the results we've generated.
+  Status print_status = PrintTopLabels(outputs, labels);
+  if (!print_status.ok()) {
+    LOG(ERROR) << "Running print failed: " << print_status;
+    return -1;
+  }
+
+  return 0;
+}

tensorflow/examples/multibox_detector/BUILD

@@ -0,0 +1,36 @@
+# Description:
+#   TensorFlow C++ inference example for labeling images.
+
+load("//tensorflow:tensorflow.bzl", "tf_cc_binary")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    licenses = ["notice"],
+)
+
+tf_cc_binary(
+    name = "detect_objects",
+    srcs = [
+        "main.cc",
+    ],
+    linkopts = ["-lm"],
+    deps = [
+        "//tensorflow/cc:cc_ops",
+        "//tensorflow/cc:ops",
+        "//tensorflow/cc:scope",
+        "//tensorflow/core:core_cpu",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:framework_internal",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core:tensorflow",
+        "//tensorflow/core/platform:numbers",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
+        "@xla//xla/tsl/platform:status",
+        "@xla//xla/tsl/platform:types",
+        "@xla//xla/tsl/util:command_line_flags",
+    ],
+)

tensorflow/examples/multibox_detector/README.md

@@ -0,0 +1,71 @@
+# TensorFlow C++ MultiBox Object Detection Demo
+
+This example shows how you can load a pre-trained TensorFlow network and use it
+to detect objects in images in C++. For an alternate implementation see the
+[Android TensorFlow demo](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/examples/android)
+
+## Description
+
+This demo uses a model based on [Scalable Object Detection using Deep NeuralNetworks](https://arxiv.org/abs/1312.2249) to detect people in images passed in from
+the command line. This is the same model also used in the Android TensorFlow
+demo for real-time person detection and tracking in the camera preview.
+
+## To build/install/run
+
+The TensorFlow `GraphDef` that contains the model definition and weights is not
+packaged in the repo because of its size. Instead, you must first download the
+file to the `data` directory in the source tree:
+
+```bash
+$ wget https://storage.googleapis.com/download.tensorflow.org/models/mobile_multibox_v1a.zip -O tensorflow/examples/multibox_detector/data/mobile_multibox_v1a.zip
+
+$ unzip tensorflow/examples/multibox_detector/data/mobile_multibox_v1a.zip -d tensorflow/examples/multibox_detector/data/
+```
+
+Then, as long as you've managed to build the main TensorFlow framework, you
+should have everything you need to run this example installed already.
+
+Once extracted, see the box priors file in the data directory. This file
+contains means and standard deviations for all 784 possible detections,
+normalized from 0-1 in left top right bottom order.
+
+To build it, run this command:
+
+```bash
+$ bazel build --config opt tensorflow/examples/multibox_detector/...
+```
+
+That should build a binary executable that you can then run like this:
+
+```bash
+$ bazel-bin/tensorflow/examples/multibox_detector/detect_objects --image_out=$HOME/x20/surfers_labeled.png
+```
+
+This uses the default example image that ships with the framework, and should
+output something similar to this:
+
+```
+I0125 18:24:13.804047    8677 main.cc:293] ===== Top 5 Detections ======
+I0125 18:24:13.804058    8677 main.cc:307] Detection 0: L:324.542 T:76.5764 R:373.26 B:214.957 (635) score: 0.267425
+I0125 18:24:13.804077    8677 main.cc:307] Detection 1: L:332.896 T:76.2751 R:372.116 B:204.614 (523) score: 0.245334
+I0125 18:24:13.804087    8677 main.cc:307] Detection 2: L:306.605 T:76.2228 R:371.356 B:217.32 (634) score: 0.216121
+I0125 18:24:13.804096    8677 main.cc:307] Detection 3: L:143.918 T:86.0909 R:187.333 B:195.885 (387) score: 0.171368
+I0125 18:24:13.804104    8677 main.cc:307] Detection 4: L:144.915 T:86.2675 R:185.243 B:165.246 (219) score: 0.169244
+```
+
+In this case, we're using a public domain stock image of surfers walking on the
+beach, and the top two few detections are of the two on the right. Adding more
+detections with --num_detections=N will also include the surfer on the left,
+and eventually non-person boxes below a certain threshold.
+
+You can visually inspect the detections by viewing the resulting png file
+'~/surfers_labeled.png'.
+
+Next, try it out on your own images by supplying the --image= argument, e.g.
+
+```bash
+$ bazel-bin/tensorflow/examples/multibox_detector/detect_objects --image=my_image.png
+```
+
+For another implementation of this work, you can check out the [Android
+TensorFlow demo](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/examples/android).

tensorflow/examples/multibox_detector/main.cc

@@ -0,0 +1,454 @@
+/* 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 <stdio.h>
+#include <string.h>
+
+#include <algorithm>
+#include <cmath>
+#include <cstdint>
+#include <fstream>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "absl/log/check.h"
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "absl/strings/match.h"
+#include "absl/strings/numbers.h"
+#include "tensorflow/cc/framework/ops.h"
+#include "tensorflow/cc/framework/scope.h"
+#include "tensorflow/cc/ops/array_ops.h"
+#include "tensorflow/cc/ops/const_op.h"
+#include "tensorflow/cc/ops/image_ops.h"
+#include "tensorflow/cc/ops/io_ops.h"
+#include "tensorflow/cc/ops/math_ops.h"
+#include "tensorflow/cc/ops/nn_ops.h"
+#include "xla/tsl/platform/env.h"
+#include "xla/tsl/platform/errors.h"
+#include "xla/tsl/platform/types.h"
+#include "xla/tsl/util/command_line_flags.h"
+#include "tensorflow/core/framework/graph.pb.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_types.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/lib/core/stringpiece.h"
+#include "tensorflow/core/lib/strings/str_util.h"
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/init_main.h"
+#include "tensorflow/core/platform/logging.h"
+#include "tensorflow/core/platform/numbers.h"
+#include "tensorflow/core/platform/path.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/public/session.h"
+#include "tensorflow/core/public/session_options.h"
+#include "tensorflow/core/util/command_line_flags.h"
+
+// These are all common classes it's handy to reference with no namespace.
+using tensorflow::Flag;
+using tensorflow::Tensor;
+using tensorflow::Status;
+using tensorflow::string;
+using tensorflow::int32;
+using tensorflow::uint8;
+
+// Takes a file name, and loads a list of comma-separated box priors from it,
+// one per line, and returns a vector of the values.
+Status ReadLocationsFile(const string& file_name, std::vector<float>* result,
+                         size_t* found_label_count) {
+  std::ifstream file(file_name);
+  if (!file) {
+    return absl::NotFoundError(
+        absl::StrCat("Labels file ", file_name, " not found."));
+  }
+  result->clear();
+  string line;
+  while (std::getline(file, line)) {
+    std::vector<string> string_tokens = tensorflow::str_util::Split(line, ',');
+    result->reserve(string_tokens.size());
+    for (const string& string_token : string_tokens) {
+      float number;
+      CHECK(absl::SimpleAtof(string_token, &number));
+      result->push_back(number);
+    }
+  }
+  *found_label_count = result->size();
+  return absl::OkStatus();
+}
+
+// Given an image file name, read in the data, try to decode it as an image,
+// resize it to the requested size, and then scale the values as desired.
+Status ReadTensorFromImageFile(const string& file_name, const int input_height,
+                               const int input_width, const float input_mean,
+                               const float input_std,
+                               std::vector<Tensor>* out_tensors) {
+  auto root = tensorflow::Scope::NewRootScope();
+  using namespace ::tensorflow::ops;  // NOLINT(build/namespaces)
+
+  string input_name = "file_reader";
+  string original_name = "identity";
+  string output_name = "normalized";
+  auto file_reader =
+      tensorflow::ops::ReadFile(root.WithOpName(input_name), file_name);
+  // Now try to figure out what kind of file it is and decode it.
+  const int wanted_channels = 3;
+  tensorflow::Output image_reader;
+  if (absl::EndsWith(file_name, ".png")) {
+    image_reader = DecodePng(root.WithOpName("png_reader"), file_reader,
+                             DecodePng::Channels(wanted_channels));
+  } else if (absl::EndsWith(file_name, ".gif")) {
+    image_reader = DecodeGif(root.WithOpName("gif_reader"), file_reader);
+  } else {
+    // Assume if it's neither a PNG nor a GIF then it must be a JPEG.
+    image_reader = DecodeJpeg(root.WithOpName("jpeg_reader"), file_reader,
+                              DecodeJpeg::Channels(wanted_channels));
+  }
+
+  // Also return identity so that we can know the original dimensions and
+  // optionally save the image out with bounding boxes overlaid.
+  auto original_image = Identity(root.WithOpName(original_name), image_reader);
+
+  // Now cast the image data to float so we can do normal math on it.
+  auto float_caster = Cast(root.WithOpName("float_caster"), original_image,
+                           tensorflow::DT_FLOAT);
+  // The convention for image ops in TensorFlow is that all images are expected
+  // to be in batches, so that they're four-dimensional arrays with indices of
+  // [batch, height, width, channel]. Because we only have a single image, we
+  // have to add a batch dimension of 1 to the start with ExpandDims().
+  auto dims_expander = ExpandDims(root, float_caster, 0);
+
+  // Bilinearly resize the image to fit the required dimensions.
+  auto resized = ResizeBilinear(
+      root, dims_expander,
+      Const(root.WithOpName("size"), {input_height, input_width}));
+  // Subtract the mean and divide by the scale.
+  Div give_me_a_name(root.WithOpName(output_name),
+                     Sub(root, resized, {input_mean}), {input_std});
+
+  // This runs the GraphDef network definition that we've just constructed, and
+  // returns the results in the output tensor.
+  tensorflow::GraphDef graph;
+  TF_RETURN_IF_ERROR(root.ToGraphDef(&graph));
+
+  std::unique_ptr<tensorflow::Session> session(
+      tensorflow::NewSession(tensorflow::SessionOptions()));
+  TF_RETURN_IF_ERROR(session->Create(graph));
+  TF_RETURN_IF_ERROR(
+      session->Run({}, {output_name, original_name}, {}, out_tensors));
+  return absl::OkStatus();
+}
+
+Status SaveImage(const Tensor& tensor, const string& file_path) {
+  LOG(INFO) << "Saving image to " << file_path;
+  CHECK(absl::EndsWith(file_path, ".png"))
+      << "Only saving of png files is supported.";
+
+  auto root = tensorflow::Scope::NewRootScope();
+  using namespace ::tensorflow::ops;  // NOLINT(build/namespaces)
+
+  string encoder_name = "encode";
+  string output_name = "file_writer";
+
+  tensorflow::Output image_encoder =
+      EncodePng(root.WithOpName(encoder_name), tensor);
+  tensorflow::ops::WriteFile file_saver = tensorflow::ops::WriteFile(
+      root.WithOpName(output_name), file_path, image_encoder);
+
+  tensorflow::GraphDef graph;
+  TF_RETURN_IF_ERROR(root.ToGraphDef(&graph));
+
+  std::unique_ptr<tensorflow::Session> session(
+      tensorflow::NewSession(tensorflow::SessionOptions()));
+  TF_RETURN_IF_ERROR(session->Create(graph));
+  std::vector<Tensor> outputs;
+  TF_RETURN_IF_ERROR(session->Run({}, {}, {output_name}, &outputs));
+
+  return absl::OkStatus();
+}
+
+// Reads a model graph definition from disk, and creates a session object you
+// can use to run it.
+Status LoadGraph(const string& graph_file_name,
+                 std::unique_ptr<tensorflow::Session>* session) {
+  tensorflow::GraphDef graph_def;
+  Status load_graph_status =
+      ReadBinaryProto(tensorflow::Env::Default(), graph_file_name, &graph_def);
+  if (!load_graph_status.ok()) {
+    return absl::NotFoundError(absl::StrCat("Failed to load compute graph at '",
+                                            graph_file_name, "'"));
+  }
+  session->reset(tensorflow::NewSession(tensorflow::SessionOptions()));
+  Status session_create_status = (*session)->Create(graph_def);
+  if (!session_create_status.ok()) {
+    return session_create_status;
+  }
+  return absl::OkStatus();
+}
+
+// Analyzes the output of the MultiBox graph to retrieve the highest scores and
+// their positions in the tensor, which correspond to individual box detections.
+Status GetTopDetections(const std::vector<Tensor>& outputs, int how_many_labels,
+                        Tensor* indices, Tensor* scores) {
+  auto root = tensorflow::Scope::NewRootScope();
+  using namespace ::tensorflow::ops;  // NOLINT(build/namespaces)
+
+  string output_name = "top_k";
+  TopK give_me_a_name(root.WithOpName(output_name), outputs[0],
+                      how_many_labels);
+  // This runs the GraphDef network definition that we've just constructed, and
+  // returns the results in the output tensors.
+  tensorflow::GraphDef graph;
+  TF_RETURN_IF_ERROR(root.ToGraphDef(&graph));
+
+  std::unique_ptr<tensorflow::Session> session(
+      tensorflow::NewSession(tensorflow::SessionOptions()));
+  TF_RETURN_IF_ERROR(session->Create(graph));
+  // The TopK node returns two outputs, the scores and their original indices,
+  // so we have to append :0 and :1 to specify them both.
+  std::vector<Tensor> out_tensors;
+  TF_RETURN_IF_ERROR(session->Run({}, {output_name + ":0", output_name + ":1"},
+                                  {}, &out_tensors));
+  *scores = out_tensors[0];
+  *indices = out_tensors[1];
+  return absl::OkStatus();
+}
+
+// Converts an encoded location to an actual box placement with the provided
+// box priors.
+void DecodeLocation(const float* encoded_location, const float* box_priors,
+                    float* decoded_location) {
+  bool non_zero = false;
+  for (int i = 0; i < 4; ++i) {
+    const float curr_encoding = encoded_location[i];
+    non_zero = non_zero || curr_encoding != 0.0f;
+
+    const float mean = box_priors[i * 2];
+    const float std_dev = box_priors[i * 2 + 1];
+
+    float currentLocation = curr_encoding * std_dev + mean;
+
+    currentLocation = std::max(currentLocation, 0.0f);
+    currentLocation = std::min(currentLocation, 1.0f);
+    decoded_location[i] = currentLocation;
+  }
+
+  if (!non_zero) {
+    LOG(WARNING) << "No non-zero encodings; check log for inference errors.";
+  }
+}
+
+float DecodeScore(float encoded_score) {
+  return 1 / (1 + std::exp(-encoded_score));
+}
+
+void DrawBox(const int image_width, const int image_height, int left, int top,
+             int right, int bottom, tensorflow::TTypes<uint8>::Flat* image) {
+  tensorflow::TTypes<uint8>::Flat image_ref = *image;
+
+  top = std::max(0, std::min(image_height - 1, top));
+  bottom = std::max(0, std::min(image_height - 1, bottom));
+
+  left = std::max(0, std::min(image_width - 1, left));
+  right = std::max(0, std::min(image_width - 1, right));
+
+  for (int i = 0; i < 3; ++i) {
+    uint8 val = i == 2 ? 255 : 0;
+    for (int x = left; x <= right; ++x) {
+      image_ref((top * image_width + x) * 3 + i) = val;
+      image_ref((bottom * image_width + x) * 3 + i) = val;
+    }
+    for (int y = top; y <= bottom; ++y) {
+      image_ref((y * image_width + left) * 3 + i) = val;
+      image_ref((y * image_width + right) * 3 + i) = val;
+    }
+  }
+}
+
+// Given the output of a model run, and the name of a file containing the labels
+// this prints out the top five highest-scoring values.
+Status PrintTopDetections(const std::vector<Tensor>& outputs,
+                          const string& labels_file_name,
+                          const int num_boxes,
+                          const int num_detections,
+                          const string& image_file_name,
+                          Tensor* original_tensor) {
+  std::vector<float> locations;
+  size_t label_count;
+  Status read_labels_status =
+      ReadLocationsFile(labels_file_name, &locations, &label_count);
+  if (!read_labels_status.ok()) {
+    LOG(ERROR) << read_labels_status;
+    return read_labels_status;
+  }
+  CHECK_EQ(label_count, num_boxes * 8);
+
+  const int how_many_labels =
+      std::min(num_detections, static_cast<int>(label_count));
+  Tensor indices;
+  Tensor scores;
+  TF_RETURN_IF_ERROR(
+      GetTopDetections(outputs, how_many_labels, &indices, &scores));
+
+  tensorflow::TTypes<float>::Flat scores_flat = scores.flat<float>();
+
+  tensorflow::TTypes<int32>::Flat indices_flat = indices.flat<int32>();
+
+  const Tensor& encoded_locations = outputs[1];
+  auto locations_encoded = encoded_locations.flat<float>();
+
+  LOG(INFO) << original_tensor->DebugString();
+  const int image_width = original_tensor->shape().dim_size(1);
+  const int image_height = original_tensor->shape().dim_size(0);
+
+  tensorflow::TTypes<uint8>::Flat image_flat = original_tensor->flat<uint8>();
+
+  LOG(INFO) << "===== Top " << how_many_labels << " Detections ======";
+  for (int pos = 0; pos < how_many_labels; ++pos) {
+    const int label_index = indices_flat(pos);
+    const float score = scores_flat(pos);
+
+    float decoded_location[4];
+    DecodeLocation(&locations_encoded(label_index * 4),
+                   &locations[label_index * 8], decoded_location);
+
+    float left = decoded_location[0] * image_width;
+    float top = decoded_location[1] * image_height;
+    float right = decoded_location[2] * image_width;
+    float bottom = decoded_location[3] * image_height;
+
+    LOG(INFO) << "Detection " << pos << ": "
+              << "L:" << left << " "
+              << "T:" << top << " "
+              << "R:" << right << " "
+              << "B:" << bottom << " "
+              << "(" << label_index << ") score: " << DecodeScore(score);
+
+    DrawBox(image_width, image_height, left, top, right, bottom, &image_flat);
+  }
+
+  if (!image_file_name.empty()) {
+    return SaveImage(*original_tensor, image_file_name);
+  }
+  return absl::OkStatus();
+}
+
+int main(int argc, char* argv[]) {
+  // These are the command-line flags the program can understand.
+  // They define where the graph and input data is located, and what kind of
+  // input the model expects. If you train your own model, or use something
+  // other than multibox_model you'll need to update these.
+  string image =
+      "tensorflow/examples/multibox_detector/data/surfers.jpg";
+  string graph =
+      "tensorflow/examples/multibox_detector/data/"
+      "multibox_model.pb";
+  string box_priors =
+      "tensorflow/examples/multibox_detector/data/"
+      "multibox_location_priors.txt";
+  int32_t input_width = 224;
+  int32_t input_height = 224;
+  int32_t input_mean = 128;
+  int32_t input_std = 128;
+  int32_t num_detections = 5;
+  int32_t num_boxes = 784;
+  string input_layer = "ResizeBilinear";
+  string output_location_layer = "output_locations/Reshape";
+  string output_score_layer = "output_scores/Reshape";
+  string root_dir = "";
+  string image_out = "";
+
+  std::vector<Flag> flag_list = {
+      Flag("image", &image, "image to be processed"),
+      Flag("image_out", &image_out,
+           "location to save output image, if desired"),
+      Flag("graph", &graph, "graph to be executed"),
+      Flag("box_priors", &box_priors, "name of file containing box priors"),
+      Flag("input_width", &input_width, "resize image to this width in pixels"),
+      Flag("input_height", &input_height,
+           "resize image to this height in pixels"),
+      Flag("input_mean", &input_mean, "scale pixel values to this mean"),
+      Flag("input_std", &input_std, "scale pixel values to this std deviation"),
+      Flag("num_detections", &num_detections,
+           "number of top detections to return"),
+      Flag("num_boxes", &num_boxes,
+           "number of boxes defined by the location file"),
+      Flag("input_layer", &input_layer, "name of input layer"),
+      Flag("output_location_layer", &output_location_layer,
+           "name of location output layer"),
+      Flag("output_score_layer", &output_score_layer,
+           "name of score output layer"),
+      Flag("root_dir", &root_dir,
+           "interpret image and graph file names relative to this directory"),
+  };
+
+  string usage = tensorflow::Flags::Usage(argv[0], flag_list);
+  const bool parse_result = tensorflow::Flags::Parse(&argc, argv, flag_list);
+  if (!parse_result) {
+    LOG(ERROR) << usage;
+    return -1;
+  }
+
+  // We need to call this to set up global state for TensorFlow.
+  tensorflow::port::InitMain(argv[0], &argc, &argv);
+  if (argc > 1) {
+    LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage;
+    return -1;
+  }
+
+  // First we load and initialize the model.
+  std::unique_ptr<tensorflow::Session> session;
+  string graph_path = tensorflow::io::JoinPath(root_dir, graph);
+  Status load_graph_status = LoadGraph(graph_path, &session);
+  if (!load_graph_status.ok()) {
+    LOG(ERROR) << load_graph_status;
+    return -1;
+  }
+
+  // Get the image from disk as a float array of numbers, resized and normalized
+  // to the specifications the main graph expects.
+  std::vector<Tensor> image_tensors;
+  string image_path = tensorflow::io::JoinPath(root_dir, image);
+
+  Status read_tensor_status =
+      ReadTensorFromImageFile(image_path, input_height, input_width, input_mean,
+                              input_std, &image_tensors);
+  if (!read_tensor_status.ok()) {
+    LOG(ERROR) << read_tensor_status;
+    return -1;
+  }
+  const Tensor& resized_tensor = image_tensors[0];
+
+  // Actually run the image through the model.
+  std::vector<Tensor> outputs;
+  Status run_status =
+      session->Run({{input_layer, resized_tensor}},
+                   {output_score_layer, output_location_layer}, {}, &outputs);
+  if (!run_status.ok()) {
+    LOG(ERROR) << "Running model failed: " << run_status;
+    return -1;
+  }
+
+  Status print_status = PrintTopDetections(outputs, box_priors, num_boxes,
+                                           num_detections, image_out,
+                                           &image_tensors[1]);
+
+  if (!print_status.ok()) {
+    LOG(ERROR) << "Running print failed: " << print_status;
+    return -1;
+  }
+  return 0;
+}

tensorflow/examples/speech_commands/BUILD

@@ -0,0 +1,463 @@
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+load("@xla//third_party/rules_python/python:py_binary.bzl", "py_binary")
+load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
+load("//tensorflow:tensorflow.bzl", "tf_cc_binary", "tf_cc_test")
+load("//tensorflow:tensorflow.default.bzl", "tf_py_strict_test")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = [
+        "//visibility:public",
+    ],
+    licenses = ["notice"],
+)
+
+py_binary(
+    name = "accuracy_utils_py",
+    srcs = ["accuracy_utils.py"],
+    main = "accuracy_utils.py",
+    strict_deps = True,
+    deps = [
+        "//tensorflow:tensorflow_py",
+        "//third_party/py/numpy",
+    ],
+)
+
+py_library(
+    name = "accuracy_utils_lib",
+    srcs = ["accuracy_utils.py"],
+    strict_deps = True,
+    deps = [
+        "//tensorflow:tensorflow_py",
+        "//third_party/py/numpy",
+    ],
+)
+
+py_binary(
+    name = "recognize_commands_py",
+    srcs = ["recognize_commands.py"],
+    main = "recognize_commands.py",
+    strict_deps = True,
+    deps = ["//third_party/py/numpy"],
+)
+
+py_library(
+    name = "recognize_commands_lib",
+    srcs = ["recognize_commands.py"],
+    strict_deps = True,
+    deps = ["//third_party/py/numpy"],
+)
+
+py_binary(
+    name = "test_streaming_accuracy_py",
+    srcs = ["test_streaming_accuracy.py"],
+    main = "test_streaming_accuracy.py",
+    strict_deps = True,
+    deps = [
+        ":accuracy_utils_lib",
+        ":recognize_commands_lib",
+        "//tensorflow:tensorflow_py",
+        "//tensorflow/python/ops:io_ops",
+        "//third_party/py/numpy",
+    ],
+)
+
+py_library(
+    name = "models",
+    srcs = [
+        "models.py",
+    ],
+    strict_deps = True,
+    deps = ["//tensorflow:tensorflow_py"],
+)
+
+tf_py_strict_test(
+    name = "models_test",
+    size = "medium",
+    srcs = ["models_test.py"],
+    deps = [
+        ":models",
+        "//tensorflow:tensorflow_py_no_contrib",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+py_library(
+    name = "input_data",
+    srcs = [
+        "input_data.py",
+    ],
+    strict_deps = True,
+    deps = [
+        "//tensorflow:tensorflow_py",
+        "//tensorflow/python/ops:audio_ops_gen",
+        "//tensorflow/python/ops:io_ops",
+        "//tensorflow/python/platform:gfile",
+        "//tensorflow/python/util:compat",
+        "//third_party/py/numpy",
+    ],
+)
+
+tf_py_strict_test(
+    name = "input_data_test",
+    size = "small",
+    srcs = ["input_data_test.py"],
+    tags = [
+        "v1only",  # uses contrib
+    ],
+    deps = [
+        ":input_data",
+        ":models",
+        "//tensorflow:tensorflow_py_no_contrib",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+py_binary(
+    name = "train",
+    srcs = ["train.py"],
+    strict_deps = True,
+    deps = [":train_main_lib"],
+)
+
+py_library(
+    name = "train_main_lib",
+    srcs = [
+        "train.py",
+    ],
+    strict_deps = True,
+    deps = [
+        ":input_data",
+        ":models",
+        "//tensorflow:tensorflow_py",
+        "//tensorflow/python/platform:gfile",
+        "//third_party/py/numpy",
+    ],
+)
+
+tf_py_strict_test(
+    name = "train_test",
+    size = "small",
+    srcs = ["train_test.py"],
+    tags = [
+        "v1only",  # uses contrib
+    ],
+    deps = [
+        ":train_main_lib",
+        "//tensorflow:tensorflow_py_no_contrib",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/platform:gfile",
+    ],
+)
+
+py_binary(
+    name = "freeze",
+    srcs = ["freeze.py"],
+    strict_deps = True,
+    deps = [":freeze_main_lib"],
+)
+
+py_library(
+    name = "freeze_main_lib",
+    srcs = ["freeze.py"],
+    strict_deps = True,
+    deps = [
+        ":input_data",
+        ":models",
+        "//tensorflow:tensorflow_py_no_contrib",
+        "//tensorflow/python/ops:audio_ops_gen",
+    ],
+)
+
+py_library(
+    name = "freeze_lib",
+    srcs = [
+        "freeze.py",
+    ],
+    strict_deps = True,
+    tags = [
+        "no_pip",  # b/131330719
+    ],
+    deps = [
+        ":input_data",
+        ":models",
+        "//tensorflow:tensorflow_py",
+        "//tensorflow/python/ops:audio_ops_gen",
+    ],
+)
+
+tf_py_strict_test(
+    name = "freeze_test",
+    size = "small",
+    srcs = ["freeze_test.py"],
+    tags = [
+        "v1only",  # uses contrib
+    ],
+    deps = [
+        ":freeze_main_lib",
+        "//tensorflow/python/framework:convert_to_constants",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+py_binary(
+    name = "wav_to_features",
+    srcs = ["wav_to_features.py"],
+    strict_deps = True,
+    deps = [":wav_to_features_main_lib"],
+)
+
+py_library(
+    name = "wav_to_features_main_lib",
+    srcs = ["wav_to_features.py"],
+    strict_deps = True,
+    deps = [
+        ":input_data",
+        ":models",
+        "//tensorflow:tensorflow_py_no_contrib",
+        "//tensorflow/python/platform:gfile",
+    ],
+)
+
+py_library(
+    name = "wav_to_features_lib",
+    srcs = [
+        "wav_to_features.py",
+    ],
+    strict_deps = True,
+    deps = [
+        ":input_data",
+        ":models",
+        "//tensorflow:tensorflow_py",
+        "//tensorflow/python/platform:gfile",
+    ],
+)
+
+tf_py_strict_test(
+    name = "wav_to_features_test",
+    size = "small",
+    srcs = ["wav_to_features_test.py"],
+    tags = [
+        "v1only",  # uses contrib
+    ],
+    deps = [
+        ":wav_to_features_main_lib",
+        "//tensorflow:tensorflow_py_no_contrib",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+py_binary(
+    name = "generate_streaming_test_wav",
+    srcs = ["generate_streaming_test_wav.py"],
+    strict_deps = True,
+    deps = [":generate_streaming_test_wav_main_lib"],
+)
+
+py_library(
+    name = "generate_streaming_test_wav_main_lib",
+    srcs = ["generate_streaming_test_wav.py"],
+    strict_deps = True,
+    deps = [
+        ":input_data",
+        ":models",
+        "//tensorflow:tensorflow_py_no_contrib",
+        "//third_party/py/numpy",
+    ],
+)
+
+py_library(
+    name = "generate_streaming_test_wav_lib",
+    srcs = [
+        "generate_streaming_test_wav.py",
+    ],
+    strict_deps = True,
+    deps = [
+        ":input_data",
+        ":models",
+        "//tensorflow:tensorflow_py",
+        "//third_party/py/numpy",
+    ],
+)
+
+tf_py_strict_test(
+    name = "generate_streaming_test_wav_test",
+    size = "small",
+    srcs = ["generate_streaming_test_wav_test.py"],
+    tags = [
+        "v1only",  # uses contrib
+    ],
+    deps = [
+        ":generate_streaming_test_wav_main_lib",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+tf_cc_binary(
+    name = "label_wav_cc",
+    srcs = [
+        "label_wav.cc",
+    ],
+    deps = [
+        "//tensorflow/core:core_cpu",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:framework_internal",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core:tensorflow",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/status",
+        "@xla//xla/tsl/platform:status",
+        "@xla//xla/tsl/platform:types",
+        "@xla//xla/tsl/util:command_line_flags",
+    ],
+)
+
+py_binary(
+    name = "label_wav",
+    srcs = ["label_wav.py"],
+    strict_deps = True,
+    deps = [":label_wav_main_lib"],
+)
+
+py_library(
+    name = "label_wav_main_lib",
+    srcs = ["label_wav.py"],
+    strict_deps = True,
+    deps = ["//tensorflow:tensorflow_py_no_contrib"],
+)
+
+py_library(
+    name = "label_wav_lib",
+    srcs = [
+        "label_wav.py",
+    ],
+    strict_deps = True,
+    deps = [
+        "//tensorflow:tensorflow_py",
+    ],
+)
+
+tf_py_strict_test(
+    name = "label_wav_test",
+    size = "medium",
+    srcs = ["label_wav_test.py"],
+    tags = [
+        "v1only",  # uses contrib
+    ],
+    deps = [
+        ":label_wav_main_lib",
+        "//tensorflow:tensorflow_py_no_contrib",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+cc_library(
+    name = "recognize_commands",
+    srcs = [
+        "recognize_commands.cc",
+    ],
+    hdrs = [
+        "recognize_commands.h",
+    ],
+    deps = [
+        "//tensorflow/core:framework",
+        "//tensorflow/core:framework_internal",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:tensorflow",
+        "@com_google_absl//absl/status",
+    ],
+)
+
+tf_cc_test(
+    name = "recognize_commands_test",
+    size = "medium",
+    srcs = [
+        "recognize_commands_test.cc",
+    ],
+    deps = [
+        ":recognize_commands",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:lib_internal",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+        "//tensorflow/core:testlib",
+    ],
+)
+
+cc_library(
+    name = "accuracy_utils",
+    srcs = [
+        "accuracy_utils.cc",
+    ],
+    hdrs = [
+        "accuracy_utils.h",
+    ],
+    deps = [
+        "//tensorflow/core:framework",
+        "//tensorflow/core:framework_internal",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:tensorflow",
+        "//tensorflow/core/platform:numbers",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
+    ],
+)
+
+tf_cc_test(
+    name = "accuracy_utils_test",
+    size = "medium",
+    srcs = [
+        "accuracy_utils_test.cc",
+    ],
+    deps = [
+        ":accuracy_utils",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:lib_internal",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+    ],
+)
+
+tf_cc_binary(
+    name = "test_streaming_accuracy",
+    srcs = [
+        "test_streaming_accuracy.cc",
+    ],
+    deps = [
+        ":accuracy_utils",
+        ":recognize_commands",
+        "//tensorflow/core:core_cpu",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:framework_internal",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:lib_internal",
+        "//tensorflow/core:protos_all_cc",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/status",
+        "@xla//xla/tsl/platform:types",
+    ],
+)
+
+py_library(
+    name = "test_lib",
+    strict_deps = True,
+    deps = [
+        ":freeze_main_lib",
+        ":generate_streaming_test_wav_main_lib",
+        ":input_data",
+        ":label_wav_main_lib",
+        ":models",
+        ":train_main_lib",
+        ":wav_to_features_main_lib",
+    ],
+)

tensorflow/examples/speech_commands/README.md

@@ -0,0 +1,4 @@
+# Speech Commands Example
+
+This is a basic speech recognition example. For more information, see the
+tutorial at https://www.tensorflow.org/tutorials/audio/simple_audio.

tensorflow/examples/speech_commands/accuracy_utils.cc

@@ -0,0 +1,153 @@
+/* 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/examples/speech_commands/accuracy_utils.h"
+
+#include <algorithm>
+#include <cstdint>
+#include <fstream>
+#include <iomanip>
+#include <ios>
+#include <limits>
+#include <string>
+#include <unordered_set>
+#include <utility>
+#include <vector>
+
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "absl/strings/numbers.h"
+#include "tensorflow/core/lib/strings/str_util.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/numbers.h"
+#include "tensorflow/core/platform/types.h"
+
+namespace tensorflow {
+
+absl::Status ReadGroundTruthFile(
+    const std::string& file_name,
+    std::vector<std::pair<std::string, int64_t>>* result) {
+  std::ifstream file(file_name);
+  if (!file) {
+    return absl::NotFoundError(
+        absl::StrCat("Ground truth file '", file_name, "' not found."));
+  }
+  result->clear();
+  std::string line;
+  while (std::getline(file, line)) {
+    std::vector<std::string> pieces = tensorflow::str_util::Split(line, ',');
+    if (pieces.size() != 2) {
+      continue;
+    }
+    float timestamp;
+    if (!absl::SimpleAtof(pieces[1], &timestamp)) {
+      return absl::InvalidArgumentError(
+          absl::StrCat("Wrong number format at line: ", line));
+    }
+    std::string label = pieces[0];
+    auto timestamp_int64 = static_cast<int64_t>(timestamp);
+    result->push_back({label, timestamp_int64});
+  }
+  std::sort(result->begin(), result->end(),
+            [](const std::pair<std::string, int64_t>& left,
+               const std::pair<std::string, int64_t>& right) {
+              return left.second < right.second;
+            });
+  return absl::OkStatus();
+}
+
+void CalculateAccuracyStats(
+    const std::vector<std::pair<std::string, int64_t>>& ground_truth_list,
+    const std::vector<std::pair<std::string, int64_t>>& found_words,
+    int64_t up_to_time_ms, int64_t time_tolerance_ms,
+    StreamingAccuracyStats* stats) {
+  int64_t latest_possible_time;
+  if (up_to_time_ms == -1) {
+    latest_possible_time = std::numeric_limits<int64_t>::max();
+  } else {
+    latest_possible_time = up_to_time_ms + time_tolerance_ms;
+  }
+  stats->how_many_ground_truth_words = 0;
+  for (const std::pair<std::string, int64_t>& ground_truth :
+       ground_truth_list) {
+    const int64_t ground_truth_time = ground_truth.second;
+    if (ground_truth_time > latest_possible_time) {
+      break;
+    }
+    ++stats->how_many_ground_truth_words;
+  }
+
+  stats->how_many_false_positives = 0;
+  stats->how_many_correct_words = 0;
+  stats->how_many_wrong_words = 0;
+  std::unordered_set<int64_t> has_ground_truth_been_matched;
+  for (const std::pair<std::string, int64_t>& found_word : found_words) {
+    const std::string& found_label = found_word.first;
+    const int64_t found_time = found_word.second;
+    const int64_t earliest_time = found_time - time_tolerance_ms;
+    const int64_t latest_time = found_time + time_tolerance_ms;
+    bool has_match_been_found = false;
+    for (const std::pair<std::string, int64_t>& ground_truth :
+         ground_truth_list) {
+      const int64_t ground_truth_time = ground_truth.second;
+      if ((ground_truth_time > latest_time) ||
+          (ground_truth_time > latest_possible_time)) {
+        break;
+      }
+      if (ground_truth_time < earliest_time) {
+        continue;
+      }
+      const std::string& ground_truth_label = ground_truth.first;
+      if ((ground_truth_label == found_label) &&
+          (has_ground_truth_been_matched.count(ground_truth_time) == 0)) {
+        ++stats->how_many_correct_words;
+      } else {
+        ++stats->how_many_wrong_words;
+      }
+      has_ground_truth_been_matched.insert(ground_truth_time);
+      has_match_been_found = true;
+      break;
+    }
+    if (!has_match_been_found) {
+      ++stats->how_many_false_positives;
+    }
+  }
+  stats->how_many_ground_truth_matched = has_ground_truth_been_matched.size();
+}
+
+void PrintAccuracyStats(const StreamingAccuracyStats& stats) {
+  if (stats.how_many_ground_truth_words == 0) {
+    LOG(INFO) << "No ground truth yet, " << stats.how_many_false_positives
+              << " false positives";
+  } else {
+    float any_match_percentage =
+        (stats.how_many_ground_truth_matched * 100.0f) /
+        stats.how_many_ground_truth_words;
+    float correct_match_percentage = (stats.how_many_correct_words * 100.0f) /
+                                     stats.how_many_ground_truth_words;
+    float wrong_match_percentage = (stats.how_many_wrong_words * 100.0f) /
+                                   stats.how_many_ground_truth_words;
+    float false_positive_percentage =
+        (stats.how_many_false_positives * 100.0f) /
+        stats.how_many_ground_truth_words;
+
+    LOG(INFO) << std::setprecision(1) << std::fixed << any_match_percentage
+              << "% matched, " << correct_match_percentage << "% correctly, "
+              << wrong_match_percentage << "% wrongly, "
+              << false_positive_percentage << "% false positives ";
+  }
+}
+
+}  // namespace tensorflow

tensorflow/examples/speech_commands/accuracy_utils.h

@@ -0,0 +1,65 @@
+/* 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_EXAMPLES_SPEECH_COMMANDS_ACCURACY_UTILS_H_
+#define TENSORFLOW_EXAMPLES_SPEECH_COMMANDS_ACCURACY_UTILS_H_
+
+#include <cstdint>
+#include <utility>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/platform/types.h"
+
+namespace tensorflow {
+
+struct StreamingAccuracyStats {
+  StreamingAccuracyStats()
+      : how_many_ground_truth_words(0),
+        how_many_ground_truth_matched(0),
+        how_many_false_positives(0),
+        how_many_correct_words(0),
+        how_many_wrong_words(0) {}
+  int32_t how_many_ground_truth_words;
+  int32_t how_many_ground_truth_matched;
+  int32_t how_many_false_positives;
+  int32_t how_many_correct_words;
+  int32_t how_many_wrong_words;
+};
+
+// Takes a file name, and loads a list of expected word labels and times from
+// it, as comma-separated variables.
+absl::Status ReadGroundTruthFile(
+    const std::string& file_name,
+    std::vector<std::pair<std::string, int64_t>>* result);
+
+// Given ground truth labels and corresponding predictions found by a model,
+// figure out how many were correct. Takes a time limit, so that only
+// predictions up to a point in time are considered, in case we're evaluating
+// accuracy when the model has only been run on part of the stream.
+void CalculateAccuracyStats(
+    const std::vector<std::pair<std::string, int64_t>>& ground_truth_list,
+    const std::vector<std::pair<std::string, int64_t>>& found_words,
+    int64_t up_to_time_ms, int64_t time_tolerance_ms,
+    StreamingAccuracyStats* stats);
+
+// Writes a human-readable description of the statistics to stdout.
+void PrintAccuracyStats(const StreamingAccuracyStats& stats);
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_EXAMPLES_SPEECH_COMMANDS_ACCURACY_UTILS_H_

tensorflow/examples/speech_commands/accuracy_utils.py

@@ -0,0 +1,150 @@
+# 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.
+# ==============================================================================
+"""Utils for getting accuracy statistics."""
+import numpy as np
+import tensorflow as tf
+
+
+class StreamingAccuracyStats(object):
+  """Get streaming accuracy statistics every time a new command is founded.
+
+    Attributes:
+      _how_many_gt: How many ground truths.
+      _how_many_gt_matched: How many ground truths have been matched.
+      _how_many_fp: How many commands have been fired as false positive.
+      _how_many_c: How many commands have been fired correctly.
+      _how_many_w: How many commands have been fired wrongly.
+      _gt_occurrence: A list to record which commands and when it occurs in the
+        input audio stream.
+      _previous_c: A variable to record the last status of _how_many_c.
+      _previous_w: A variable to record the last status of _how_many_w.
+      _previous_fp: A variable to record the last status of _how_many_fp.
+  """
+
+  def __init__(self):
+    """Init StreamingAccuracyStats with void or zero values."""
+    self._how_many_gt = 0
+    self._how_many_gt_matched = 0
+    self._how_many_fp = 0
+    self._how_many_c = 0
+    self._how_many_w = 0
+    self._gt_occurrence = []
+    self._previous_c = 0
+    self._previous_w = 0
+    self._previous_fp = 0
+
+  def read_ground_truth_file(self, file_name):
+    """Load ground truth and timestamp pairs and store it in time order."""
+    with open(file_name, 'r') as f:
+      for line in f:
+        line_split = line.strip().split(',')
+        if len(line_split) != 2:
+          continue
+        timestamp = round(float(line_split[1]))
+        label = line_split[0]
+        self._gt_occurrence.append([label, timestamp])
+    self._gt_occurrence = sorted(self._gt_occurrence, key=lambda item: item[1])
+
+  def delta(self):
+    """Compute delta of StreamingAccuracyStats against last status."""
+    fp_delta = self._how_many_fp - self._previous_fp
+    w_delta = self._how_many_w - self._previous_w
+    c_delta = self._how_many_c - self._previous_c
+    if fp_delta == 1:
+      recognition_state = '(False Positive)'
+    elif c_delta == 1:
+      recognition_state = '(Correct)'
+    elif w_delta == 1:
+      recognition_state = '(Wrong)'
+    else:
+      raise ValueError('Unexpected state in statistics')
+    # Update the previous status
+    self._previous_c = self._how_many_c
+    self._previous_w = self._how_many_w
+    self._previous_fp = self._how_many_fp
+    return recognition_state
+
+  def calculate_accuracy_stats(self, found_words, up_to_time_ms,
+                               time_tolerance_ms):
+    """Calculate accuracy statistics when a new commands is founded.
+
+    Given ground truth and corresponding predictions founded by
+    model, figure out how many were correct. Take a tolerance time, so that only
+    predictions up to a point in time are considered.
+
+    Args:
+        found_words: A list of all founded commands up to now.
+        up_to_time_ms: End timestamp of this audio piece.
+        time_tolerance_ms: The tolerance milliseconds before and after
+          up_to_time_ms to match a ground truth.
+    """
+    if up_to_time_ms == -1:
+      latest_possible_time = np.inf
+    else:
+      latest_possible_time = up_to_time_ms + time_tolerance_ms
+    self._how_many_gt = 0
+    for ground_truth in self._gt_occurrence:
+      ground_truth_time = ground_truth[1]
+      if ground_truth_time > latest_possible_time:
+        break
+      self._how_many_gt += 1
+    self._how_many_fp = 0
+    self._how_many_c = 0
+    self._how_many_w = 0
+    has_gt_matched = []
+    for found_word in found_words:
+      found_label = found_word[0]
+      found_time = found_word[1]
+      earliest_time = found_time - time_tolerance_ms
+      latest_time = found_time + time_tolerance_ms
+      has_matched_been_found = False
+      for ground_truth in self._gt_occurrence:
+        ground_truth_time = ground_truth[1]
+        if (ground_truth_time > latest_time or
+            ground_truth_time > latest_possible_time):
+          break
+        if ground_truth_time < earliest_time:
+          continue
+        ground_truth_label = ground_truth[0]
+        if (
+            ground_truth_label == found_label
+            and ground_truth_time not in has_gt_matched
+        ):
+          self._how_many_c += 1
+        else:
+          self._how_many_w += 1
+        has_gt_matched.append(ground_truth_time)
+        has_matched_been_found = True
+        break
+      if not has_matched_been_found:
+        self._how_many_fp += 1
+    self._how_many_gt_matched = len(has_gt_matched)
+
+  def print_accuracy_stats(self):
+    """Write a human-readable description of the statistics to stdout."""
+    if self._how_many_gt == 0:
+      tf.compat.v1.logging.info('No ground truth yet, {}false positives'.format(
+          self._how_many_fp))
+    else:
+      any_match_percentage = self._how_many_gt_matched / self._how_many_gt * 100
+      correct_match_percentage = self._how_many_c / self._how_many_gt * 100
+      wrong_match_percentage = self._how_many_w / self._how_many_gt * 100
+      false_positive_percentage = self._how_many_fp / self._how_many_gt * 100
+      tf.compat.v1.logging.info(
+          '{:.1f}% matched, {:.1f}% correct, {:.1f}% wrong, '
+          '{:.1f}% false positive'.format(any_match_percentage,
+                                          correct_match_percentage,
+                                          wrong_match_percentage,
+                                          false_positive_percentage))

tensorflow/examples/speech_commands/accuracy_utils_test.cc

@@ -0,0 +1,63 @@
+/* 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/examples/speech_commands/accuracy_utils.h"
+
+#include <cstdint>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "xla/tsl/lib/core/status_test_util.h"
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/path.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/platform/types.h"
+
+namespace tensorflow {
+
+TEST(AccuracyUtilsTest, ReadGroundTruthFile) {
+  std::string file_name = tensorflow::io::JoinPath(
+      tensorflow::testing::TmpDir(), "ground_truth.txt");
+  std::string file_data = "a,10\nb,12\n";
+  TF_ASSERT_OK(WriteStringToFile(Env::Default(), file_name, file_data));
+
+  std::vector<std::pair<std::string, int64_t>> ground_truth;
+  TF_ASSERT_OK(ReadGroundTruthFile(file_name, &ground_truth));
+  ASSERT_EQ(2, ground_truth.size());
+  EXPECT_EQ("a", ground_truth[0].first);
+  EXPECT_EQ(10, ground_truth[0].second);
+  EXPECT_EQ("b", ground_truth[1].first);
+  EXPECT_EQ(12, ground_truth[1].second);
+}
+
+TEST(AccuracyUtilsTest, CalculateAccuracyStats) {
+  StreamingAccuracyStats stats;
+  CalculateAccuracyStats({{"a", 1000}, {"b", 9000}},
+                         {{"a", 1200}, {"b", 5000}, {"a", 8700}}, 10000, 500,
+                         &stats);
+  EXPECT_EQ(2, stats.how_many_ground_truth_words);
+  EXPECT_EQ(2, stats.how_many_ground_truth_matched);
+  EXPECT_EQ(1, stats.how_many_false_positives);
+  EXPECT_EQ(1, stats.how_many_correct_words);
+  EXPECT_EQ(1, stats.how_many_wrong_words);
+}
+
+TEST(AccuracyUtilsTest, PrintAccuracyStats) {
+  StreamingAccuracyStats stats;
+  PrintAccuracyStats(stats);
+}
+
+}  // namespace tensorflow

tensorflow/examples/speech_commands/freeze.py

@@ -0,0 +1,307 @@
+# 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.
+# ==============================================================================
+r"""Converts a trained checkpoint into a frozen model for mobile inference.
+
+Once you've trained a model using the `train.py` script, you can use this tool
+to convert it into a binary GraphDef file that can be loaded into the Android,
+iOS, or Raspberry Pi example code. Here's an example of how to run it:
+
+bazel run tensorflow/examples/speech_commands/freeze -- \
+--sample_rate=16000 --dct_coefficient_count=40 --window_size_ms=20 \
+--window_stride_ms=10 --clip_duration_ms=1000 \
+--model_architecture=conv \
+--start_checkpoint=/tmp/speech_commands_train/conv.ckpt-1300 \
+--output_file=/tmp/my_frozen_graph.pb
+
+One thing to watch out for is that you need to pass in the same arguments for
+`sample_rate` and other command line variables here as you did for the training
+script.
+
+The resulting graph has an input for WAV-encoded data named 'wav_data', one for
+raw PCM data (as floats in the range -1.0 to 1.0) called 'decoded_sample_data',
+and the output is called 'labels_softmax'.
+
+"""
+import argparse
+import os.path
+import sys
+
+import tensorflow as tf
+
+import input_data
+import models
+from tensorflow.python.ops import gen_audio_ops as audio_ops
+
+# If it's available, load the specialized feature generator. If this doesn't
+# work, try building with bazel instead of running the Python script directly.
+# bazel run tensorflow/examples/speech_commands:freeze_graph
+try:
+  from tensorflow.lite.experimental.microfrontend.python.ops import audio_microfrontend_op as frontend_op  # pylint:disable=g-import-not-at-top
+except ImportError:
+  frontend_op = None
+
+FLAGS = None
+
+
+def create_inference_graph(wanted_words, sample_rate, clip_duration_ms,
+                           clip_stride_ms, window_size_ms, window_stride_ms,
+                           feature_bin_count, model_architecture, preprocess):
+  """Creates an audio model with the nodes needed for inference.
+
+  Uses the supplied arguments to create a model, and inserts the input and
+  output nodes that are needed to use the graph for inference.
+
+  Args:
+    wanted_words: Comma-separated list of the words we're trying to recognize.
+    sample_rate: How many samples per second are in the input audio files.
+    clip_duration_ms: How many samples to analyze for the audio pattern.
+    clip_stride_ms: How often to run recognition. Useful for models with cache.
+    window_size_ms: Time slice duration to estimate frequencies from.
+    window_stride_ms: How far apart time slices should be.
+    feature_bin_count: Number of frequency bands to analyze.
+    model_architecture: Name of the kind of model to generate.
+    preprocess: How the spectrogram is processed to produce features, for
+      example 'mfcc', 'average', or 'micro'.
+
+  Returns:
+    Input and output tensor objects.
+
+  Raises:
+    Exception: If the preprocessing mode isn't recognized.
+  """
+
+  words_list = input_data.prepare_words_list(wanted_words.split(','))
+  model_settings = models.prepare_model_settings(
+      len(words_list), sample_rate, clip_duration_ms, window_size_ms,
+      window_stride_ms, feature_bin_count, preprocess)
+  runtime_settings = {'clip_stride_ms': clip_stride_ms}
+
+  wav_data_placeholder = tf.compat.v1.placeholder(tf.string, [],
+                                                  name='wav_data')
+  decoded_sample_data = tf.audio.decode_wav(
+      wav_data_placeholder,
+      desired_channels=1,
+      desired_samples=model_settings['desired_samples'],
+      name='decoded_sample_data')
+  spectrogram = audio_ops.audio_spectrogram(
+      decoded_sample_data.audio,
+      window_size=model_settings['window_size_samples'],
+      stride=model_settings['window_stride_samples'],
+      magnitude_squared=True)
+
+  if preprocess == 'average':
+    fingerprint_input = tf.nn.pool(
+        input=tf.expand_dims(spectrogram, -1),
+        window_shape=[1, model_settings['average_window_width']],
+        strides=[1, model_settings['average_window_width']],
+        pooling_type='AVG',
+        padding='SAME')
+  elif preprocess == 'mfcc':
+    fingerprint_input = audio_ops.mfcc(
+        spectrogram,
+        sample_rate,
+        dct_coefficient_count=model_settings['fingerprint_width'])
+  elif preprocess == 'micro':
+    if not frontend_op:
+      raise Exception(
+          'Micro frontend op is currently not available when running TensorFlow'
+          ' directly from Python, you need to build and run through Bazel, for'
+          ' example'
+          ' `bazel run tensorflow/examples/speech_commands:freeze_graph`')
+    sample_rate = model_settings['sample_rate']
+    window_size_ms = (model_settings['window_size_samples'] *
+                      1000) / sample_rate
+    window_step_ms = (model_settings['window_stride_samples'] *
+                      1000) / sample_rate
+    int16_input = tf.cast(
+        tf.multiply(decoded_sample_data.audio, 32767), tf.int16)
+    micro_frontend = frontend_op.audio_microfrontend(
+        int16_input,
+        sample_rate=sample_rate,
+        window_size=window_size_ms,
+        window_step=window_step_ms,
+        num_channels=model_settings['fingerprint_width'],
+        out_scale=1,
+        out_type=tf.float32)
+    fingerprint_input = tf.multiply(micro_frontend, (10.0 / 256.0))
+  else:
+    raise Exception('Unknown preprocess mode "%s" (should be "mfcc",'
+                    ' "average", or "micro")' % (preprocess))
+
+  fingerprint_size = model_settings['fingerprint_size']
+  reshaped_input = tf.reshape(fingerprint_input, [-1, fingerprint_size])
+
+  logits = models.create_model(
+      reshaped_input, model_settings, model_architecture, is_training=False,
+      runtime_settings=runtime_settings)
+
+  # Create an output to use for inference.
+  softmax = tf.nn.softmax(logits, name='labels_softmax')
+
+  return reshaped_input, softmax
+
+
+def save_graph_def(file_name, frozen_graph_def):
+  """Writes a graph def file out to disk.
+
+  Args:
+    file_name: Where to save the file.
+    frozen_graph_def: GraphDef proto object to save.
+  """
+  tf.io.write_graph(
+      frozen_graph_def,
+      os.path.dirname(file_name),
+      os.path.basename(file_name),
+      as_text=False)
+  tf.compat.v1.logging.info('Saved frozen graph to %s', file_name)
+
+
+def save_saved_model(file_name, sess, input_tensor, output_tensor):
+  """Writes a SavedModel out to disk.
+
+  Args:
+    file_name: Where to save the file.
+    sess: TensorFlow session containing the graph.
+    input_tensor: Tensor object defining the input's properties.
+    output_tensor: Tensor object defining the output's properties.
+  """
+  # Store the frozen graph as a SavedModel for v2 compatibility.
+  builder = tf.compat.v1.saved_model.builder.SavedModelBuilder(file_name)
+  tensor_info_inputs = {
+      'input': tf.compat.v1.saved_model.utils.build_tensor_info(input_tensor)
+  }
+  tensor_info_outputs = {
+      'output': tf.compat.v1.saved_model.utils.build_tensor_info(output_tensor)
+  }
+  signature = (
+      tf.compat.v1.saved_model.signature_def_utils.build_signature_def(
+          inputs=tensor_info_inputs,
+          outputs=tensor_info_outputs,
+          method_name=tf.compat.v1.saved_model.signature_constants
+          .PREDICT_METHOD_NAME))
+  builder.add_meta_graph_and_variables(
+      sess,
+      [tf.compat.v1.saved_model.tag_constants.SERVING],
+      signature_def_map={
+          tf.compat.v1.saved_model.signature_constants
+          .DEFAULT_SERVING_SIGNATURE_DEF_KEY:
+              signature,
+      },
+  )
+  builder.save()
+
+
+def main(_):
+  if FLAGS.quantize:
+    try:
+      _ = tf.contrib
+    except AttributeError as e:
+      msg = e.args[0]
+      msg += ('\n\n The --quantize option still requires contrib, which is not '
+              'part of TensorFlow 2.0. Please install a previous version:'
+              '\n    `pip install tensorflow<=1.15`')
+      e.args = (msg,)
+      raise e
+
+  # Create the model and load its weights.
+  sess = tf.compat.v1.InteractiveSession()
+  input_tensor, output_tensor = create_inference_graph(
+      FLAGS.wanted_words, FLAGS.sample_rate, FLAGS.clip_duration_ms,
+      FLAGS.clip_stride_ms, FLAGS.window_size_ms, FLAGS.window_stride_ms,
+      FLAGS.feature_bin_count, FLAGS.model_architecture, FLAGS.preprocess)
+  if FLAGS.quantize:
+    tf.contrib.quantize.create_eval_graph()
+  models.load_variables_from_checkpoint(sess, FLAGS.start_checkpoint)
+
+  # Turn all the variables into inline constants inside the graph and save it.
+  frozen_graph_def = tf.compat.v1.graph_util.convert_variables_to_constants(
+      sess, sess.graph_def, ['labels_softmax'])
+
+  if FLAGS.save_format == 'graph_def':
+    save_graph_def(FLAGS.output_file, frozen_graph_def)
+  elif FLAGS.save_format == 'saved_model':
+    save_saved_model(FLAGS.output_file, sess, input_tensor, output_tensor)
+  else:
+    raise Exception('Unknown save format "%s" (should be "graph_def" or'
+                    ' "saved_model")' % (FLAGS.save_format))
+
+
+if __name__ == '__main__':
+  parser = argparse.ArgumentParser()
+  parser.add_argument(
+      '--sample_rate',
+      type=int,
+      default=16000,
+      help='Expected sample rate of the wavs',)
+  parser.add_argument(
+      '--clip_duration_ms',
+      type=int,
+      default=1000,
+      help='Expected duration in milliseconds of the wavs',)
+  parser.add_argument(
+      '--clip_stride_ms',
+      type=int,
+      default=30,
+      help='How often to run recognition. Useful for models with cache.',)
+  parser.add_argument(
+      '--window_size_ms',
+      type=float,
+      default=30.0,
+      help='How long each spectrogram timeslice is',)
+  parser.add_argument(
+      '--window_stride_ms',
+      type=float,
+      default=10.0,
+      help='How long the stride is between spectrogram timeslices',)
+  parser.add_argument(
+      '--feature_bin_count',
+      type=int,
+      default=40,
+      help='How many bins to use for the MFCC fingerprint',
+  )
+  parser.add_argument(
+      '--start_checkpoint',
+      type=str,
+      default='',
+      help='If specified, restore this pretrained model before any training.')
+  parser.add_argument(
+      '--model_architecture',
+      type=str,
+      default='conv',
+      help='What model architecture to use')
+  parser.add_argument(
+      '--wanted_words',
+      type=str,
+      default='yes,no,up,down,left,right,on,off,stop,go',
+      help='Words to use (others will be added to an unknown label)',)
+  parser.add_argument(
+      '--output_file', type=str, help='Where to save the frozen graph.')
+  parser.add_argument(
+      '--quantize',
+      type=bool,
+      default=False,
+      help='Whether to train the model for eight-bit deployment')
+  parser.add_argument(
+      '--preprocess',
+      type=str,
+      default='mfcc',
+      help='Spectrogram processing mode. Can be "mfcc" or "average"')
+  parser.add_argument(
+      '--save_format',
+      type=str,
+      default='graph_def',
+      help='How to save the result. Can be "graph_def" or "saved_model"')
+  FLAGS, unparsed = parser.parse_known_args()
+  tf.compat.v1.app.run(main=main, argv=[sys.argv[0]] + unparsed)

tensorflow/examples/speech_commands/freeze_test.py

@@ -0,0 +1,129 @@
+# 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.
+# ==============================================================================
+"""Tests for data input for speech commands."""
+
+import os.path
+
+from tensorflow.examples.speech_commands import freeze
+from tensorflow.python.framework import convert_to_constants
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops.variables import global_variables_initializer
+from tensorflow.python.platform import test
+
+
+class FreezeTest(test.TestCase):
+
+  @test_util.run_deprecated_v1
+  def testCreateInferenceGraphWithMfcc(self):
+    with self.cached_session() as sess:
+      freeze.create_inference_graph(
+          wanted_words='a,b,c,d',
+          sample_rate=16000,
+          clip_duration_ms=1000.0,
+          clip_stride_ms=30.0,
+          window_size_ms=30.0,
+          window_stride_ms=10.0,
+          feature_bin_count=40,
+          model_architecture='conv',
+          preprocess='mfcc')
+      self.assertIsNotNone(sess.graph.get_tensor_by_name('wav_data:0'))
+      self.assertIsNotNone(
+          sess.graph.get_tensor_by_name('decoded_sample_data:0'))
+      self.assertIsNotNone(sess.graph.get_tensor_by_name('labels_softmax:0'))
+      ops = [node.op for node in sess.graph_def.node]
+      self.assertEqual(1, ops.count('Mfcc'))
+
+  @test_util.run_deprecated_v1
+  def testCreateInferenceGraphWithoutMfcc(self):
+    with self.cached_session() as sess:
+      freeze.create_inference_graph(
+          wanted_words='a,b,c,d',
+          sample_rate=16000,
+          clip_duration_ms=1000.0,
+          clip_stride_ms=30.0,
+          window_size_ms=30.0,
+          window_stride_ms=10.0,
+          feature_bin_count=40,
+          model_architecture='conv',
+          preprocess='average')
+      self.assertIsNotNone(sess.graph.get_tensor_by_name('wav_data:0'))
+      self.assertIsNotNone(
+          sess.graph.get_tensor_by_name('decoded_sample_data:0'))
+      self.assertIsNotNone(sess.graph.get_tensor_by_name('labels_softmax:0'))
+      ops = [node.op for node in sess.graph_def.node]
+      self.assertEqual(0, ops.count('Mfcc'))
+
+  @test_util.run_deprecated_v1
+  def testCreateInferenceGraphWithMicro(self):
+    with self.cached_session() as sess:
+      freeze.create_inference_graph(
+          wanted_words='a,b,c,d',
+          sample_rate=16000,
+          clip_duration_ms=1000.0,
+          clip_stride_ms=30.0,
+          window_size_ms=30.0,
+          window_stride_ms=10.0,
+          feature_bin_count=40,
+          model_architecture='conv',
+          preprocess='micro')
+      self.assertIsNotNone(sess.graph.get_tensor_by_name('wav_data:0'))
+      self.assertIsNotNone(
+          sess.graph.get_tensor_by_name('decoded_sample_data:0'))
+      self.assertIsNotNone(sess.graph.get_tensor_by_name('labels_softmax:0'))
+
+  @test_util.run_deprecated_v1
+  def testFeatureBinCount(self):
+    with self.cached_session() as sess:
+      freeze.create_inference_graph(
+          wanted_words='a,b,c,d',
+          sample_rate=16000,
+          clip_duration_ms=1000.0,
+          clip_stride_ms=30.0,
+          window_size_ms=30.0,
+          window_stride_ms=10.0,
+          feature_bin_count=80,
+          model_architecture='conv',
+          preprocess='average')
+      self.assertIsNotNone(sess.graph.get_tensor_by_name('wav_data:0'))
+      self.assertIsNotNone(
+          sess.graph.get_tensor_by_name('decoded_sample_data:0'))
+      self.assertIsNotNone(sess.graph.get_tensor_by_name('labels_softmax:0'))
+      ops = [node.op for node in sess.graph_def.node]
+      self.assertEqual(0, ops.count('Mfcc'))
+
+  @test_util.run_deprecated_v1
+  def testCreateSavedModel(self):
+    tmp_dir = self.get_temp_dir()
+    saved_model_path = os.path.join(tmp_dir, 'saved_model')
+    with self.cached_session() as sess:
+      input_tensor, output_tensor = freeze.create_inference_graph(
+          wanted_words='a,b,c,d',
+          sample_rate=16000,
+          clip_duration_ms=1000.0,
+          clip_stride_ms=30.0,
+          window_size_ms=30.0,
+          window_stride_ms=10.0,
+          feature_bin_count=40,
+          model_architecture='conv',
+          preprocess='micro')
+      global_variables_initializer().run()
+      convert_to_constants.convert_variables_to_constants(
+          sess, sess.graph_def, ['labels_softmax'])
+      freeze.save_saved_model(saved_model_path, sess, input_tensor,
+                              output_tensor)
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/examples/speech_commands/generate_streaming_test_wav.py

@@ -0,0 +1,281 @@
+# 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.
+# ==============================================================================
+r"""Saves out a .wav file with synthesized conversational data and labels.
+
+The best way to estimate the real-world performance of an audio recognition
+model is by running it against a continuous stream of data, the way that it
+would be used in an application. Training evaluations are only run against
+discrete individual samples, so the results aren't as realistic.
+
+To make it easy to run evaluations against audio streams, this script uses
+samples from the testing partition of the data set, mixes them in at random
+positions together with background noise, and saves out the result as one long
+audio file.
+
+Here's an example of generating a test file:
+
+bazel run tensorflow/examples/speech_commands:generate_streaming_test_wav -- \
+--data_dir=/tmp/my_wavs --background_dir=/tmp/my_backgrounds \
+--background_volume=0.1 --test_duration_seconds=600 \
+--output_audio_file=/tmp/streaming_test.wav \
+--output_labels_file=/tmp/streaming_test_labels.txt
+
+Once you've created a streaming audio file, you can then use the
+test_streaming_accuracy tool to calculate accuracy metrics for a model.
+"""
+import argparse
+import math
+import sys
+
+import numpy as np
+import tensorflow as tf
+
+import input_data
+import models
+
+FLAGS = None
+
+
+def mix_in_audio_sample(track_data, track_offset, sample_data, sample_offset,
+                        clip_duration, sample_volume, ramp_in, ramp_out):
+  """Mixes the sample data into the main track at the specified offset.
+
+  Args:
+    track_data: Numpy array holding main audio data. Modified in-place.
+    track_offset: Where to mix the sample into the main track.
+    sample_data: Numpy array of audio data to mix into the main track.
+    sample_offset: Where to start in the audio sample.
+    clip_duration: How long the sample segment is.
+    sample_volume: Loudness to mix the sample in at.
+    ramp_in: Length in samples of volume increase stage.
+    ramp_out: Length in samples of volume decrease stage.
+  """
+  ramp_out_index = clip_duration - ramp_out
+  track_end = min(track_offset + clip_duration, track_data.shape[0])
+  track_end = min(track_end,
+                  track_offset + (sample_data.shape[0] - sample_offset))
+  sample_range = track_end - track_offset
+  for i in range(sample_range):
+    if i < ramp_in:
+      envelope_scale = i / ramp_in
+    elif i > ramp_out_index:
+      envelope_scale = (clip_duration - i) / ramp_out
+    else:
+      envelope_scale = 1
+    sample_input = sample_data[sample_offset + i]
+    track_data[track_offset
+               + i] += sample_input * envelope_scale * sample_volume
+
+
+def main(_):
+  words_list = input_data.prepare_words_list(FLAGS.wanted_words.split(','))
+  model_settings = models.prepare_model_settings(
+      len(words_list), FLAGS.sample_rate, FLAGS.clip_duration_ms,
+      FLAGS.window_size_ms, FLAGS.window_stride_ms, FLAGS.feature_bin_count,
+      'mfcc')
+  audio_processor = input_data.AudioProcessor(
+      '', FLAGS.data_dir, FLAGS.silence_percentage, 10,
+      FLAGS.wanted_words.split(','), FLAGS.validation_percentage,
+      FLAGS.testing_percentage, model_settings, FLAGS.data_dir)
+
+  output_audio_sample_count = FLAGS.sample_rate * FLAGS.test_duration_seconds
+  output_audio = np.zeros((output_audio_sample_count,), dtype=np.float32)
+
+  # Set up background audio.
+  background_crossover_ms = 500
+  background_segment_duration_ms = (
+      FLAGS.clip_duration_ms + background_crossover_ms)
+  background_segment_duration_samples = int(
+      (background_segment_duration_ms * FLAGS.sample_rate) / 1000)
+  background_segment_stride_samples = int(
+      (FLAGS.clip_duration_ms * FLAGS.sample_rate) / 1000)
+  background_ramp_samples = int(
+      ((background_crossover_ms / 2) * FLAGS.sample_rate) / 1000)
+
+  # Mix the background audio into the main track.
+  how_many_backgrounds = int(
+      math.ceil(output_audio_sample_count / background_segment_stride_samples))
+  for i in range(how_many_backgrounds):
+    output_offset = int(i * background_segment_stride_samples)
+    background_index = np.random.randint(len(audio_processor.background_data))
+    background_samples = audio_processor.background_data[background_index]
+    background_offset = np.random.randint(
+        0, len(background_samples) - model_settings['desired_samples'])
+    background_volume = np.random.uniform(0, FLAGS.background_volume)
+    mix_in_audio_sample(output_audio, output_offset, background_samples,
+                        background_offset, background_segment_duration_samples,
+                        background_volume, background_ramp_samples,
+                        background_ramp_samples)
+
+  # Mix the words into the main track, noting their labels and positions.
+  output_labels = []
+  word_stride_ms = FLAGS.clip_duration_ms + FLAGS.word_gap_ms
+  word_stride_samples = int((word_stride_ms * FLAGS.sample_rate) / 1000)
+  clip_duration_samples = int(
+      (FLAGS.clip_duration_ms * FLAGS.sample_rate) / 1000)
+  word_gap_samples = int((FLAGS.word_gap_ms * FLAGS.sample_rate) / 1000)
+  how_many_words = int(
+      math.floor(output_audio_sample_count / word_stride_samples))
+  all_test_data, all_test_labels = audio_processor.get_unprocessed_data(
+      -1, model_settings, 'testing')
+  for i in range(how_many_words):
+    output_offset = (
+        int(i * word_stride_samples) + np.random.randint(word_gap_samples))
+    output_offset_ms = (output_offset * 1000) / FLAGS.sample_rate
+    is_unknown = np.random.randint(100) < FLAGS.unknown_percentage
+    if is_unknown:
+      wanted_label = input_data.UNKNOWN_WORD_LABEL
+    else:
+      wanted_label = words_list[2 + np.random.randint(len(words_list) - 2)]
+    test_data_start = np.random.randint(len(all_test_data))
+    found_sample_data = None
+    index_lookup = np.arange(len(all_test_data), dtype=np.int32)
+    np.random.shuffle(index_lookup)
+    for test_data_offset in range(len(all_test_data)):
+      test_data_index = index_lookup[(
+          test_data_start + test_data_offset) % len(all_test_data)]
+      current_label = all_test_labels[test_data_index]
+      if current_label == wanted_label:
+        found_sample_data = all_test_data[test_data_index]
+        break
+    mix_in_audio_sample(output_audio, output_offset, found_sample_data, 0,
+                        clip_duration_samples, 1.0, 500, 500)
+    output_labels.append({'label': wanted_label, 'time': output_offset_ms})
+
+  input_data.save_wav_file(FLAGS.output_audio_file, output_audio,
+                           FLAGS.sample_rate)
+  tf.compat.v1.logging.info('Saved streaming test wav to %s',
+                            FLAGS.output_audio_file)
+
+  with open(FLAGS.output_labels_file, 'w') as f:
+    for output_label in output_labels:
+      f.write('%s, %f\n' % (output_label['label'], output_label['time']))
+  tf.compat.v1.logging.info('Saved streaming test labels to %s',
+                            FLAGS.output_labels_file)
+
+
+if __name__ == '__main__':
+  parser = argparse.ArgumentParser()
+  parser.add_argument(
+      '--data_url',
+      type=str,
+      # pylint: disable=line-too-long
+      default='https://storage.googleapis.com/download.tensorflow.org/data/speech_commands_v0.01.tar.gz',
+      # pylint: enable=line-too-long
+      help='Location of speech training data')
+  parser.add_argument(
+      '--data_dir',
+      type=str,
+      default='/tmp/speech_dataset',
+      help="""\
+      Where to download the speech training data to.
+      """)
+  parser.add_argument(
+      '--background_dir',
+      type=str,
+      default='',
+      help="""\
+      Path to a directory of .wav files to mix in as background noise during training.
+      """)
+  parser.add_argument(
+      '--background_volume',
+      type=float,
+      default=0.1,
+      help="""\
+      How loud the background noise should be, between 0 and 1.
+      """)
+  parser.add_argument(
+      '--background_frequency',
+      type=float,
+      default=0.8,
+      help="""\
+      How many of the training samples have background noise mixed in.
+      """)
+  parser.add_argument(
+      '--silence_percentage',
+      type=float,
+      default=10.0,
+      help="""\
+      How much of the training data should be silence.
+      """)
+  parser.add_argument(
+      '--testing_percentage',
+      type=int,
+      default=10,
+      help='What percentage of wavs to use as a test set.')
+  parser.add_argument(
+      '--validation_percentage',
+      type=int,
+      default=10,
+      help='What percentage of wavs to use as a validation set.')
+  parser.add_argument(
+      '--sample_rate',
+      type=int,
+      default=16000,
+      help='Expected sample rate of the wavs.',)
+  parser.add_argument(
+      '--clip_duration_ms',
+      type=int,
+      default=1000,
+      help='Expected duration in milliseconds of the wavs.',)
+  parser.add_argument(
+      '--window_size_ms',
+      type=float,
+      default=30.0,
+      help='How long each spectrogram timeslice is',)
+  parser.add_argument(
+      '--window_stride_ms',
+      type=float,
+      default=10.0,
+      help='How long the stride is between spectrogram timeslices',)
+  parser.add_argument(
+      '--feature_bin_count',
+      type=int,
+      default=40,
+      help='How many bins to use for the MFCC fingerprint',
+  )
+  parser.add_argument(
+      '--wanted_words',
+      type=str,
+      default='yes,no,up,down,left,right,on,off,stop,go',
+      help='Words to use (others will be added to an unknown label)',)
+  parser.add_argument(
+      '--output_audio_file',
+      type=str,
+      default='/tmp/speech_commands_train/streaming_test.wav',
+      help='File to save the generated test audio to.')
+  parser.add_argument(
+      '--output_labels_file',
+      type=str,
+      default='/tmp/speech_commands_train/streaming_test_labels.txt',
+      help='File to save the generated test labels to.')
+  parser.add_argument(
+      '--test_duration_seconds',
+      type=int,
+      default=600,
+      help='How long the generated test audio file should be.',)
+  parser.add_argument(
+      '--word_gap_ms',
+      type=int,
+      default=2000,
+      help='How long the average gap should be between words.',)
+  parser.add_argument(
+      '--unknown_percentage',
+      type=int,
+      default=30,
+      help='What percentage of words should be unknown.')
+
+  FLAGS, unparsed = parser.parse_known_args()
+  tf.compat.v1.app.run(main=main, argv=[sys.argv[0]] + unparsed)

tensorflow/examples/speech_commands/generate_streaming_test_wav_test.py

@@ -0,0 +1,35 @@
+# 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.
+# ==============================================================================
+"""Tests for test file generation for speech commands."""
+
+import numpy as np
+
+from tensorflow.examples.speech_commands import generate_streaming_test_wav
+from tensorflow.python.platform import test
+
+
+class GenerateStreamingTestWavTest(test.TestCase):
+
+  def testMixInAudioSample(self):
+    track_data = np.zeros([10000])
+    sample_data = np.ones([1000])
+    generate_streaming_test_wav.mix_in_audio_sample(
+        track_data, 2000, sample_data, 0, 1000, 1.0, 100, 100)
+    self.assertNear(1.0, track_data[2500], 0.0001)
+    self.assertNear(0.0, track_data[3500], 0.0001)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/examples/speech_commands/input_data.py

@@ -0,0 +1,679 @@
+# 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.
+# ==============================================================================
+"""Model definitions for simple speech recognition.
+
+"""
+import hashlib
+import math
+import os.path
+import random
+import re
+import sys
+import tarfile
+
+import numpy as np
+import urllib
+import tensorflow as tf
+
+from tensorflow.python.ops import gen_audio_ops as audio_ops
+from tensorflow.python.ops import io_ops
+from tensorflow.python.platform import gfile
+from tensorflow.python.util import compat
+
+tf.compat.v1.disable_eager_execution()
+
+# If it's available, load the specialized feature generator. If this doesn't
+# work, try building with bazel instead of running the Python script directly.
+try:
+  from tensorflow.lite.experimental.microfrontend.python.ops import audio_microfrontend_op as frontend_op  # pylint:disable=g-import-not-at-top
+except ImportError:
+  frontend_op = None
+
+MAX_NUM_WAVS_PER_CLASS = 2**27 - 1  # ~134M
+SILENCE_LABEL = '_silence_'
+SILENCE_INDEX = 0
+UNKNOWN_WORD_LABEL = '_unknown_'
+UNKNOWN_WORD_INDEX = 1
+BACKGROUND_NOISE_DIR_NAME = '_background_noise_'
+RANDOM_SEED = 59185
+
+
+def prepare_words_list(wanted_words):
+  """Prepends common tokens to the custom word list.
+
+  Args:
+    wanted_words: List of strings containing the custom words.
+
+  Returns:
+    List with the standard silence and unknown tokens added.
+  """
+  return [SILENCE_LABEL, UNKNOWN_WORD_LABEL] + wanted_words
+
+
+def which_set(filename, validation_percentage, testing_percentage):
+  """Determines which data partition the file should belong to.
+
+  We want to keep files in the same training, validation, or testing sets even
+  if new ones are added over time. This makes it less likely that testing
+  samples will accidentally be reused in training when long runs are restarted
+  for example. To keep this stability, a hash of the filename is taken and used
+  to determine which set it should belong to. This determination only depends on
+  the name and the set proportions, so it won't change as other files are added.
+
+  It's also useful to associate particular files as related (for example words
+  spoken by the same person), so anything after '_nohash_' in a filename is
+  ignored for set determination. This ensures that 'bobby_nohash_0.wav' and
+  'bobby_nohash_1.wav' are always in the same set, for example.
+
+  Args:
+    filename: File path of the data sample.
+    validation_percentage: How much of the data set to use for validation.
+    testing_percentage: How much of the data set to use for testing.
+
+  Returns:
+    String, one of 'training', 'validation', or 'testing'.
+  """
+  base_name = os.path.basename(filename)
+  # We want to ignore anything after '_nohash_' in the file name when
+  # deciding which set to put a wav in, so the data set creator has a way of
+  # grouping wavs that are close variations of each other.
+  hash_name = re.sub(r'_nohash_.*$', '', base_name)
+  # This looks a bit magical, but we need to decide whether this file should
+  # go into the training, testing, or validation sets, and we want to keep
+  # existing files in the same set even if more files are subsequently
+  # added.
+  # To do that, we need a stable way of deciding based on just the file name
+  # itself, so we do a hash of that and then use that to generate a
+  # probability value that we use to assign it.
+  hash_name_hashed = hashlib.sha1(compat.as_bytes(hash_name)).hexdigest()
+  percentage_hash = ((int(hash_name_hashed, 16) %
+                      (MAX_NUM_WAVS_PER_CLASS + 1)) *
+                     (100.0 / MAX_NUM_WAVS_PER_CLASS))
+  if percentage_hash < validation_percentage:
+    result = 'validation'
+  elif percentage_hash < (testing_percentage + validation_percentage):
+    result = 'testing'
+  else:
+    result = 'training'
+  return result
+
+
+def load_wav_file(filename):
+  """Loads an audio file and returns a float PCM-encoded array of samples.
+
+  Args:
+    filename: Path to the .wav file to load.
+
+  Returns:
+    Numpy array holding the sample data as floats between -1.0 and 1.0.
+  """
+  with tf.compat.v1.Session(graph=tf.Graph()) as sess:
+    wav_filename_placeholder = tf.compat.v1.placeholder(tf.string, [])
+    wav_loader = io_ops.read_file(wav_filename_placeholder)
+    wav_decoder = tf.audio.decode_wav(wav_loader, desired_channels=1)
+    return sess.run(
+        wav_decoder,
+        feed_dict={wav_filename_placeholder: filename}).audio.flatten()
+
+
+def save_wav_file(filename, wav_data, sample_rate):
+  """Saves audio sample data to a .wav audio file.
+
+  Args:
+    filename: Path to save the file to.
+    wav_data: 2D array of float PCM-encoded audio data.
+    sample_rate: Samples per second to encode in the file.
+  """
+  with tf.compat.v1.Session(graph=tf.Graph()) as sess:
+    wav_filename_placeholder = tf.compat.v1.placeholder(tf.string, [])
+    sample_rate_placeholder = tf.compat.v1.placeholder(tf.int32, [])
+    wav_data_placeholder = tf.compat.v1.placeholder(tf.float32, [None, 1])
+    wav_encoder = tf.audio.encode_wav(wav_data_placeholder,
+                                      sample_rate_placeholder)
+    wav_saver = io_ops.write_file(wav_filename_placeholder, wav_encoder)
+    sess.run(
+        wav_saver,
+        feed_dict={
+            wav_filename_placeholder: filename,
+            sample_rate_placeholder: sample_rate,
+            wav_data_placeholder: np.reshape(wav_data, (-1, 1))
+        })
+
+
+def get_features_range(model_settings):
+  """Returns the expected min/max for generated features.
+
+  Args:
+    model_settings: Information about the current model being trained.
+
+  Returns:
+    Min/max float pair holding the range of features.
+
+  Raises:
+    Exception: If preprocessing mode isn't recognized.
+  """
+  # TODO(petewarden): These values have been derived from the observed ranges
+  # of spectrogram and MFCC inputs. If the preprocessing pipeline changes,
+  # they may need to be updated.
+  if model_settings['preprocess'] == 'average':
+    features_min = 0.0
+    features_max = 127.5
+  elif model_settings['preprocess'] == 'mfcc':
+    features_min = -247.0
+    features_max = 30.0
+  elif model_settings['preprocess'] == 'micro':
+    features_min = 0.0
+    features_max = 26.0
+  else:
+    raise Exception('Unknown preprocess mode "%s" (should be "mfcc",'
+                    ' "average", or "micro")' % (model_settings['preprocess']))
+  return features_min, features_max
+
+
+class AudioProcessor(object):
+  """Handles loading, partitioning, and preparing audio training data."""
+
+  def __init__(self, data_url, data_dir, silence_percentage, unknown_percentage,
+               wanted_words, validation_percentage, testing_percentage,
+               model_settings, summaries_dir):
+    if data_dir:
+      self.data_dir = data_dir
+      self.maybe_download_and_extract_dataset(data_url, data_dir)
+      self.prepare_data_index(silence_percentage, unknown_percentage,
+                              wanted_words, validation_percentage,
+                              testing_percentage)
+      self.prepare_background_data()
+    self.prepare_processing_graph(model_settings, summaries_dir)
+
+  def maybe_download_and_extract_dataset(self, data_url, dest_directory):
+    """Download and extract data set tar file.
+
+    If the data set we're using doesn't already exist, this function
+    downloads it from the TensorFlow.org website and unpacks it into a
+    directory.
+    If the data_url is none, don't download anything and expect the data
+    directory to contain the correct files already.
+
+    Args:
+      data_url: Web location of the tar file containing the data set.
+      dest_directory: File path to extract data to.
+    """
+    if not data_url:
+      return
+    if not gfile.Exists(dest_directory):
+      os.makedirs(dest_directory)
+    filename = data_url.split('/')[-1]
+    filepath = os.path.join(dest_directory, filename)
+    if not gfile.Exists(filepath):
+
+      def _progress(count, block_size, total_size):
+        sys.stdout.write(
+            '\r>> Downloading %s %.1f%%' %
+            (filename, float(count * block_size) / float(total_size) * 100.0))
+        sys.stdout.flush()
+
+      try:
+        filepath, _ = urllib.request.urlretrieve(data_url, filepath, _progress)
+      except Exception:
+        tf.compat.v1.logging.error(
+            'Failed to download URL: {0} to folder: {1}. Please make sure you '
+            'have enough free space and an internet connection'.format(
+                data_url, filepath))
+        raise
+      print()
+      statinfo = os.stat(filepath)
+      tf.compat.v1.logging.info(
+          'Successfully downloaded {0} ({1} bytes)'.format(
+              filename, statinfo.st_size))
+      tarfile.open(filepath, 'r:gz').extractall(dest_directory)
+
+  def prepare_data_index(self, silence_percentage, unknown_percentage,
+                         wanted_words, validation_percentage,
+                         testing_percentage):
+    """Prepares a list of the samples organized by set and label.
+
+    The training loop needs a list of all the available data, organized by
+    which partition it should belong to, and with ground truth labels attached.
+    This function analyzes the folders below the `data_dir`, figures out the
+    right
+    labels for each file based on the name of the subdirectory it belongs to,
+    and uses a stable hash to assign it to a data set partition.
+
+    Args:
+      silence_percentage: How much of the resulting data should be background.
+      unknown_percentage: How much should be audio outside the wanted classes.
+      wanted_words: Labels of the classes we want to be able to recognize.
+      validation_percentage: How much of the data set to use for validation.
+      testing_percentage: How much of the data set to use for testing.
+
+    Returns:
+      Dictionary containing a list of file information for each set partition,
+      and a lookup map for each class to determine its numeric index.
+
+    Raises:
+      Exception: If expected files are not found.
+    """
+    # Make sure the shuffling and picking of unknowns is deterministic.
+    random.seed(RANDOM_SEED)
+    wanted_words_index = {}
+    for index, wanted_word in enumerate(wanted_words):
+      wanted_words_index[wanted_word] = index + 2
+    self.data_index = {'validation': [], 'testing': [], 'training': []}
+    unknown_index = {'validation': [], 'testing': [], 'training': []}
+    all_words = {}
+    # Look through all the subfolders to find audio samples
+    search_path = os.path.join(self.data_dir, '*', '*.wav')
+    for wav_path in gfile.Glob(search_path):
+      _, word = os.path.split(os.path.dirname(wav_path))
+      word = word.lower()
+      # Treat the '_background_noise_' folder as a special case, since we expect
+      # it to contain long audio samples we mix in to improve training.
+      if word == BACKGROUND_NOISE_DIR_NAME:
+        continue
+      all_words[word] = True
+      set_index = which_set(wav_path, validation_percentage, testing_percentage)
+      # If it's a known class, store its detail, otherwise add it to the list
+      # we'll use to train the unknown label.
+      if word in wanted_words_index:
+        self.data_index[set_index].append({'label': word, 'file': wav_path})
+      else:
+        unknown_index[set_index].append({'label': word, 'file': wav_path})
+    if not all_words:
+      raise Exception('No .wavs found at ' + search_path)
+    for index, wanted_word in enumerate(wanted_words):
+      if wanted_word not in all_words:
+        raise Exception('Expected to find ' + wanted_word +
+                        ' in labels but only found ' +
+                        ', '.join(all_words.keys()))
+    # We need an arbitrary file to load as the input for the silence samples.
+    # It's multiplied by zero later, so the content doesn't matter.
+    silence_wav_path = self.data_index['training'][0]['file']
+    for set_index in ['validation', 'testing', 'training']:
+      set_size = len(self.data_index[set_index])
+      silence_size = int(math.ceil(set_size * silence_percentage / 100))
+      for _ in range(silence_size):
+        self.data_index[set_index].append({
+            'label': SILENCE_LABEL,
+            'file': silence_wav_path
+        })
+      # Pick some unknowns to add to each partition of the data set.
+      random.shuffle(unknown_index[set_index])
+      unknown_size = int(math.ceil(set_size * unknown_percentage / 100))
+      self.data_index[set_index].extend(unknown_index[set_index][:unknown_size])
+    # Make sure the ordering is random.
+    for set_index in ['validation', 'testing', 'training']:
+      random.shuffle(self.data_index[set_index])
+    # Prepare the rest of the result data structure.
+    self.words_list = prepare_words_list(wanted_words)
+    self.word_to_index = {}
+    for word in all_words:
+      if word in wanted_words_index:
+        self.word_to_index[word] = wanted_words_index[word]
+      else:
+        self.word_to_index[word] = UNKNOWN_WORD_INDEX
+    self.word_to_index[SILENCE_LABEL] = SILENCE_INDEX
+
+  def prepare_background_data(self):
+    """Searches a folder for background noise audio, and loads it into memory.
+
+    It's expected that the background audio samples will be in a subdirectory
+    named '_background_noise_' inside the 'data_dir' folder, as .wavs that match
+    the sample rate of the training data, but can be much longer in duration.
+
+    If the '_background_noise_' folder doesn't exist at all, this isn't an
+    error, it's just taken to mean that no background noise augmentation should
+    be used. If the folder does exist, but it's empty, that's treated as an
+    error.
+
+    Returns:
+      List of raw PCM-encoded audio samples of background noise.
+
+    Raises:
+      Exception: If files aren't found in the folder.
+    """
+    self.background_data = []
+    background_dir = os.path.join(self.data_dir, BACKGROUND_NOISE_DIR_NAME)
+    if not gfile.Exists(background_dir):
+      return self.background_data
+    with tf.compat.v1.Session(graph=tf.Graph()) as sess:
+      wav_filename_placeholder = tf.compat.v1.placeholder(tf.string, [])
+      wav_loader = io_ops.read_file(wav_filename_placeholder)
+      wav_decoder = tf.audio.decode_wav(wav_loader, desired_channels=1)
+      search_path = os.path.join(self.data_dir, BACKGROUND_NOISE_DIR_NAME,
+                                 '*.wav')
+      for wav_path in gfile.Glob(search_path):
+        wav_data = sess.run(
+            wav_decoder,
+            feed_dict={wav_filename_placeholder: wav_path}).audio.flatten()
+        self.background_data.append(wav_data)
+      if not self.background_data:
+        raise Exception('No background wav files were found in ' + search_path)
+
+  def prepare_processing_graph(self, model_settings, summaries_dir):
+    """Builds a TensorFlow graph to apply the input distortions.
+
+    Creates a graph that loads a WAVE file, decodes it, scales the volume,
+    shifts it in time, adds in background noise, calculates a spectrogram, and
+    then builds an MFCC fingerprint from that.
+
+    This must be called with an active TensorFlow session running, and it
+    creates multiple placeholder inputs, and one output:
+
+      - wav_filename_placeholder_: Filename of the WAV to load.
+      - foreground_volume_placeholder_: How loud the main clip should be.
+      - time_shift_padding_placeholder_: Where to pad the clip.
+      - time_shift_offset_placeholder_: How much to move the clip in time.
+      - background_data_placeholder_: PCM sample data for background noise.
+      - background_volume_placeholder_: Loudness of mixed-in background.
+      - output_: Output 2D fingerprint of processed audio.
+
+    Args:
+      model_settings: Information about the current model being trained.
+      summaries_dir: Path to save training summary information to.
+
+    Raises:
+      ValueError: If the preprocessing mode isn't recognized.
+      Exception: If the preprocessor wasn't compiled in.
+    """
+    with tf.compat.v1.get_default_graph().name_scope('data'):
+      desired_samples = model_settings['desired_samples']
+      self.wav_filename_placeholder_ = tf.compat.v1.placeholder(
+          tf.string, [], name='wav_filename')
+      wav_loader = io_ops.read_file(self.wav_filename_placeholder_)
+      wav_decoder = tf.audio.decode_wav(
+          wav_loader, desired_channels=1, desired_samples=desired_samples)
+      # Allow the audio sample's volume to be adjusted.
+      self.foreground_volume_placeholder_ = tf.compat.v1.placeholder(
+          tf.float32, [], name='foreground_volume')
+      scaled_foreground = tf.multiply(wav_decoder.audio,
+                                      self.foreground_volume_placeholder_)
+      # Shift the sample's start position, and pad any gaps with zeros.
+      self.time_shift_padding_placeholder_ = tf.compat.v1.placeholder(
+          tf.int32, [2, 2], name='time_shift_padding')
+      self.time_shift_offset_placeholder_ = tf.compat.v1.placeholder(
+          tf.int32, [2], name='time_shift_offset')
+      padded_foreground = tf.pad(
+          tensor=scaled_foreground,
+          paddings=self.time_shift_padding_placeholder_,
+          mode='CONSTANT')
+      sliced_foreground = tf.slice(padded_foreground,
+                                   self.time_shift_offset_placeholder_,
+                                   [desired_samples, -1])
+      # Mix in background noise.
+      self.background_data_placeholder_ = tf.compat.v1.placeholder(
+          tf.float32, [desired_samples, 1], name='background_data')
+      self.background_volume_placeholder_ = tf.compat.v1.placeholder(
+          tf.float32, [], name='background_volume')
+      background_mul = tf.multiply(self.background_data_placeholder_,
+                                   self.background_volume_placeholder_)
+      background_add = tf.add(background_mul, sliced_foreground)
+      background_clamp = tf.clip_by_value(background_add, -1.0, 1.0)
+      # Run the spectrogram and MFCC ops to get a 2D 'fingerprint' of the audio.
+      spectrogram = audio_ops.audio_spectrogram(
+          background_clamp,
+          window_size=model_settings['window_size_samples'],
+          stride=model_settings['window_stride_samples'],
+          magnitude_squared=True)
+      tf.compat.v1.summary.image(
+          'spectrogram', tf.expand_dims(spectrogram, -1), max_outputs=1)
+      # The number of buckets in each FFT row in the spectrogram will depend on
+      # how many input samples there are in each window. This can be quite
+      # large, with a 160 sample window producing 127 buckets for example. We
+      # don't need this level of detail for classification, so we often want to
+      # shrink them down to produce a smaller result. That's what this section
+      # implements. One method is to use average pooling to merge adjacent
+      # buckets, but a more sophisticated approach is to apply the MFCC
+      # algorithm to shrink the representation.
+      if model_settings['preprocess'] == 'average':
+        self.output_ = tf.nn.pool(
+            input=tf.expand_dims(spectrogram, -1),
+            window_shape=[1, model_settings['average_window_width']],
+            strides=[1, model_settings['average_window_width']],
+            pooling_type='AVG',
+            padding='SAME')
+        tf.compat.v1.summary.image('shrunk_spectrogram',
+                                   self.output_,
+                                   max_outputs=1)
+      elif model_settings['preprocess'] == 'mfcc':
+        self.output_ = audio_ops.mfcc(
+            spectrogram,
+            wav_decoder.sample_rate,
+            dct_coefficient_count=model_settings['fingerprint_width'])
+        tf.compat.v1.summary.image(
+            'mfcc', tf.expand_dims(self.output_, -1), max_outputs=1)
+      elif model_settings['preprocess'] == 'micro':
+        if not frontend_op:
+          raise Exception(
+              'Micro frontend op is currently not available when running'
+              ' TensorFlow directly from Python, you need to build and run'
+              ' through Bazel')
+        sample_rate = model_settings['sample_rate']
+        window_size_ms = (model_settings['window_size_samples'] *
+                          1000) / sample_rate
+        window_step_ms = (model_settings['window_stride_samples'] *
+                          1000) / sample_rate
+        int16_input = tf.cast(tf.multiply(background_clamp, 32768), tf.int16)
+        micro_frontend = frontend_op.audio_microfrontend(
+            int16_input,
+            sample_rate=sample_rate,
+            window_size=window_size_ms,
+            window_step=window_step_ms,
+            num_channels=model_settings['fingerprint_width'],
+            out_scale=1,
+            out_type=tf.float32)
+        self.output_ = tf.multiply(micro_frontend, (10.0 / 256.0))
+        tf.compat.v1.summary.image(
+            'micro',
+            tf.expand_dims(tf.expand_dims(self.output_, -1), 0),
+            max_outputs=1)
+      else:
+        raise ValueError('Unknown preprocess mode "%s" (should be "mfcc", '
+                         ' "average", or "micro")' %
+                         (model_settings['preprocess']))
+
+      # Merge all the summaries and write them out to /tmp/retrain_logs (by
+      # default)
+      self.merged_summaries_ = tf.compat.v1.summary.merge_all(scope='data')
+      if summaries_dir:
+        self.summary_writer_ = tf.compat.v1.summary.FileWriter(
+            summaries_dir + '/data', tf.compat.v1.get_default_graph())
+
+  def set_size(self, mode):
+    """Calculates the number of samples in the dataset partition.
+
+    Args:
+      mode: Which partition, must be 'training', 'validation', or 'testing'.
+
+    Returns:
+      Number of samples in the partition.
+    """
+    return len(self.data_index[mode])
+
+  def get_data(self, how_many, offset, model_settings, background_frequency,
+               background_volume_range, time_shift, mode, sess):
+    """Gather samples from the data set, applying transformations as needed.
+
+    When the mode is 'training', a random selection of samples will be returned,
+    otherwise the first N clips in the partition will be used. This ensures that
+    validation always uses the same samples, reducing noise in the metrics.
+
+    Args:
+      how_many: Desired number of samples to return. -1 means the entire
+        contents of this partition.
+      offset: Where to start when fetching deterministically.
+      model_settings: Information about the current model being trained.
+      background_frequency: How many clips will have background noise, 0.0 to
+        1.0.
+      background_volume_range: How loud the background noise will be.
+      time_shift: How much to randomly shift the clips by in time.
+      mode: Which partition to use, must be 'training', 'validation', or
+        'testing'.
+      sess: TensorFlow session that was active when processor was created.
+
+    Returns:
+      List of sample data for the transformed samples, and list of label indexes
+
+    Raises:
+      ValueError: If background samples are too short.
+    """
+    # Pick one of the partitions to choose samples from.
+    candidates = self.data_index[mode]
+    if how_many == -1:
+      sample_count = len(candidates)
+    else:
+      sample_count = max(0, min(how_many, len(candidates) - offset))
+    # Data and labels will be populated and returned.
+    data = np.zeros((sample_count, model_settings['fingerprint_size']))
+    labels = np.zeros(sample_count)
+    desired_samples = model_settings['desired_samples']
+    use_background = self.background_data and (mode == 'training')
+    pick_deterministically = (mode != 'training')
+    # Use the processing graph we created earlier to repeatedly to generate the
+    # final output sample data we'll use in training.
+    for i in range(offset, offset + sample_count):
+      # Pick which audio sample to use.
+      if how_many == -1 or pick_deterministically:
+        sample_index = i
+      else:
+        sample_index = np.random.randint(len(candidates))
+      sample = candidates[sample_index]
+      # If we're time shifting, set up the offset for this sample.
+      if time_shift > 0:
+        time_shift_amount = np.random.randint(-time_shift, time_shift)
+      else:
+        time_shift_amount = 0
+      if time_shift_amount > 0:
+        time_shift_padding = [[time_shift_amount, 0], [0, 0]]
+        time_shift_offset = [0, 0]
+      else:
+        time_shift_padding = [[0, -time_shift_amount], [0, 0]]
+        time_shift_offset = [-time_shift_amount, 0]
+      input_dict = {
+          self.wav_filename_placeholder_: sample['file'],
+          self.time_shift_padding_placeholder_: time_shift_padding,
+          self.time_shift_offset_placeholder_: time_shift_offset,
+      }
+      # Choose a section of background noise to mix in.
+      if use_background or sample['label'] == SILENCE_LABEL:
+        background_index = np.random.randint(len(self.background_data))
+        background_samples = self.background_data[background_index]
+        if len(background_samples) <= model_settings['desired_samples']:
+          raise ValueError(
+              'Background sample is too short! Need more than %d'
+              ' samples but only %d were found' %
+              (model_settings['desired_samples'], len(background_samples)))
+        background_offset = np.random.randint(
+            0, len(background_samples) - model_settings['desired_samples'])
+        background_clipped = background_samples[background_offset:(
+            background_offset + desired_samples)]
+        background_reshaped = background_clipped.reshape([desired_samples, 1])
+        if sample['label'] == SILENCE_LABEL:
+          background_volume = np.random.uniform(0, 1)
+        elif np.random.uniform(0, 1) < background_frequency:
+          background_volume = np.random.uniform(0, background_volume_range)
+        else:
+          background_volume = 0
+      else:
+        background_reshaped = np.zeros([desired_samples, 1])
+        background_volume = 0
+      input_dict[self.background_data_placeholder_] = background_reshaped
+      input_dict[self.background_volume_placeholder_] = background_volume
+      # If we want silence, mute out the main sample but leave the background.
+      if sample['label'] == SILENCE_LABEL:
+        input_dict[self.foreground_volume_placeholder_] = 0
+      else:
+        input_dict[self.foreground_volume_placeholder_] = 1
+      # Run the graph to produce the output audio.
+      summary, data_tensor = sess.run(
+          [self.merged_summaries_, self.output_], feed_dict=input_dict)
+      self.summary_writer_.add_summary(summary)
+      data[i - offset, :] = data_tensor.flatten()
+      label_index = self.word_to_index[sample['label']]
+      labels[i - offset] = label_index
+    return data, labels
+
+  def get_features_for_wav(self, wav_filename, model_settings, sess):
+    """Applies the feature transformation process to the input_wav.
+
+    Runs the feature generation process (generally producing a spectrogram from
+    the input samples) on the WAV file. This can be useful for testing and
+    verifying implementations being run on other platforms.
+
+    Args:
+      wav_filename: The path to the input audio file.
+      model_settings: Information about the current model being trained.
+      sess: TensorFlow session that was active when processor was created.
+
+    Returns:
+      Numpy data array containing the generated features.
+    """
+    desired_samples = model_settings['desired_samples']
+    input_dict = {
+        self.wav_filename_placeholder_: wav_filename,
+        self.time_shift_padding_placeholder_: [[0, 0], [0, 0]],
+        self.time_shift_offset_placeholder_: [0, 0],
+        self.background_data_placeholder_: np.zeros([desired_samples, 1]),
+        self.background_volume_placeholder_: 0,
+        self.foreground_volume_placeholder_: 1,
+    }
+    # Run the graph to produce the output audio.
+    data_tensor = sess.run([self.output_], feed_dict=input_dict)
+    return data_tensor
+
+  def get_unprocessed_data(self, how_many, model_settings, mode):
+    """Retrieve sample data for the given partition, with no transformations.
+
+    Args:
+      how_many: Desired number of samples to return. -1 means the entire
+        contents of this partition.
+      model_settings: Information about the current model being trained.
+      mode: Which partition to use, must be 'training', 'validation', or
+        'testing'.
+
+    Returns:
+      List of sample data for the samples, and list of labels in one-hot form.
+    """
+    candidates = self.data_index[mode]
+    if how_many == -1:
+      sample_count = len(candidates)
+    else:
+      sample_count = how_many
+    desired_samples = model_settings['desired_samples']
+    words_list = self.words_list
+    data = np.zeros((sample_count, desired_samples))
+    labels = []
+    with tf.compat.v1.Session(graph=tf.Graph()) as sess:
+      wav_filename_placeholder = tf.compat.v1.placeholder(tf.string, [])
+      wav_loader = io_ops.read_file(wav_filename_placeholder)
+      wav_decoder = tf.audio.decode_wav(
+          wav_loader, desired_channels=1, desired_samples=desired_samples)
+      foreground_volume_placeholder = tf.compat.v1.placeholder(tf.float32, [])
+      scaled_foreground = tf.multiply(wav_decoder.audio,
+                                      foreground_volume_placeholder)
+      for i in range(sample_count):
+        if how_many == -1:
+          sample_index = i
+        else:
+          sample_index = np.random.randint(len(candidates))
+        sample = candidates[sample_index]
+        input_dict = {wav_filename_placeholder: sample['file']}
+        if sample['label'] == SILENCE_LABEL:
+          input_dict[foreground_volume_placeholder] = 0
+        else:
+          input_dict[foreground_volume_placeholder] = 1
+        data[i, :] = sess.run(scaled_foreground, feed_dict=input_dict).flatten()
+        label_index = self.word_to_index[sample['label']]
+        labels.append(words_list[label_index])
+    return data, labels

tensorflow/examples/speech_commands/input_data_test.py

@@ -0,0 +1,286 @@
+# 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.
+# ==============================================================================
+"""Tests for data input for speech commands."""
+
+import os
+
+import numpy as np
+import tensorflow as tf
+
+
+from tensorflow.examples.speech_commands import input_data
+from tensorflow.examples.speech_commands import models
+from tensorflow.python.framework import test_util
+from tensorflow.python.platform import test
+
+
+class InputDataTest(test.TestCase):
+
+  def _getWavData(self):
+    with self.cached_session():
+      sample_data = tf.zeros([32000, 2])
+      wav_encoder = tf.audio.encode_wav(sample_data, 16000)
+      wav_data = self.evaluate(wav_encoder)
+    return wav_data
+
+  def _saveTestWavFile(self, filename, wav_data):
+    with open(filename, "wb") as f:
+      f.write(wav_data)
+
+  def _saveWavFolders(self, root_dir, labels, how_many):
+    wav_data = self._getWavData()
+    for label in labels:
+      dir_name = os.path.join(root_dir, label)
+      os.mkdir(dir_name)
+      for i in range(how_many):
+        file_path = os.path.join(dir_name, "some_audio_%d.wav" % i)
+        self._saveTestWavFile(file_path, wav_data)
+
+  def _model_settings(self):
+    return {
+        "desired_samples": 160,
+        "fingerprint_size": 40,
+        "label_count": 4,
+        "window_size_samples": 100,
+        "window_stride_samples": 100,
+        "fingerprint_width": 40,
+        "preprocess": "mfcc",
+    }
+
+  def _runGetDataTest(self, preprocess, window_length_ms):
+    tmp_dir = self.get_temp_dir()
+    wav_dir = os.path.join(tmp_dir, "wavs")
+    os.mkdir(wav_dir)
+    self._saveWavFolders(wav_dir, ["a", "b", "c"], 100)
+    background_dir = os.path.join(wav_dir, "_background_noise_")
+    os.mkdir(background_dir)
+    wav_data = self._getWavData()
+    for i in range(10):
+      file_path = os.path.join(background_dir, "background_audio_%d.wav" % i)
+      self._saveTestWavFile(file_path, wav_data)
+    model_settings = models.prepare_model_settings(
+        4, 16000, 1000, window_length_ms, 20, 40, preprocess)
+    with self.cached_session() as sess:
+      audio_processor = input_data.AudioProcessor(
+          "", wav_dir, 10, 10, ["a", "b"], 10, 10, model_settings, tmp_dir)
+      result_data, result_labels = audio_processor.get_data(
+          10, 0, model_settings, 0.3, 0.1, 100, "training", sess)
+      self.assertEqual(10, len(result_data))
+      self.assertEqual(10, len(result_labels))
+
+  def testPrepareWordsList(self):
+    words_list = ["a", "b"]
+    self.assertGreater(
+        len(input_data.prepare_words_list(words_list)), len(words_list))
+
+  def testWhichSet(self):
+    self.assertEqual(
+        input_data.which_set("foo.wav", 10, 10),
+        input_data.which_set("foo.wav", 10, 10))
+    self.assertEqual(
+        input_data.which_set("foo_nohash_0.wav", 10, 10),
+        input_data.which_set("foo_nohash_1.wav", 10, 10))
+
+  @test_util.run_deprecated_v1
+  def testPrepareDataIndex(self):
+    tmp_dir = self.get_temp_dir()
+    self._saveWavFolders(tmp_dir, ["a", "b", "c"], 100)
+    audio_processor = input_data.AudioProcessor("", tmp_dir, 10, 10,
+                                                ["a", "b"], 10, 10,
+                                                self._model_settings(), tmp_dir)
+    self.assertLess(0, audio_processor.set_size("training"))
+    self.assertIn("training", audio_processor.data_index)
+    self.assertIn("validation", audio_processor.data_index)
+    self.assertIn("testing", audio_processor.data_index)
+    self.assertEqual(input_data.UNKNOWN_WORD_INDEX,
+                     audio_processor.word_to_index["c"])
+
+  def testPrepareDataIndexEmpty(self):
+    tmp_dir = self.get_temp_dir()
+    self._saveWavFolders(tmp_dir, ["a", "b", "c"], 0)
+    with self.assertRaises(Exception) as e:
+      _ = input_data.AudioProcessor("", tmp_dir, 10, 10, ["a", "b"], 10, 10,
+                                    self._model_settings(), tmp_dir)
+    self.assertIn("No .wavs found", str(e.exception))
+
+  def testPrepareDataIndexMissing(self):
+    tmp_dir = self.get_temp_dir()
+    self._saveWavFolders(tmp_dir, ["a", "b", "c"], 100)
+    with self.assertRaises(Exception) as e:
+      _ = input_data.AudioProcessor("", tmp_dir, 10, 10, ["a", "b", "d"], 10,
+                                    10, self._model_settings(), tmp_dir)
+    self.assertIn("Expected to find", str(e.exception))
+
+  @test_util.run_deprecated_v1
+  def testPrepareBackgroundData(self):
+    tmp_dir = self.get_temp_dir()
+    background_dir = os.path.join(tmp_dir, "_background_noise_")
+    os.mkdir(background_dir)
+    wav_data = self._getWavData()
+    for i in range(10):
+      file_path = os.path.join(background_dir, "background_audio_%d.wav" % i)
+      self._saveTestWavFile(file_path, wav_data)
+    self._saveWavFolders(tmp_dir, ["a", "b", "c"], 100)
+    audio_processor = input_data.AudioProcessor("", tmp_dir, 10, 10,
+                                                ["a", "b"], 10, 10,
+                                                self._model_settings(), tmp_dir)
+    self.assertEqual(10, len(audio_processor.background_data))
+
+  def testLoadWavFile(self):
+    tmp_dir = self.get_temp_dir()
+    file_path = os.path.join(tmp_dir, "load_test.wav")
+    wav_data = self._getWavData()
+    self._saveTestWavFile(file_path, wav_data)
+    sample_data = input_data.load_wav_file(file_path)
+    self.assertIsNotNone(sample_data)
+
+  def testSaveWavFile(self):
+    tmp_dir = self.get_temp_dir()
+    file_path = os.path.join(tmp_dir, "load_test.wav")
+    save_data = np.zeros([16000, 1])
+    input_data.save_wav_file(file_path, save_data, 16000)
+    loaded_data = input_data.load_wav_file(file_path)
+    self.assertIsNotNone(loaded_data)
+    self.assertEqual(16000, len(loaded_data))
+
+  @test_util.run_deprecated_v1
+  def testPrepareProcessingGraph(self):
+    tmp_dir = self.get_temp_dir()
+    wav_dir = os.path.join(tmp_dir, "wavs")
+    os.mkdir(wav_dir)
+    self._saveWavFolders(wav_dir, ["a", "b", "c"], 100)
+    background_dir = os.path.join(wav_dir, "_background_noise_")
+    os.mkdir(background_dir)
+    wav_data = self._getWavData()
+    for i in range(10):
+      file_path = os.path.join(background_dir, "background_audio_%d.wav" % i)
+      self._saveTestWavFile(file_path, wav_data)
+    model_settings = {
+        "desired_samples": 160,
+        "fingerprint_size": 40,
+        "label_count": 4,
+        "window_size_samples": 100,
+        "window_stride_samples": 100,
+        "fingerprint_width": 40,
+        "preprocess": "mfcc",
+    }
+    audio_processor = input_data.AudioProcessor("", wav_dir, 10, 10, ["a", "b"],
+                                                10, 10, model_settings, tmp_dir)
+    self.assertIsNotNone(audio_processor.wav_filename_placeholder_)
+    self.assertIsNotNone(audio_processor.foreground_volume_placeholder_)
+    self.assertIsNotNone(audio_processor.time_shift_padding_placeholder_)
+    self.assertIsNotNone(audio_processor.time_shift_offset_placeholder_)
+    self.assertIsNotNone(audio_processor.background_data_placeholder_)
+    self.assertIsNotNone(audio_processor.background_volume_placeholder_)
+    self.assertIsNotNone(audio_processor.output_)
+
+  @test_util.run_deprecated_v1
+  def testGetDataAverage(self):
+    self._runGetDataTest("average", 10)
+
+  @test_util.run_deprecated_v1
+  def testGetDataAverageLongWindow(self):
+    self._runGetDataTest("average", 30)
+
+  @test_util.run_deprecated_v1
+  def testGetDataMfcc(self):
+    self._runGetDataTest("mfcc", 30)
+
+  @test_util.run_deprecated_v1
+  def testGetDataMicro(self):
+    self._runGetDataTest("micro", 20)
+
+  @test_util.run_deprecated_v1
+  def testGetUnprocessedData(self):
+    tmp_dir = self.get_temp_dir()
+    wav_dir = os.path.join(tmp_dir, "wavs")
+    os.mkdir(wav_dir)
+    self._saveWavFolders(wav_dir, ["a", "b", "c"], 100)
+    model_settings = {
+        "desired_samples": 160,
+        "fingerprint_size": 40,
+        "label_count": 4,
+        "window_size_samples": 100,
+        "window_stride_samples": 100,
+        "fingerprint_width": 40,
+        "preprocess": "mfcc",
+    }
+    audio_processor = input_data.AudioProcessor("", wav_dir, 10, 10, ["a", "b"],
+                                                10, 10, model_settings, tmp_dir)
+    result_data, result_labels = audio_processor.get_unprocessed_data(
+        10, model_settings, "training")
+    self.assertEqual(10, len(result_data))
+    self.assertEqual(10, len(result_labels))
+
+  @test_util.run_deprecated_v1
+  def testGetFeaturesForWav(self):
+    tmp_dir = self.get_temp_dir()
+    wav_dir = os.path.join(tmp_dir, "wavs")
+    os.mkdir(wav_dir)
+    self._saveWavFolders(wav_dir, ["a", "b", "c"], 1)
+    desired_samples = 1600
+    model_settings = {
+        "desired_samples": desired_samples,
+        "fingerprint_size": 40,
+        "label_count": 4,
+        "window_size_samples": 100,
+        "window_stride_samples": 100,
+        "fingerprint_width": 40,
+        "average_window_width": 6,
+        "preprocess": "average",
+    }
+    with self.cached_session() as sess:
+      audio_processor = input_data.AudioProcessor(
+          "", wav_dir, 10, 10, ["a", "b"], 10, 10, model_settings, tmp_dir)
+      sample_data = np.zeros([desired_samples, 1])
+      for i in range(desired_samples):
+        phase = i % 4
+        if phase == 0:
+          sample_data[i, 0] = 0
+        elif phase == 1:
+          sample_data[i, 0] = -1
+        elif phase == 2:
+          sample_data[i, 0] = 0
+        elif phase == 3:
+          sample_data[i, 0] = 1
+      test_wav_path = os.path.join(tmp_dir, "test_wav.wav")
+      input_data.save_wav_file(test_wav_path, sample_data, 16000)
+
+      results = audio_processor.get_features_for_wav(test_wav_path,
+                                                     model_settings, sess)
+      spectrogram = results[0]
+      self.assertEqual(1, spectrogram.shape[0])
+      self.assertEqual(16, spectrogram.shape[1])
+      self.assertEqual(11, spectrogram.shape[2])
+      self.assertNear(0, spectrogram[0, 0, 0], 0.1)
+      self.assertNear(200, spectrogram[0, 0, 5], 0.1)
+
+  def testGetFeaturesRange(self):
+    model_settings = {
+        "preprocess": "average",
+    }
+    features_min, _ = input_data.get_features_range(model_settings)
+    self.assertNear(0.0, features_min, 1e-5)
+
+  def testGetMfccFeaturesRange(self):
+    model_settings = {
+        "preprocess": "mfcc",
+    }
+    features_min, features_max = input_data.get_features_range(model_settings)
+    self.assertLess(features_min, features_max)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/examples/speech_commands/label_wav.cc

@@ -0,0 +1,194 @@
+/* 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 <algorithm>
+#include <cstdint>
+#include <fstream>
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "xla/tsl/platform/env.h"
+#include "xla/tsl/platform/types.h"
+#include "xla/tsl/util/command_line_flags.h"
+#include "tensorflow/core/framework/graph.pb.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/tensor_types.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/init_main.h"
+#include "tensorflow/core/platform/logging.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/platform/tstring.h"
+#include "tensorflow/core/public/session.h"
+#include "tensorflow/core/public/session_options.h"
+#include "tensorflow/core/util/command_line_flags.h"
+
+// These are all common classes it's handy to reference with no namespace.
+using tensorflow::Flag;
+using tensorflow::int32;
+using tensorflow::Status;
+using tensorflow::string;
+using tensorflow::Tensor;
+using tensorflow::tstring;
+
+namespace {
+
+// Reads a model graph definition from disk, and creates a session object you
+// can use to run it.
+Status LoadGraph(const string& graph_file_name,
+                 std::unique_ptr<tensorflow::Session>* session) {
+  tensorflow::GraphDef graph_def;
+  Status load_graph_status =
+      ReadBinaryProto(tensorflow::Env::Default(), graph_file_name, &graph_def);
+  if (!load_graph_status.ok()) {
+    return absl::NotFoundError(absl::StrCat("Failed to load compute graph at '",
+                                            graph_file_name, "'"));
+  }
+  session->reset(tensorflow::NewSession(tensorflow::SessionOptions()));
+  Status session_create_status = (*session)->Create(graph_def);
+  if (!session_create_status.ok()) {
+    return session_create_status;
+  }
+  return absl::OkStatus();
+}
+
+// Takes a file name, and loads a list of labels from it, one per line, and
+// returns a vector of the strings.
+Status ReadLabelsFile(const string& file_name, std::vector<string>* result) {
+  std::ifstream file(file_name);
+  if (!file) {
+    return absl::NotFoundError(
+        absl::StrCat("Labels file ", file_name, " not found."));
+  }
+  result->clear();
+  string line;
+  while (std::getline(file, line)) {
+    result->push_back(line);
+  }
+  return absl::OkStatus();
+}
+
+// Analyzes the output of the graph to retrieve the highest scores and
+// their positions in the tensor.
+void GetTopLabels(const std::vector<Tensor>& outputs, int how_many_labels,
+                  Tensor* out_indices, Tensor* out_scores) {
+  const Tensor& unsorted_scores_tensor = outputs[0];
+  auto unsorted_scores_flat = unsorted_scores_tensor.flat<float>();
+  std::vector<std::pair<int, float>> scores;
+  scores.reserve(unsorted_scores_flat.size());
+  for (int i = 0; i < unsorted_scores_flat.size(); ++i) {
+    scores.push_back(std::pair<int, float>({i, unsorted_scores_flat(i)}));
+  }
+  std::sort(scores.begin(), scores.end(),
+            [](const std::pair<int, float>& left,
+               const std::pair<int, float>& right) {
+              return left.second > right.second;
+            });
+  scores.resize(how_many_labels);
+  Tensor sorted_indices(tensorflow::DT_INT32, {how_many_labels});
+  Tensor sorted_scores(tensorflow::DT_FLOAT, {how_many_labels});
+  for (int i = 0; i < scores.size(); ++i) {
+    sorted_indices.flat<int>()(i) = scores[i].first;
+    sorted_scores.flat<float>()(i) = scores[i].second;
+  }
+  *out_indices = sorted_indices;
+  *out_scores = sorted_scores;
+}
+
+}  // namespace
+
+int main(int argc, char* argv[]) {
+  string wav = "";
+  string graph = "";
+  string labels = "";
+  string input_name = "wav_data";
+  string output_name = "labels_softmax";
+  int32_t how_many_labels = 3;
+  std::vector<Flag> flag_list = {
+      Flag("wav", &wav, "audio file to be identified"),
+      Flag("graph", &graph, "model to be executed"),
+      Flag("labels", &labels, "path to file containing labels"),
+      Flag("input_name", &input_name, "name of input node in model"),
+      Flag("output_name", &output_name, "name of output node in model"),
+      Flag("how_many_labels", &how_many_labels, "number of results to show"),
+  };
+  string usage = tensorflow::Flags::Usage(argv[0], flag_list);
+  const bool parse_result = tensorflow::Flags::Parse(&argc, argv, flag_list);
+  if (!parse_result) {
+    LOG(ERROR) << usage;
+    return -1;
+  }
+
+  // We need to call this to set up global state for TensorFlow.
+  tensorflow::port::InitMain(argv[0], &argc, &argv);
+  if (argc > 1) {
+    LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage;
+    return -1;
+  }
+
+  // First we load and initialize the model.
+  std::unique_ptr<tensorflow::Session> session;
+  Status load_graph_status = LoadGraph(graph, &session);
+  if (!load_graph_status.ok()) {
+    LOG(ERROR) << load_graph_status;
+    return -1;
+  }
+
+  std::vector<string> labels_list;
+  Status read_labels_status = ReadLabelsFile(labels, &labels_list);
+  if (!read_labels_status.ok()) {
+    LOG(ERROR) << read_labels_status;
+    return -1;
+  }
+
+  string wav_string;
+  Status read_wav_status = tensorflow::ReadFileToString(
+      tensorflow::Env::Default(), wav, &wav_string);
+  if (!read_wav_status.ok()) {
+    LOG(ERROR) << read_wav_status;
+    return -1;
+  }
+  Tensor wav_tensor(tensorflow::DT_STRING, tensorflow::TensorShape({}));
+  wav_tensor.scalar<tstring>()() = wav_string;
+
+  // Actually run the audio through the model.
+  std::vector<Tensor> outputs;
+  Status run_status =
+      session->Run({{input_name, wav_tensor}}, {output_name}, {}, &outputs);
+  if (!run_status.ok()) {
+    LOG(ERROR) << "Running model failed: " << run_status;
+    return -1;
+  }
+
+  Tensor indices;
+  Tensor scores;
+  GetTopLabels(outputs, how_many_labels, &indices, &scores);
+  tensorflow::TTypes<float>::Flat scores_flat = scores.flat<float>();
+  tensorflow::TTypes<int32>::Flat indices_flat = indices.flat<int32>();
+  for (int pos = 0; pos < how_many_labels; ++pos) {
+    const int label_index = indices_flat(pos);
+    const float score = scores_flat(pos);
+    LOG(INFO) << labels_list[label_index] << " (" << label_index
+              << "): " << score;
+  }
+
+  return 0;
+}

tensorflow/examples/speech_commands/label_wav.py

@@ -0,0 +1,125 @@
+# 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.
+# ==============================================================================
+r"""Runs a trained audio graph against a WAVE file and reports the results.
+
+The model, labels and .wav file specified in the arguments will be loaded, and
+then the predictions from running the model against the audio data will be
+printed to the console. This is a useful script for sanity checking trained
+models, and as an example of how to use an audio model from Python.
+
+Here's an example of running it:
+
+python tensorflow/examples/speech_commands/label_wav.py \
+--graph=/tmp/my_frozen_graph.pb \
+--labels=/tmp/speech_commands_train/conv_labels.txt \
+--wav=/tmp/speech_dataset/left/a5d485dc_nohash_0.wav
+
+"""
+import argparse
+import sys
+
+import tensorflow as tf
+
+
+FLAGS = None
+
+
+def load_graph(filename):
+  """Unpersists graph from file as default graph."""
+  with tf.io.gfile.GFile(filename, 'rb') as f:
+    graph_def = tf.compat.v1.GraphDef()
+    graph_def.ParseFromString(f.read())
+    tf.import_graph_def(graph_def, name='')
+
+
+def load_labels(filename):
+  """Read in labels, one label per line."""
+  return [line.rstrip() for line in tf.io.gfile.GFile(filename)]
+
+
+def run_graph(wav_data, labels, input_layer_name, output_layer_name,
+              num_top_predictions):
+  """Runs the audio data through the graph and prints predictions."""
+  with tf.compat.v1.Session() as sess:
+    # Feed the audio data as input to the graph.
+    #   predictions  will contain a two-dimensional array, where one
+    #   dimension represents the input image count, and the other has
+    #   predictions per class
+    softmax_tensor = sess.graph.get_tensor_by_name(output_layer_name)
+    predictions, = sess.run(softmax_tensor, {input_layer_name: wav_data})
+
+    # Sort to show labels in order of confidence
+    top_k = predictions.argsort()[-num_top_predictions:][::-1]
+    for node_id in top_k:
+      human_string = labels[node_id]
+      score = predictions[node_id]
+      print('%s (score = %.5f)' % (human_string, score))
+
+    return 0
+
+
+def label_wav(wav, labels, graph, input_name, output_name, how_many_labels):
+  """Loads the model and labels, and runs the inference to print predictions."""
+  if not wav or not tf.io.gfile.exists(wav):
+    raise ValueError('Audio file does not exist at {0}'.format(wav))
+  if not labels or not tf.io.gfile.exists(labels):
+    raise ValueError('Labels file does not exist at {0}'.format(labels))
+
+  if not graph or not tf.io.gfile.exists(graph):
+    raise ValueError('Graph file does not exist at {0}'.format(graph))
+
+  labels_list = load_labels(labels)
+
+  # load graph, which is stored in the default session
+  load_graph(graph)
+
+  with open(wav, 'rb') as wav_file:
+    wav_data = wav_file.read()
+
+  run_graph(wav_data, labels_list, input_name, output_name, how_many_labels)
+
+
+def main(_):
+  """Entry point for script, converts flags to arguments."""
+  label_wav(FLAGS.wav, FLAGS.labels, FLAGS.graph, FLAGS.input_name,
+            FLAGS.output_name, FLAGS.how_many_labels)
+
+
+if __name__ == '__main__':
+  parser = argparse.ArgumentParser()
+  parser.add_argument(
+      '--wav', type=str, default='', help='Audio file to be identified.')
+  parser.add_argument(
+      '--graph', type=str, default='', help='Model to use for identification.')
+  parser.add_argument(
+      '--labels', type=str, default='', help='Path to file containing labels.')
+  parser.add_argument(
+      '--input_name',
+      type=str,
+      default='wav_data:0',
+      help='Name of WAVE data input node in model.')
+  parser.add_argument(
+      '--output_name',
+      type=str,
+      default='labels_softmax:0',
+      help='Name of node outputting a prediction in the model.')
+  parser.add_argument(
+      '--how_many_labels',
+      type=int,
+      default=3,
+      help='Number of results to show.')
+
+  FLAGS, unparsed = parser.parse_known_args()
+  tf.compat.v1.app.run(main=main, argv=[sys.argv[0]] + unparsed)

tensorflow/examples/speech_commands/label_wav_dir.py

@@ -0,0 +1,128 @@
+# 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.
+# ==============================================================================
+r"""Runs a trained audio graph against WAVE files and reports the results.
+
+The model, labels and .wav files specified in the arguments will be loaded, and
+then the predictions from running the model against the audio data will be
+printed to the console. This is a useful script for sanity checking trained
+models, and as an example of how to use an audio model from Python.
+
+Here's an example of running it:
+
+python tensorflow/examples/speech_commands/label_wav_dir.py \
+--graph=/tmp/my_frozen_graph.pb \
+--labels=/tmp/speech_commands_train/conv_labels.txt \
+--wav_dir=/tmp/speech_dataset/left
+
+"""
+import argparse
+import glob
+import sys
+
+import tensorflow as tf
+
+
+FLAGS = None
+
+
+def load_graph(filename):
+  """Unpersists graph from file as default graph."""
+  with tf.io.gfile.GFile(filename, 'rb') as f:
+    graph_def = tf.compat.v1.GraphDef()
+    graph_def.ParseFromString(f.read())
+    tf.import_graph_def(graph_def, name='')
+
+
+def load_labels(filename):
+  """Read in labels, one label per line."""
+  return [line.rstrip() for line in tf.io.gfile.GFile(filename)]
+
+
+def run_graph(wav_dir, labels, input_layer_name, output_layer_name,
+              num_top_predictions):
+  """Runs the audio data through the graph and prints predictions."""
+  with tf.compat.v1.Session() as sess:
+    # Feed the audio data as input to the graph.
+    #   predictions  will contain a two-dimensional array, where one
+    #   dimension represents the input image count, and the other has
+    #   predictions per class
+    for wav_path in glob.glob(wav_dir + '/*.wav'):
+      if not wav_path or not tf.io.gfile.exists(wav_path):
+        raise ValueError('Audio file does not exist at {0}'.format(wav_path))
+      with open(wav_path, 'rb') as wav_file:
+        wav_data = wav_file.read()
+
+      softmax_tensor = sess.graph.get_tensor_by_name(output_layer_name)
+      predictions, = sess.run(softmax_tensor, {input_layer_name: wav_data})
+
+      # Sort to show labels in order of confidence
+      print('\n%s' % (wav_path.split('/')[-1]))
+      top_k = predictions.argsort()[-num_top_predictions:][::-1]
+      for node_id in top_k:
+        human_string = labels[node_id]
+        score = predictions[node_id]
+        print('%s (score = %.5f)' % (human_string, score))
+
+    return 0
+
+
+def label_wav(wav_dir, labels, graph, input_name, output_name, how_many_labels):
+  """Loads the model and labels, and runs the inference to print predictions."""
+  if not labels or not tf.io.gfile.exists(labels):
+    raise ValueError('Labels file does not exist at {0}'.format(labels))
+
+  if not graph or not tf.io.gfile.exists(graph):
+    raise ValueError('Graph file does not exist at {0}'.format(graph))
+
+  labels_list = load_labels(labels)
+
+  # load graph, which is stored in the default session
+  load_graph(graph)
+
+  run_graph(wav_dir, labels_list, input_name, output_name, how_many_labels)
+
+
+def main(_):
+  """Entry point for script, converts flags to arguments."""
+  label_wav(FLAGS.wav_dir, FLAGS.labels, FLAGS.graph, FLAGS.input_name,
+            FLAGS.output_name, FLAGS.how_many_labels)
+
+
+if __name__ == '__main__':
+  parser = argparse.ArgumentParser()
+  parser.add_argument(
+      '--wav_dir', type=str, default='', help='Audio file to be identified.')
+  parser.add_argument(
+      '--graph', type=str, default='', help='Model to use for identification.')
+  parser.add_argument(
+      '--labels', type=str, default='', help='Path to file containing labels.')
+  parser.add_argument(
+      '--input_name',
+      type=str,
+      default='wav_data:0',
+      help='Name of WAVE data input node in model.')
+  parser.add_argument(
+      '--output_name',
+      type=str,
+      default='labels_softmax:0',
+      help='Name of node outputting a prediction in the model.')
+  parser.add_argument(
+      '--how_many_labels',
+      type=int,
+      default=3,
+      help='Number of results to show.')
+
+  FLAGS, unparsed = parser.parse_known_args()
+  tf.compat.v1.app.run(main=main, argv=[sys.argv[0]] + unparsed)

tensorflow/examples/speech_commands/label_wav_test.py

@@ -0,0 +1,59 @@
+# 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.
+# ==============================================================================
+"""Tests for WAVE file labeling tool."""
+
+import os
+
+import tensorflow as tf
+
+from tensorflow.examples.speech_commands import label_wav
+from tensorflow.python.platform import test
+
+
+class LabelWavTest(test.TestCase):
+
+  def _getWavData(self):
+    with self.cached_session():
+      sample_data = tf.zeros([1000, 2])
+      wav_encoder = tf.audio.encode_wav(sample_data, 16000)
+      wav_data = self.evaluate(wav_encoder)
+    return wav_data
+
+  def _saveTestWavFile(self, filename, wav_data):
+    with open(filename, "wb") as f:
+      f.write(wav_data)
+
+  def testLabelWav(self):
+    tmp_dir = self.get_temp_dir()
+    wav_data = self._getWavData()
+    wav_filename = os.path.join(tmp_dir, "wav_file.wav")
+    self._saveTestWavFile(wav_filename, wav_data)
+    input_name = "test_input"
+    output_name = "test_output"
+    graph_filename = os.path.join(tmp_dir, "test_graph.pb")
+    with tf.compat.v1.Session() as sess:
+      tf.compat.v1.placeholder(tf.string, name=input_name)
+      tf.zeros([1, 3], name=output_name)
+      with open(graph_filename, "wb") as f:
+        f.write(sess.graph.as_graph_def().SerializeToString())
+    labels_filename = os.path.join(tmp_dir, "test_labels.txt")
+    with open(labels_filename, "w") as f:
+      f.write("a\nb\nc\n")
+    label_wav.label_wav(wav_filename, labels_filename, graph_filename,
+                        input_name + ":0", output_name + ":0", 3)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/examples/speech_commands/models.py

@@ -0,0 +1,893 @@
+# 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.
+# ==============================================================================
+"""Model definitions for simple speech recognition.
+
+"""
+import math
+
+import tensorflow as tf
+
+
+def _next_power_of_two(x):
+  """Calculates the smallest enclosing power of two for an input.
+
+  Args:
+    x: Positive float or integer number.
+
+  Returns:
+    Next largest power of two integer.
+  """
+  return 1 if x == 0 else 2**(int(x) - 1).bit_length()
+
+
+def prepare_model_settings(label_count, sample_rate, clip_duration_ms,
+                           window_size_ms, window_stride_ms, feature_bin_count,
+                           preprocess):
+  """Calculates common settings needed for all models.
+
+  Args:
+    label_count: How many classes are to be recognized.
+    sample_rate: Number of audio samples per second.
+    clip_duration_ms: Length of each audio clip to be analyzed.
+    window_size_ms: Duration of frequency analysis window.
+    window_stride_ms: How far to move in time between frequency windows.
+    feature_bin_count: Number of frequency bins to use for analysis.
+    preprocess: How the spectrogram is processed to produce features.
+
+  Returns:
+    Dictionary containing common settings.
+
+  Raises:
+    ValueError: If the preprocessing mode isn't recognized.
+  """
+  desired_samples = int(sample_rate * clip_duration_ms / 1000)
+  window_size_samples = int(sample_rate * window_size_ms / 1000)
+  window_stride_samples = int(sample_rate * window_stride_ms / 1000)
+  length_minus_window = (desired_samples - window_size_samples)
+  if length_minus_window < 0:
+    spectrogram_length = 0
+  else:
+    spectrogram_length = 1 + int(length_minus_window / window_stride_samples)
+  if preprocess == 'average':
+    fft_bin_count = 1 + (_next_power_of_two(window_size_samples) / 2)
+    average_window_width = int(math.floor(fft_bin_count / feature_bin_count))
+    fingerprint_width = int(math.ceil(fft_bin_count / average_window_width))
+  elif preprocess == 'mfcc':
+    average_window_width = -1
+    fingerprint_width = feature_bin_count
+  elif preprocess == 'micro':
+    average_window_width = -1
+    fingerprint_width = feature_bin_count
+  else:
+    raise ValueError('Unknown preprocess mode "%s" (should be "mfcc",'
+                     ' "average", or "micro")' % (preprocess))
+  fingerprint_size = fingerprint_width * spectrogram_length
+  return {
+      'desired_samples': desired_samples,
+      'window_size_samples': window_size_samples,
+      'window_stride_samples': window_stride_samples,
+      'spectrogram_length': spectrogram_length,
+      'fingerprint_width': fingerprint_width,
+      'fingerprint_size': fingerprint_size,
+      'label_count': label_count,
+      'sample_rate': sample_rate,
+      'preprocess': preprocess,
+      'average_window_width': average_window_width,
+  }
+
+
+def create_model(fingerprint_input, model_settings, model_architecture,
+                 is_training, runtime_settings=None):
+  """Builds a model of the requested architecture compatible with the settings.
+
+  There are many possible ways of deriving predictions from a spectrogram
+  input, so this function provides an abstract interface for creating different
+  kinds of models in a black-box way. You need to pass in a TensorFlow node as
+  the 'fingerprint' input, and this should output a batch of 1D features that
+  describe the audio. Typically this will be derived from a spectrogram that's
+  been run through an MFCC, but in theory it can be any feature vector of the
+  size specified in model_settings['fingerprint_size'].
+
+  The function will build the graph it needs in the current TensorFlow graph,
+  and return the tensorflow output that will contain the 'logits' input to the
+  softmax prediction process. If training flag is on, it will also return a
+  placeholder node that can be used to control the dropout amount.
+
+  See the implementations below for the possible model architectures that can be
+  requested.
+
+  Args:
+    fingerprint_input: TensorFlow node that will output audio feature vectors.
+    model_settings: Dictionary of information about the model.
+    model_architecture: String specifying which kind of model to create.
+    is_training: Whether the model is going to be used for training.
+    runtime_settings: Dictionary of information about the runtime.
+
+  Returns:
+    TensorFlow node outputting logits results, and optionally a dropout
+    placeholder.
+
+  Raises:
+    Exception: If the architecture type isn't recognized.
+  """
+  if model_architecture == 'single_fc':
+    return create_single_fc_model(fingerprint_input, model_settings,
+                                  is_training)
+  elif model_architecture == 'conv':
+    return create_conv_model(fingerprint_input, model_settings, is_training)
+  elif model_architecture == 'low_latency_conv':
+    return create_low_latency_conv_model(fingerprint_input, model_settings,
+                                         is_training)
+  elif model_architecture == 'low_latency_svdf':
+    return create_low_latency_svdf_model(fingerprint_input, model_settings,
+                                         is_training, runtime_settings)
+  elif model_architecture == 'tiny_conv':
+    return create_tiny_conv_model(fingerprint_input, model_settings,
+                                  is_training)
+  elif model_architecture == 'tiny_embedding_conv':
+    return create_tiny_embedding_conv_model(fingerprint_input, model_settings,
+                                            is_training)
+  else:
+    raise Exception('model_architecture argument "' + model_architecture +
+                    '" not recognized, should be one of "single_fc", "conv",' +
+                    ' "low_latency_conv, "low_latency_svdf",' +
+                    ' "tiny_conv", or "tiny_embedding_conv"')
+
+
+def load_variables_from_checkpoint(sess, start_checkpoint):
+  """Utility function to centralize checkpoint restoration.
+
+  Args:
+    sess: TensorFlow session.
+    start_checkpoint: Path to saved checkpoint on disk.
+  """
+  saver = tf.compat.v1.train.Saver(tf.compat.v1.global_variables())
+  saver.restore(sess, start_checkpoint)
+
+
+def create_single_fc_model(fingerprint_input, model_settings, is_training):
+  """Builds a model with a single hidden fully-connected layer.
+
+  This is a very simple model with just one matmul and bias layer. As you'd
+  expect, it doesn't produce very accurate results, but it is very fast and
+  simple, so it's useful for sanity testing.
+
+  Here's the layout of the graph:
+
+  (fingerprint_input)
+          v
+      [MatMul]<-(weights)
+          v
+      [BiasAdd]<-(bias)
+          v
+
+  Args:
+    fingerprint_input: TensorFlow node that will output audio feature vectors.
+    model_settings: Dictionary of information about the model.
+    is_training: Whether the model is going to be used for training.
+
+  Returns:
+    TensorFlow node outputting logits results, and optionally a dropout
+    placeholder.
+  """
+  if is_training:
+    dropout_rate = tf.compat.v1.placeholder(tf.float32, name='dropout_rate')
+  fingerprint_size = model_settings['fingerprint_size']
+  label_count = model_settings['label_count']
+  weights = tf.compat.v1.get_variable(
+      name='weights',
+      initializer=tf.compat.v1.truncated_normal_initializer(stddev=0.001),
+      shape=[fingerprint_size, label_count])
+  bias = tf.compat.v1.get_variable(name='bias',
+                                   initializer=tf.compat.v1.zeros_initializer,
+                                   shape=[label_count])
+  logits = tf.matmul(fingerprint_input, weights) + bias
+  if is_training:
+    return logits, dropout_rate
+  else:
+    return logits
+
+
+def create_conv_model(fingerprint_input, model_settings, is_training):
+  """Builds a standard convolutional model.
+
+  This is roughly the network labeled as 'cnn-trad-fpool3' in the
+  'Convolutional Neural Networks for Small-footprint Keyword Spotting' paper:
+  http://www.isca-speech.org/archive/interspeech_2015/papers/i15_1478.pdf
+
+  Here's the layout of the graph:
+
+  (fingerprint_input)
+          v
+      [Conv2D]<-(weights)
+          v
+      [BiasAdd]<-(bias)
+          v
+        [Relu]
+          v
+      [MaxPool]
+          v
+      [Conv2D]<-(weights)
+          v
+      [BiasAdd]<-(bias)
+          v
+        [Relu]
+          v
+      [MaxPool]
+          v
+      [MatMul]<-(weights)
+          v
+      [BiasAdd]<-(bias)
+          v
+
+  This produces fairly good quality results, but can involve a large number of
+  weight parameters and computations. For a cheaper alternative from the same
+  paper with slightly less accuracy, see 'low_latency_conv' below.
+
+  During training, dropout nodes are introduced after each relu, controlled by a
+  placeholder.
+
+  Args:
+    fingerprint_input: TensorFlow node that will output audio feature vectors.
+    model_settings: Dictionary of information about the model.
+    is_training: Whether the model is going to be used for training.
+
+  Returns:
+    TensorFlow node outputting logits results, and optionally a dropout
+    placeholder.
+  """
+  if is_training:
+    dropout_rate = tf.compat.v1.placeholder(tf.float32, name='dropout_rate')
+  input_frequency_size = model_settings['fingerprint_width']
+  input_time_size = model_settings['spectrogram_length']
+  fingerprint_4d = tf.reshape(fingerprint_input,
+                              [-1, input_time_size, input_frequency_size, 1])
+  first_filter_width = 8
+  first_filter_height = 20
+  first_filter_count = 64
+  first_weights = tf.compat.v1.get_variable(
+      name='first_weights',
+      initializer=tf.compat.v1.truncated_normal_initializer(stddev=0.01),
+      shape=[first_filter_height, first_filter_width, 1, first_filter_count])
+  first_bias = tf.compat.v1.get_variable(
+      name='first_bias',
+      initializer=tf.compat.v1.zeros_initializer,
+      shape=[first_filter_count])
+
+  first_conv = tf.nn.conv2d(input=fingerprint_4d,
+                            filters=first_weights,
+                            strides=[1, 1, 1, 1],
+                            padding='SAME') + first_bias
+  first_relu = tf.nn.relu(first_conv)
+  if is_training:
+    first_dropout = tf.nn.dropout(first_relu, rate=dropout_rate)
+  else:
+    first_dropout = first_relu
+  max_pool = tf.nn.max_pool2d(input=first_dropout,
+                              ksize=[1, 2, 2, 1],
+                              strides=[1, 2, 2, 1],
+                              padding='SAME')
+  second_filter_width = 4
+  second_filter_height = 10
+  second_filter_count = 64
+  second_weights = tf.compat.v1.get_variable(
+      name='second_weights',
+      initializer=tf.compat.v1.truncated_normal_initializer(stddev=0.01),
+      shape=[
+          second_filter_height, second_filter_width, first_filter_count,
+          second_filter_count
+      ])
+  second_bias = tf.compat.v1.get_variable(
+      name='second_bias',
+      initializer=tf.compat.v1.zeros_initializer,
+      shape=[second_filter_count])
+  second_conv = tf.nn.conv2d(input=max_pool,
+                             filters=second_weights,
+                             strides=[1, 1, 1, 1],
+                             padding='SAME') + second_bias
+  second_relu = tf.nn.relu(second_conv)
+  if is_training:
+    second_dropout = tf.nn.dropout(second_relu, rate=dropout_rate)
+  else:
+    second_dropout = second_relu
+  second_conv_shape = second_dropout.get_shape()
+  second_conv_output_width = second_conv_shape[2]
+  second_conv_output_height = second_conv_shape[1]
+  second_conv_element_count = int(
+      second_conv_output_width * second_conv_output_height *
+      second_filter_count)
+  flattened_second_conv = tf.reshape(second_dropout,
+                                     [-1, second_conv_element_count])
+  label_count = model_settings['label_count']
+  final_fc_weights = tf.compat.v1.get_variable(
+      name='final_fc_weights',
+      initializer=tf.compat.v1.truncated_normal_initializer(stddev=0.01),
+      shape=[second_conv_element_count, label_count])
+  final_fc_bias = tf.compat.v1.get_variable(
+      name='final_fc_bias',
+      initializer=tf.compat.v1.zeros_initializer,
+      shape=[label_count])
+  final_fc = tf.matmul(flattened_second_conv, final_fc_weights) + final_fc_bias
+  if is_training:
+    return final_fc, dropout_rate
+  else:
+    return final_fc
+
+
+def create_low_latency_conv_model(fingerprint_input, model_settings,
+                                  is_training):
+  """Builds a convolutional model with low compute requirements.
+
+  This is roughly the network labeled as 'cnn-one-fstride4' in the
+  'Convolutional Neural Networks for Small-footprint Keyword Spotting' paper:
+  http://www.isca-speech.org/archive/interspeech_2015/papers/i15_1478.pdf
+
+  Here's the layout of the graph:
+
+  (fingerprint_input)
+          v
+      [Conv2D]<-(weights)
+          v
+      [BiasAdd]<-(bias)
+          v
+        [Relu]
+          v
+      [MatMul]<-(weights)
+          v
+      [BiasAdd]<-(bias)
+          v
+      [MatMul]<-(weights)
+          v
+      [BiasAdd]<-(bias)
+          v
+      [MatMul]<-(weights)
+          v
+      [BiasAdd]<-(bias)
+          v
+
+  This produces slightly lower quality results than the 'conv' model, but needs
+  fewer weight parameters and computations.
+
+  During training, dropout nodes are introduced after the relu, controlled by a
+  placeholder.
+
+  Args:
+    fingerprint_input: TensorFlow node that will output audio feature vectors.
+    model_settings: Dictionary of information about the model.
+    is_training: Whether the model is going to be used for training.
+
+  Returns:
+    TensorFlow node outputting logits results, and optionally a dropout
+    placeholder.
+  """
+  if is_training:
+    dropout_rate = tf.compat.v1.placeholder(tf.float32, name='dropout_rate')
+  input_frequency_size = model_settings['fingerprint_width']
+  input_time_size = model_settings['spectrogram_length']
+  fingerprint_4d = tf.reshape(fingerprint_input,
+                              [-1, input_time_size, input_frequency_size, 1])
+  first_filter_width = 8
+  first_filter_height = input_time_size
+  first_filter_count = 186
+  first_filter_stride_x = 1
+  first_filter_stride_y = 1
+  first_weights = tf.compat.v1.get_variable(
+      name='first_weights',
+      initializer=tf.compat.v1.truncated_normal_initializer(stddev=0.01),
+      shape=[first_filter_height, first_filter_width, 1, first_filter_count])
+  first_bias = tf.compat.v1.get_variable(
+      name='first_bias',
+      initializer=tf.compat.v1.zeros_initializer,
+      shape=[first_filter_count])
+  first_conv = tf.nn.conv2d(
+      input=fingerprint_4d,
+      filters=first_weights,
+      strides=[1, first_filter_stride_y, first_filter_stride_x, 1],
+      padding='VALID') + first_bias
+  first_relu = tf.nn.relu(first_conv)
+  if is_training:
+    first_dropout = tf.nn.dropout(first_relu, rate=dropout_rate)
+  else:
+    first_dropout = first_relu
+  first_conv_output_width = math.floor(
+      (input_frequency_size - first_filter_width + first_filter_stride_x) /
+      first_filter_stride_x)
+  first_conv_output_height = math.floor(
+      (input_time_size - first_filter_height + first_filter_stride_y) /
+      first_filter_stride_y)
+  first_conv_element_count = int(
+      first_conv_output_width * first_conv_output_height * first_filter_count)
+  flattened_first_conv = tf.reshape(first_dropout,
+                                    [-1, first_conv_element_count])
+  first_fc_output_channels = 128
+  first_fc_weights = tf.compat.v1.get_variable(
+      name='first_fc_weights',
+      initializer=tf.compat.v1.truncated_normal_initializer(stddev=0.01),
+      shape=[first_conv_element_count, first_fc_output_channels])
+  first_fc_bias = tf.compat.v1.get_variable(
+      name='first_fc_bias',
+      initializer=tf.compat.v1.zeros_initializer,
+      shape=[first_fc_output_channels])
+  first_fc = tf.matmul(flattened_first_conv, first_fc_weights) + first_fc_bias
+  if is_training:
+    second_fc_input = tf.nn.dropout(first_fc, rate=dropout_rate)
+  else:
+    second_fc_input = first_fc
+  second_fc_output_channels = 128
+  second_fc_weights = tf.compat.v1.get_variable(
+      name='second_fc_weights',
+      initializer=tf.compat.v1.truncated_normal_initializer(stddev=0.01),
+      shape=[first_fc_output_channels, second_fc_output_channels])
+  second_fc_bias = tf.compat.v1.get_variable(
+      name='second_fc_bias',
+      initializer=tf.compat.v1.zeros_initializer,
+      shape=[second_fc_output_channels])
+  second_fc = tf.matmul(second_fc_input, second_fc_weights) + second_fc_bias
+  if is_training:
+    final_fc_input = tf.nn.dropout(second_fc, rate=dropout_rate)
+  else:
+    final_fc_input = second_fc
+  label_count = model_settings['label_count']
+  final_fc_weights = tf.compat.v1.get_variable(
+      name='final_fc_weights',
+      initializer=tf.compat.v1.truncated_normal_initializer(stddev=0.01),
+      shape=[second_fc_output_channels, label_count])
+  final_fc_bias = tf.compat.v1.get_variable(
+      name='final_fc_bias',
+      initializer=tf.compat.v1.zeros_initializer,
+      shape=[label_count])
+  final_fc = tf.matmul(final_fc_input, final_fc_weights) + final_fc_bias
+  if is_training:
+    return final_fc, dropout_rate
+  else:
+    return final_fc
+
+
+def create_low_latency_svdf_model(fingerprint_input, model_settings,
+                                  is_training, runtime_settings):
+  """Builds an SVDF model with low compute requirements.
+
+  This is based in the topology presented in the 'Compressing Deep Neural
+  Networks using a Rank-Constrained Topology' paper:
+  https://static.googleusercontent.com/media/research.google.com/en//pubs/archive/43813.pdf
+
+  Here's the layout of the graph:
+
+  (fingerprint_input)
+          v
+        [SVDF]<-(weights)
+          v
+      [BiasAdd]<-(bias)
+          v
+        [Relu]
+          v
+      [MatMul]<-(weights)
+          v
+      [BiasAdd]<-(bias)
+          v
+      [MatMul]<-(weights)
+          v
+      [BiasAdd]<-(bias)
+          v
+      [MatMul]<-(weights)
+          v
+      [BiasAdd]<-(bias)
+          v
+
+  This model produces lower recognition accuracy than the 'conv' model above,
+  but requires fewer weight parameters and, significantly fewer computations.
+
+  During training, dropout nodes are introduced after the relu, controlled by a
+  placeholder.
+
+  Args:
+    fingerprint_input: TensorFlow node that will output audio feature vectors.
+    The node is expected to produce a 2D Tensor of shape:
+      [batch, model_settings['fingerprint_width'] *
+              model_settings['spectrogram_length']]
+    with the features corresponding to the same time slot arranged contiguously,
+    and the oldest slot at index [:, 0], and newest at [:, -1].
+    model_settings: Dictionary of information about the model.
+    is_training: Whether the model is going to be used for training.
+    runtime_settings: Dictionary of information about the runtime.
+
+  Returns:
+    TensorFlow node outputting logits results, and optionally a dropout
+    placeholder.
+
+  Raises:
+      ValueError: If the inputs tensor is incorrectly shaped.
+  """
+  if is_training:
+    dropout_rate = tf.compat.v1.placeholder(tf.float32, name='dropout_rate')
+
+  input_frequency_size = model_settings['fingerprint_width']
+  input_time_size = model_settings['spectrogram_length']
+
+  # Validation.
+  input_shape = fingerprint_input.get_shape()
+  if len(input_shape) != 2:
+    raise ValueError('Inputs to `SVDF` should have rank == 2.')
+  if input_shape[-1].value is None:
+    raise ValueError('The last dimension of the input to `SVDF` '
+                     'should be defined. Found `None`.')
+  if input_shape[-1].value % input_frequency_size != 0:
+    raise ValueError('The last dimension of the input to `SVDF` = {0} must be '
+                     'a multiple of the frame size = {1}'.format(
+                         input_shape.shape[-1].value, input_frequency_size))
+
+  # Set number of units (i.e. nodes) and rank.
+  rank = 2
+  num_units = 1280
+  # Number of filters: pairs of feature and time filters.
+  num_filters = rank * num_units
+  # Create the runtime memory: [num_filters, batch, input_time_size]
+  batch = 1
+  memory = tf.compat.v1.get_variable(
+      initializer=tf.compat.v1.zeros_initializer,
+      shape=[num_filters, batch, input_time_size],
+      trainable=False,
+      name='runtime-memory')
+  first_time_flag = tf.compat.v1.get_variable(
+      name='first_time_flag', dtype=tf.int32, initializer=1)
+  # Determine the number of new frames in the input, such that we only operate
+  # on those. For training we do not use the memory, and thus use all frames
+  # provided in the input.
+  # new_fingerprint_input: [batch, num_new_frames*input_frequency_size]
+  if is_training:
+    num_new_frames = input_time_size
+  else:
+    window_stride_ms = int(model_settings['window_stride_samples'] * 1000 /
+                           model_settings['sample_rate'])
+    num_new_frames = tf.cond(
+        pred=tf.equal(first_time_flag, 1),
+        true_fn=lambda: input_time_size,
+        false_fn=lambda: int(runtime_settings['clip_stride_ms'] / window_stride_ms))  # pylint:disable=line-too-long
+  first_time_flag = 0
+  new_fingerprint_input = fingerprint_input[
+      :, -num_new_frames*input_frequency_size:]
+  # Expand to add input channels dimension.
+  new_fingerprint_input = tf.expand_dims(new_fingerprint_input, 2)
+
+  # Create the frequency filters.
+  weights_frequency = tf.compat.v1.get_variable(
+      name='weights_frequency',
+      initializer=tf.compat.v1.truncated_normal_initializer(stddev=0.01),
+      shape=[input_frequency_size, num_filters])
+  # Expand to add input channels dimensions.
+  # weights_frequency: [input_frequency_size, 1, num_filters]
+  weights_frequency = tf.expand_dims(weights_frequency, 1)
+  # Convolve the 1D feature filters sliding over the time dimension.
+  # activations_time: [batch, num_new_frames, num_filters]
+  activations_time = tf.nn.conv1d(input=new_fingerprint_input,
+                                  filters=weights_frequency,
+                                  stride=input_frequency_size,
+                                  padding='VALID')
+  # Rearrange such that we can perform the batched matmul.
+  # activations_time: [num_filters, batch, num_new_frames]
+  activations_time = tf.transpose(a=activations_time, perm=[2, 0, 1])
+
+  # Runtime memory optimization.
+  if not is_training:
+    # We need to drop the activations corresponding to the oldest frames, and
+    # then add those corresponding to the new frames.
+    new_memory = memory[:, :, num_new_frames:]
+    new_memory = tf.concat([new_memory, activations_time], 2)
+    tf.compat.v1.assign(memory, new_memory)
+    activations_time = new_memory
+
+  # Create the time filters.
+  weights_time = tf.compat.v1.get_variable(
+      name='weights_time',
+      initializer=tf.compat.v1.truncated_normal_initializer(stddev=0.01),
+      shape=[num_filters, input_time_size])
+  # Apply the time filter on the outputs of the feature filters.
+  # weights_time: [num_filters, input_time_size, 1]
+  # outputs: [num_filters, batch, 1]
+  weights_time = tf.expand_dims(weights_time, 2)
+  outputs = tf.matmul(activations_time, weights_time)
+  # Split num_units and rank into separate dimensions (the remaining
+  # dimension is the input_shape[0] -i.e. batch size). This also squeezes
+  # the last dimension, since it's not used.
+  # [num_filters, batch, 1] => [num_units, rank, batch]
+  outputs = tf.reshape(outputs, [num_units, rank, -1])
+  # Sum the rank outputs per unit => [num_units, batch].
+  units_output = tf.reduce_sum(input_tensor=outputs, axis=1)
+  # Transpose to shape [batch, num_units]
+  units_output = tf.transpose(a=units_output)
+
+  # Apply bias.
+  bias = tf.compat.v1.get_variable(name='bias',
+                                   initializer=tf.compat.v1.zeros_initializer,
+                                   shape=[num_units])
+  first_bias = tf.nn.bias_add(units_output, bias)
+
+  # Relu.
+  first_relu = tf.nn.relu(first_bias)
+
+  if is_training:
+    first_dropout = tf.nn.dropout(first_relu, rate=dropout_rate)
+  else:
+    first_dropout = first_relu
+
+  first_fc_output_channels = 256
+  first_fc_weights = tf.compat.v1.get_variable(
+      name='first_fc_weights',
+      initializer=tf.compat.v1.truncated_normal_initializer(stddev=0.01),
+      shape=[num_units, first_fc_output_channels])
+  first_fc_bias = tf.compat.v1.get_variable(
+      name='first_fc_bias',
+      initializer=tf.compat.v1.zeros_initializer,
+      shape=[first_fc_output_channels])
+  first_fc = tf.matmul(first_dropout, first_fc_weights) + first_fc_bias
+  if is_training:
+    second_fc_input = tf.nn.dropout(first_fc, rate=dropout_rate)
+  else:
+    second_fc_input = first_fc
+  second_fc_output_channels = 256
+  second_fc_weights = tf.compat.v1.get_variable(
+      name='second_fc_weights',
+      initializer=tf.compat.v1.truncated_normal_initializer(stddev=0.01),
+      shape=[first_fc_output_channels, second_fc_output_channels])
+  second_fc_bias = tf.compat.v1.get_variable(
+      name='second_fc_bias',
+      initializer=tf.compat.v1.zeros_initializer,
+      shape=[second_fc_output_channels])
+  second_fc = tf.matmul(second_fc_input, second_fc_weights) + second_fc_bias
+  if is_training:
+    final_fc_input = tf.nn.dropout(second_fc, rate=dropout_rate)
+  else:
+    final_fc_input = second_fc
+  label_count = model_settings['label_count']
+  final_fc_weights = tf.compat.v1.get_variable(
+      name='final_fc_weights',
+      initializer=tf.compat.v1.truncated_normal_initializer(stddev=0.01),
+      shape=[second_fc_output_channels, label_count])
+  final_fc_bias = tf.compat.v1.get_variable(
+      name='final_fc_bias',
+      initializer=tf.compat.v1.zeros_initializer,
+      shape=[label_count])
+  final_fc = tf.matmul(final_fc_input, final_fc_weights) + final_fc_bias
+  if is_training:
+    return final_fc, dropout_rate
+  else:
+    return final_fc
+
+
+def create_tiny_conv_model(fingerprint_input, model_settings, is_training):
+  """Builds a convolutional model aimed at microcontrollers.
+
+  Devices like DSPs and microcontrollers can have very small amounts of
+  memory and limited processing power. This model is designed to use less
+  than 20KB of working RAM, and fit within 32KB of read-only (flash) memory.
+
+  Here's the layout of the graph:
+
+  (fingerprint_input)
+          v
+      [Conv2D]<-(weights)
+          v
+      [BiasAdd]<-(bias)
+          v
+        [Relu]
+          v
+      [MatMul]<-(weights)
+          v
+      [BiasAdd]<-(bias)
+          v
+
+  This doesn't produce particularly accurate results, but it's designed to be
+  used as the first stage of a pipeline, running on a low-energy piece of
+  hardware that can always be on, and then wake higher-power chips when a
+  possible utterance has been found, so that more accurate analysis can be done.
+
+  During training, a dropout node is introduced after the relu, controlled by a
+  placeholder.
+
+  Args:
+    fingerprint_input: TensorFlow node that will output audio feature vectors.
+    model_settings: Dictionary of information about the model.
+    is_training: Whether the model is going to be used for training.
+
+  Returns:
+    TensorFlow node outputting logits results, and optionally a dropout
+    placeholder.
+  """
+  if is_training:
+    dropout_rate = tf.compat.v1.placeholder(tf.float32, name='dropout_rate')
+  input_frequency_size = model_settings['fingerprint_width']
+  input_time_size = model_settings['spectrogram_length']
+  fingerprint_4d = tf.reshape(fingerprint_input,
+                              [-1, input_time_size, input_frequency_size, 1])
+  first_filter_width = 8
+  first_filter_height = 10
+  first_filter_count = 8
+  first_weights = tf.compat.v1.get_variable(
+      name='first_weights',
+      initializer=tf.compat.v1.truncated_normal_initializer(stddev=0.01),
+      shape=[first_filter_height, first_filter_width, 1, first_filter_count])
+  first_bias = tf.compat.v1.get_variable(
+      name='first_bias',
+      initializer=tf.compat.v1.zeros_initializer,
+      shape=[first_filter_count])
+  first_conv_stride_x = 2
+  first_conv_stride_y = 2
+  first_conv = tf.nn.conv2d(
+      input=fingerprint_4d, filters=first_weights,
+      strides=[1, first_conv_stride_y, first_conv_stride_x, 1],
+      padding='SAME') + first_bias
+  first_relu = tf.nn.relu(first_conv)
+  if is_training:
+    first_dropout = tf.nn.dropout(first_relu, rate=dropout_rate)
+  else:
+    first_dropout = first_relu
+  first_dropout_shape = first_dropout.get_shape()
+  first_dropout_output_width = first_dropout_shape[2]
+  first_dropout_output_height = first_dropout_shape[1]
+  first_dropout_element_count = int(
+      first_dropout_output_width * first_dropout_output_height *
+      first_filter_count)
+  flattened_first_dropout = tf.reshape(first_dropout,
+                                       [-1, first_dropout_element_count])
+  label_count = model_settings['label_count']
+  final_fc_weights = tf.compat.v1.get_variable(
+      name='final_fc_weights',
+      initializer=tf.compat.v1.truncated_normal_initializer(stddev=0.01),
+      shape=[first_dropout_element_count, label_count])
+  final_fc_bias = tf.compat.v1.get_variable(
+      name='final_fc_bias',
+      initializer=tf.compat.v1.zeros_initializer,
+      shape=[label_count])
+  final_fc = (
+      tf.matmul(flattened_first_dropout, final_fc_weights) + final_fc_bias)
+  if is_training:
+    return final_fc, dropout_rate
+  else:
+    return final_fc
+
+
+def create_tiny_embedding_conv_model(fingerprint_input, model_settings,
+                                     is_training):
+  """Builds a convolutional model aimed at microcontrollers.
+
+  Devices like DSPs and microcontrollers can have very small amounts of
+  memory and limited processing power. This model is designed to use less
+  than 20KB of working RAM, and fit within 32KB of read-only (flash) memory.
+
+  Here's the layout of the graph:
+
+  (fingerprint_input)
+          v
+      [Conv2D]<-(weights)
+          v
+      [BiasAdd]<-(bias)
+          v
+        [Relu]
+          v
+      [Conv2D]<-(weights)
+          v
+      [BiasAdd]<-(bias)
+          v
+        [Relu]
+          v
+      [Conv2D]<-(weights)
+          v
+      [BiasAdd]<-(bias)
+          v
+        [Relu]
+          v
+      [MatMul]<-(weights)
+          v
+      [BiasAdd]<-(bias)
+          v
+
+  This doesn't produce particularly accurate results, but it's designed to be
+  used as the first stage of a pipeline, running on a low-energy piece of
+  hardware that can always be on, and then wake higher-power chips when a
+  possible utterance has been found, so that more accurate analysis can be done.
+
+  During training, a dropout node is introduced after the relu, controlled by a
+  placeholder.
+
+  Args:
+    fingerprint_input: TensorFlow node that will output audio feature vectors.
+    model_settings: Dictionary of information about the model.
+    is_training: Whether the model is going to be used for training.
+
+  Returns:
+    TensorFlow node outputting logits results, and optionally a dropout
+    placeholder.
+  """
+  if is_training:
+    dropout_rate = tf.compat.v1.placeholder(tf.float32, name='dropout_rate')
+  input_frequency_size = model_settings['fingerprint_width']
+  input_time_size = model_settings['spectrogram_length']
+  fingerprint_4d = tf.reshape(fingerprint_input,
+                              [-1, input_time_size, input_frequency_size, 1])
+
+  first_filter_width = 8
+  first_filter_height = 10
+  first_filter_count = 8
+  first_weights = tf.compat.v1.get_variable(
+      name='first_weights',
+      initializer=tf.compat.v1.truncated_normal_initializer(stddev=0.01),
+      shape=[first_filter_height, first_filter_width, 1, first_filter_count])
+  first_bias = tf.compat.v1.get_variable(
+      name='first_bias',
+      initializer=tf.compat.v1.zeros_initializer,
+      shape=[first_filter_count])
+  first_conv_stride_x = 2
+  first_conv_stride_y = 2
+
+  first_conv = tf.nn.conv2d(
+      input=fingerprint_4d, filters=first_weights,
+      strides=[1, first_conv_stride_y, first_conv_stride_x, 1],
+      padding='SAME') + first_bias
+  first_relu = tf.nn.relu(first_conv)
+  if is_training:
+    first_dropout = tf.nn.dropout(first_relu, rate=dropout_rate)
+
+  else:
+    first_dropout = first_relu
+
+  second_filter_width = 8
+  second_filter_height = 10
+  second_filter_count = 8
+  second_weights = tf.compat.v1.get_variable(
+      name='second_weights',
+      initializer=tf.compat.v1.truncated_normal_initializer(stddev=0.01),
+      shape=[
+          second_filter_height, second_filter_width, first_filter_count,
+          second_filter_count
+      ])
+  second_bias = tf.compat.v1.get_variable(
+      name='second_bias',
+      initializer=tf.compat.v1.zeros_initializer,
+      shape=[second_filter_count])
+  second_conv_stride_x = 8
+  second_conv_stride_y = 8
+  second_conv = tf.nn.conv2d(
+      input=first_dropout, filters=second_weights,
+      strides=[1, second_conv_stride_y, second_conv_stride_x, 1],
+      padding='SAME') + second_bias
+  second_relu = tf.nn.relu(second_conv)
+  if is_training:
+    second_dropout = tf.nn.dropout(second_relu, rate=dropout_rate)
+  else:
+    second_dropout = second_relu
+
+  second_dropout_shape = second_dropout.get_shape()
+  second_dropout_output_width = second_dropout_shape[2]
+  second_dropout_output_height = second_dropout_shape[1]
+  second_dropout_element_count = int(second_dropout_output_width *
+                                     second_dropout_output_height *
+                                     second_filter_count)
+  flattened_second_dropout = tf.reshape(second_dropout,
+                                        [-1, second_dropout_element_count])
+  label_count = model_settings['label_count']
+  final_fc_weights = tf.compat.v1.get_variable(
+      name='final_fc_weights',
+      initializer=tf.compat.v1.truncated_normal_initializer(stddev=0.01),
+      shape=[second_dropout_element_count, label_count])
+  final_fc_bias = tf.compat.v1.get_variable(
+      name='final_fc_bias',
+      initializer=tf.compat.v1.zeros_initializer,
+      shape=[label_count])
+  final_fc = (
+      tf.matmul(flattened_second_dropout, final_fc_weights) + final_fc_bias)
+  if is_training:
+    return final_fc, dropout_rate
+  else:
+    return final_fc

tensorflow/examples/speech_commands/models_test.py

@@ -0,0 +1,116 @@
+# 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.
+# ==============================================================================
+"""Tests for speech commands models."""
+
+import tensorflow as tf
+
+from tensorflow.examples.speech_commands import models
+from tensorflow.python.framework import test_util
+from tensorflow.python.platform import test
+
+
+class ModelsTest(test.TestCase):
+
+  def _modelSettings(self):
+    return models.prepare_model_settings(
+        label_count=10,
+        sample_rate=16000,
+        clip_duration_ms=1000,
+        window_size_ms=20,
+        window_stride_ms=10,
+        feature_bin_count=40,
+        preprocess="mfcc")
+
+  def testPrepareModelSettings(self):
+    self.assertIsNotNone(
+        models.prepare_model_settings(
+            label_count=10,
+            sample_rate=16000,
+            clip_duration_ms=1000,
+            window_size_ms=20,
+            window_stride_ms=10,
+            feature_bin_count=40,
+            preprocess="mfcc"))
+
+  @test_util.run_deprecated_v1
+  def testCreateModelConvTraining(self):
+    model_settings = self._modelSettings()
+    with self.cached_session() as sess:
+      fingerprint_input = tf.zeros([1, model_settings["fingerprint_size"]])
+      logits, dropout_rate = models.create_model(
+          fingerprint_input, model_settings, "conv", True)
+      self.assertIsNotNone(logits)
+      self.assertIsNotNone(dropout_rate)
+      self.assertIsNotNone(sess.graph.get_tensor_by_name(logits.name))
+      self.assertIsNotNone(sess.graph.get_tensor_by_name(dropout_rate.name))
+
+  @test_util.run_deprecated_v1
+  def testCreateModelConvInference(self):
+    model_settings = self._modelSettings()
+    with self.cached_session() as sess:
+      fingerprint_input = tf.zeros([1, model_settings["fingerprint_size"]])
+      logits = models.create_model(fingerprint_input, model_settings, "conv",
+                                   False)
+      self.assertIsNotNone(logits)
+      self.assertIsNotNone(sess.graph.get_tensor_by_name(logits.name))
+
+  @test_util.run_deprecated_v1
+  def testCreateModelLowLatencyConvTraining(self):
+    model_settings = self._modelSettings()
+    with self.cached_session() as sess:
+      fingerprint_input = tf.zeros([1, model_settings["fingerprint_size"]])
+      logits, dropout_rate = models.create_model(
+          fingerprint_input, model_settings, "low_latency_conv", True)
+      self.assertIsNotNone(logits)
+      self.assertIsNotNone(dropout_rate)
+      self.assertIsNotNone(sess.graph.get_tensor_by_name(logits.name))
+      self.assertIsNotNone(sess.graph.get_tensor_by_name(dropout_rate.name))
+
+  @test_util.run_deprecated_v1
+  def testCreateModelFullyConnectedTraining(self):
+    model_settings = self._modelSettings()
+    with self.cached_session() as sess:
+      fingerprint_input = tf.zeros([1, model_settings["fingerprint_size"]])
+      logits, dropout_rate = models.create_model(
+          fingerprint_input, model_settings, "single_fc", True)
+      self.assertIsNotNone(logits)
+      self.assertIsNotNone(dropout_rate)
+      self.assertIsNotNone(sess.graph.get_tensor_by_name(logits.name))
+      self.assertIsNotNone(sess.graph.get_tensor_by_name(dropout_rate.name))
+
+  def testCreateModelBadArchitecture(self):
+    model_settings = self._modelSettings()
+    with self.cached_session():
+      fingerprint_input = tf.zeros([1, model_settings["fingerprint_size"]])
+      with self.assertRaises(Exception) as e:
+        models.create_model(fingerprint_input, model_settings,
+                            "bad_architecture", True)
+      self.assertIn("not recognized", str(e.exception))
+
+  @test_util.run_deprecated_v1
+  def testCreateModelTinyConvTraining(self):
+    model_settings = self._modelSettings()
+    with self.cached_session() as sess:
+      fingerprint_input = tf.zeros([1, model_settings["fingerprint_size"]])
+      logits, dropout_rate = models.create_model(
+          fingerprint_input, model_settings, "tiny_conv", True)
+      self.assertIsNotNone(logits)
+      self.assertIsNotNone(dropout_rate)
+      self.assertIsNotNone(sess.graph.get_tensor_by_name(logits.name))
+      self.assertIsNotNone(sess.graph.get_tensor_by_name(dropout_rate.name))
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/examples/speech_commands/recognize_commands.cc

@@ -0,0 +1,138 @@
+/* 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/examples/speech_commands/recognize_commands.h"
+
+#include <algorithm>
+#include <cstdint>
+#include <limits>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/types.h"
+
+namespace tensorflow {
+
+RecognizeCommands::RecognizeCommands(const std::vector<std::string>& labels,
+                                     int32_t average_window_duration_ms,
+                                     float detection_threshold,
+                                     int32_t suppression_ms,
+                                     int32_t minimum_count)
+    : labels_(labels),
+      average_window_duration_ms_(average_window_duration_ms),
+      detection_threshold_(detection_threshold),
+      suppression_ms_(suppression_ms),
+      minimum_count_(minimum_count) {
+  labels_count_ = labels.size();
+  previous_top_label_ = "_silence_";
+  previous_top_label_time_ = std::numeric_limits<int64_t>::min();
+}
+
+absl::Status RecognizeCommands::ProcessLatestResults(
+    const Tensor& latest_results, const int64_t current_time_ms,
+    std::string* found_command, float* score, bool* is_new_command) {
+  if (latest_results.NumElements() != labels_count_) {
+    return absl::InvalidArgumentError(absl::StrCat(
+        "The results for recognition should contain ", labels_count_,
+        " elements, but there are ", latest_results.NumElements()));
+  }
+
+  if ((!previous_results_.empty()) &&
+      (current_time_ms < previous_results_.front().first)) {
+    return absl::InvalidArgumentError(absl::StrCat(
+        "Results must be fed in increasing time order, but received a "
+        "timestamp of ",
+        current_time_ms, " that was earlier than the previous one of ",
+        previous_results_.front().first));
+  }
+
+  // Add the latest results to the head of the queue.
+  previous_results_.push_back({current_time_ms, latest_results});
+
+  // Prune any earlier results that are too old for the averaging window.
+  const int64_t time_limit = current_time_ms - average_window_duration_ms_;
+  while (previous_results_.front().first < time_limit) {
+    previous_results_.pop_front();
+  }
+
+  // If there are too few results, assume the result will be unreliable and
+  // bail.
+  const int64_t how_many_results = previous_results_.size();
+  const int64_t earliest_time = previous_results_.front().first;
+  const int64_t samples_duration = current_time_ms - earliest_time;
+  if ((how_many_results < minimum_count_) ||
+      (samples_duration < (average_window_duration_ms_ / 4))) {
+    *found_command = previous_top_label_;
+    *score = 0.0f;
+    *is_new_command = false;
+    return absl::OkStatus();
+  }
+
+  // Calculate the average score across all the results in the window.
+  std::vector<float> average_scores(labels_count_);
+  for (const auto& previous_result : previous_results_) {
+    const Tensor& scores_tensor = previous_result.second;
+    auto scores_flat = scores_tensor.flat<float>();
+    for (int i = 0; i < scores_flat.size(); ++i) {
+      average_scores[i] += scores_flat(i) / how_many_results;
+    }
+  }
+
+  // Sort the averaged results in descending score order.
+  std::vector<std::pair<int, float>> sorted_average_scores;
+  sorted_average_scores.reserve(labels_count_);
+  for (int i = 0; i < labels_count_; ++i) {
+    sorted_average_scores.push_back(
+        std::pair<int, float>({i, average_scores[i]}));
+  }
+  std::sort(sorted_average_scores.begin(), sorted_average_scores.end(),
+            [](const std::pair<int, float>& left,
+               const std::pair<int, float>& right) {
+              return left.second > right.second;
+            });
+
+  // See if the latest top score is enough to trigger a detection.
+  const int current_top_index = sorted_average_scores[0].first;
+  const std::string current_top_label = labels_[current_top_index];
+  const float current_top_score = sorted_average_scores[0].second;
+  // If we've recently had another label trigger, assume one that occurs too
+  // soon afterwards is a bad result.
+  int64_t time_since_last_top;
+  if ((previous_top_label_ == "_silence_") ||
+      (previous_top_label_time_ == std::numeric_limits<int64_t>::min())) {
+    time_since_last_top = std::numeric_limits<int64_t>::max();
+  } else {
+    time_since_last_top = current_time_ms - previous_top_label_time_;
+  }
+  if ((current_top_score > detection_threshold_) &&
+      (current_top_label != previous_top_label_) &&
+      (time_since_last_top > suppression_ms_)) {
+    previous_top_label_ = current_top_label;
+    previous_top_label_time_ = current_time_ms;
+    *is_new_command = true;
+  } else {
+    *is_new_command = false;
+  }
+  *found_command = current_top_label;
+  *score = current_top_score;
+
+  return absl::OkStatus();
+}
+
+}  // namespace tensorflow

tensorflow/examples/speech_commands/recognize_commands.h

@@ -0,0 +1,83 @@
+/* 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_EXAMPLES_SPEECH_COMMANDS_RECOGNIZE_COMMANDS_H_
+#define TENSORFLOW_EXAMPLES_SPEECH_COMMANDS_RECOGNIZE_COMMANDS_H_
+
+#include <cstdint>
+#include <deque>
+#include <unordered_set>
+#include <utility>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/platform/types.h"
+
+namespace tensorflow {
+
+// This class is designed to apply a very primitive decoding model on top of the
+// instantaneous results from running an audio recognition model on a single
+// window of samples. It applies smoothing over time so that noisy individual
+// label scores are averaged, increasing the confidence that apparent matches
+// are real.
+// To use it, you should create a class object with the configuration you
+// want, and then feed results from running a TensorFlow model into the
+// processing method. The timestamp for each subsequent call should be
+// increasing from the previous, since the class is designed to process a stream
+// of data over time.
+class RecognizeCommands {
+ public:
+  // labels should be a list of the strings associated with each one-hot score.
+  // The window duration controls the smoothing. Longer durations will give a
+  // higher confidence that the results are correct, but may miss some commands.
+  // The detection threshold has a similar effect, with high values increasing
+  // the precision at the cost of recall. The minimum count controls how many
+  // results need to be in the averaging window before it's seen as a reliable
+  // average. This prevents erroneous results when the averaging window is
+  // initially being populated for example. The suppression argument disables
+  // further recognitions for a set time after one has been triggered, which can
+  // help reduce spurious recognitions.
+  explicit RecognizeCommands(const std::vector<std::string>& labels,
+                             int32_t average_window_duration_ms = 1000,
+                             float detection_threshold = 0.2,
+                             int32_t suppression_ms = 500,
+                             int32_t minimum_count = 3);
+
+  // Call this with the results of running a model on sample data.
+  absl::Status ProcessLatestResults(const Tensor& latest_results,
+                                    const int64_t current_time_ms,
+                                    std::string* found_command, float* score,
+                                    bool* is_new_command);
+
+ private:
+  // Configuration
+  std::vector<std::string> labels_;
+  int32_t average_window_duration_ms_;
+  float detection_threshold_;
+  int32_t suppression_ms_;
+  int32_t minimum_count_;
+
+  // Working variables
+  std::deque<std::pair<int64_t, Tensor>> previous_results_;
+  std::string previous_top_label_;
+  int64_t labels_count_;
+  int64_t previous_top_label_time_;
+};
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_EXAMPLES_SPEECH_COMMANDS_RECOGNIZE_COMMANDS_H_

tensorflow/examples/speech_commands/recognize_commands.py

@@ -0,0 +1,197 @@
+# 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.
+# ==============================================================================
+"""Stream accuracy recognize commands."""
+import collections
+
+import numpy as np
+
+
+class RecognizeResult(object):
+  """Save recognition result temporarily.
+
+  Attributes:
+    founded_command: A string indicating the word just founded. Default value
+      is '_silence_'
+    score: A float representing the confidence of founded word. Default
+      value is zero.
+    is_new_command: A boolean indicating if the founded command is a new one
+      against the last one. Default value is False.
+  """
+
+  def __init__(self,
+               founded_command="_silence_",
+               score=0.0,
+               is_new_command=False):
+    """Construct a recognition result.
+
+    Args:
+      founded_command: A string indicating the word just founded.
+      score: A float representing the confidence of founded word.
+      is_new_command: A boolean indicating if the founded command is a new one
+        against the last one.
+    """
+    self._founded_command = founded_command
+    self._score = score
+    self._is_new_command = is_new_command
+
+  @property
+  def founded_command(self):
+    return self._founded_command
+
+  @founded_command.setter
+  def founded_command(self, value):
+    self._founded_command = value
+
+  @property
+  def score(self):
+    return self._score
+
+  @score.setter
+  def score(self, value):
+    self._score = value
+
+  @property
+  def is_new_command(self):
+    return self._is_new_command
+
+  @is_new_command.setter
+  def is_new_command(self, value):
+    self._is_new_command = value
+
+
+class RecognizeCommands(object):
+  """Smooth the inference results by using average window.
+
+  Maintain a slide window over the audio stream, which adds new result(a pair of
+  the 1.confidences of all classes and 2.the start timestamp of input audio
+  clip) directly the inference produces one and removes the most previous one
+  and other abnormal values. Then it smooth the results in the window to get
+  the most reliable command in this period.
+
+  Attributes:
+    _label: A list containing commands at corresponding lines.
+    _average_window_duration: The length of average window.
+    _detection_threshold: A confidence threshold for filtering out unreliable
+      command.
+    _suppression_ms: Milliseconds every two reliable founded commands should
+      apart.
+    _minimum_count: An integer count indicating the minimum results the average
+      window should cover.
+    _previous_results: A deque to store previous results.
+    _label_count: The length of label list.
+    _previous_top_label: Last founded command. Initial value is '_silence_'.
+    _previous_top_time: The timestamp of _previous results. Default is -np.inf.
+  """
+
+  def __init__(self, labels, average_window_duration_ms, detection_threshold,
+               suppression_ms, minimum_count):
+    """Init the RecognizeCommands with parameters used for smoothing."""
+    # Configuration
+    self._labels = labels
+    self._average_window_duration_ms = average_window_duration_ms
+    self._detection_threshold = detection_threshold
+    self._suppression_ms = suppression_ms
+    self._minimum_count = minimum_count
+    # Working Variable
+    self._previous_results = collections.deque()
+    self._label_count = len(labels)
+    self._previous_top_label = "_silence_"
+    self._previous_top_time = -np.inf
+
+  def process_latest_result(self, latest_results, current_time_ms,
+                            recognize_element):
+    """Smoothing the results in average window when a new result is added in.
+
+    Receive a new result from inference and put the founded command into
+    a RecognizeResult instance after the smoothing procedure.
+
+    Args:
+      latest_results: A list containing the confidences of all labels.
+      current_time_ms: The start timestamp of the input audio clip.
+      recognize_element: An instance of RecognizeResult to store founded
+        command, its scores and if it is a new command.
+
+    Raises:
+      ValueError: The length of this result from inference doesn't match
+        label count.
+      ValueError: The timestamp of this result is earlier than the most
+        previous one in the average window
+    """
+    if latest_results.shape[0] != self._label_count:
+      raise ValueError("The results for recognition should contain {} "
+                       "elements, but there are {} produced".format(
+                           self._label_count, latest_results.shape[0]))
+    if (
+        self._previous_results
+        and current_time_ms < self._previous_results[0][0]
+    ):
+      raise ValueError("Results must be fed in increasing time order, "
+                       "but receive a timestamp of {}, which was earlier "
+                       "than the previous one of {}".format(
+                           current_time_ms, self._previous_results[0][0]))
+
+    # Add the latest result to the head of the deque.
+    self._previous_results.append([current_time_ms, latest_results])
+
+    # Prune any earlier results that are too old for the averaging window.
+    time_limit = current_time_ms - self._average_window_duration_ms
+    while time_limit > self._previous_results[0][0]:
+      self._previous_results.popleft()
+
+    # If there are too few results, the result will be unreliable and bail.
+    how_many_results = len(self._previous_results)
+    earliest_time = self._previous_results[0][0]
+    sample_duration = current_time_ms - earliest_time
+    if (how_many_results < self._minimum_count or
+        sample_duration < self._average_window_duration_ms / 4):
+      recognize_element.founded_command = self._previous_top_label
+      recognize_element.score = 0.0
+      recognize_element.is_new_command = False
+      return
+
+    # Calculate the average score across all the results in the window.
+    average_scores = np.zeros(self._label_count)
+    for item in self._previous_results:
+      score = item[1]
+      for i in range(score.size):
+        average_scores[i] += score[i] / how_many_results
+
+    # Sort the averaged results in descending score order.
+    sorted_averaged_index_score = []
+    for i in range(self._label_count):
+      sorted_averaged_index_score.append([i, average_scores[i]])
+    sorted_averaged_index_score = sorted(
+        sorted_averaged_index_score, key=lambda p: p[1], reverse=True)
+
+    # Use the information of previous result to get current result
+    current_top_index = sorted_averaged_index_score[0][0]
+    current_top_label = self._labels[current_top_index]
+    current_top_score = sorted_averaged_index_score[0][1]
+    time_since_last_top = 0
+    if (self._previous_top_label == "_silence_" or
+        self._previous_top_time == -np.inf):
+      time_since_last_top = np.inf
+    else:
+      time_since_last_top = current_time_ms - self._previous_top_time
+    if (current_top_score > self._detection_threshold and
+        current_top_label != self._previous_top_label and
+        time_since_last_top > self._suppression_ms):
+      self._previous_top_label = current_top_label
+      self._previous_top_time = current_time_ms
+      recognize_element.is_new_command = True
+    else:
+      recognize_element.is_new_command = False
+    recognize_element.founded_command = current_top_label
+    recognize_element.score = current_top_score

tensorflow/examples/speech_commands/recognize_commands_test.cc

@@ -0,0 +1,119 @@
+/* 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/examples/speech_commands/recognize_commands.h"
+
+#include <cstdint>
+#include <string>
+
+#include "xla/tsl/lib/core/status_test_util.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_testutil.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/platform/types.h"
+
+namespace tensorflow {
+
+TEST(RecognizeCommandsTest, Basic) {
+  RecognizeCommands recognize_commands({"_silence_", "a", "b"});
+
+  Tensor results(DT_FLOAT, {3});
+  test::FillValues<float>(&results, {1.0f, 0.0f, 0.0f});
+
+  std::string found_command;
+  float score;
+  bool is_new_command;
+  TF_EXPECT_OK(recognize_commands.ProcessLatestResults(
+      results, 0, &found_command, &score, &is_new_command));
+}
+
+TEST(RecognizeCommandsTest, FindCommands) {
+  RecognizeCommands recognize_commands({"_silence_", "a", "b"}, 1000, 0.2f);
+
+  Tensor results(DT_FLOAT, {3});
+
+  test::FillValues<float>(&results, {0.0f, 1.0f, 0.0f});
+  bool has_found_new_command = false;
+  std::string new_command;
+  for (int i = 0; i < 10; ++i) {
+    std::string found_command;
+    float score;
+    bool is_new_command;
+    int64_t current_time_ms = 0 + (i * 100);
+    TF_EXPECT_OK(recognize_commands.ProcessLatestResults(
+        results, current_time_ms, &found_command, &score, &is_new_command));
+    if (is_new_command) {
+      EXPECT_FALSE(has_found_new_command);
+      has_found_new_command = true;
+      new_command = found_command;
+    }
+  }
+  EXPECT_TRUE(has_found_new_command);
+  EXPECT_EQ("a", new_command);
+
+  test::FillValues<float>(&results, {0.0f, 0.0f, 1.0f});
+  has_found_new_command = false;
+  new_command = "";
+  for (int i = 0; i < 10; ++i) {
+    std::string found_command;
+    float score;
+    bool is_new_command;
+    int64_t current_time_ms = 1000 + (i * 100);
+    TF_EXPECT_OK(recognize_commands.ProcessLatestResults(
+        results, current_time_ms, &found_command, &score, &is_new_command));
+    if (is_new_command) {
+      EXPECT_FALSE(has_found_new_command);
+      has_found_new_command = true;
+      new_command = found_command;
+    }
+  }
+  EXPECT_TRUE(has_found_new_command);
+  EXPECT_EQ("b", new_command);
+}
+
+TEST(RecognizeCommandsTest, BadInputLength) {
+  RecognizeCommands recognize_commands({"_silence_", "a", "b"}, 1000, 0.2f);
+
+  Tensor bad_results(DT_FLOAT, {2});
+  test::FillValues<float>(&bad_results, {1.0f, 0.0f});
+
+  std::string found_command;
+  float score;
+  bool is_new_command;
+  EXPECT_FALSE(recognize_commands
+                   .ProcessLatestResults(bad_results, 0, &found_command, &score,
+                                         &is_new_command)
+                   .ok());
+}
+
+TEST(RecognizeCommandsTest, BadInputTimes) {
+  RecognizeCommands recognize_commands({"_silence_", "a", "b"}, 1000, 0.2f);
+
+  Tensor results(DT_FLOAT, {3});
+  test::FillValues<float>(&results, {1.0f, 0.0f, 0.0f});
+
+  std::string found_command;
+  float score;
+  bool is_new_command;
+  TF_EXPECT_OK(recognize_commands.ProcessLatestResults(
+      results, 100, &found_command, &score, &is_new_command));
+  EXPECT_FALSE(recognize_commands
+                   .ProcessLatestResults(results, 0, &found_command, &score,
+                                         &is_new_command)
+                   .ok());
+}
+
+}  // namespace tensorflow

tensorflow/examples/speech_commands/test_streaming_accuracy.cc

@@ -0,0 +1,323 @@
+/* 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.
+==============================================================================*/
+/*
+
+Tool to create accuracy statistics from running an audio recognition model on a
+continuous stream of samples.
+
+This is designed to be an environment for running experiments on new models and
+settings to understand the effects they will have in a real application. You
+need to supply it with a long audio file containing sounds you want to recognize
+and a text file listing the labels of each sound along with the time they occur.
+With this information, and a frozen model, the tool will process the audio
+stream, apply the model, and keep track of how many mistakes and successes the
+model achieved.
+
+The matched percentage is the number of sounds that were correctly classified,
+as a percentage of the total number of sounds listed in the ground truth file.
+A correct classification is when the right label is chosen within a short time
+of the expected ground truth, where the time tolerance is controlled by the
+'time_tolerance_ms' command line flag.
+
+The wrong percentage is how many sounds triggered a detection (the classifier
+figured out it wasn't silence or background noise), but the detected class was
+wrong. This is also a percentage of the total number of ground truth sounds.
+
+The false positive percentage is how many sounds were detected when there was
+only silence or background noise. This is also expressed as a percentage of the
+total number of ground truth sounds, though since it can be large it may go
+above 100%.
+
+The easiest way to get an audio file and labels to test with is by using the
+'generate_streaming_test_wav' script. This will synthesize a test file with
+randomly placed sounds and background noise, and output a text file with the
+ground truth.
+
+If you want to test natural data, you need to use a .wav with the same sample
+rate as your model (often 16,000 samples per second), and note down where the
+sounds occur in time. Save this information out as a comma-separated text file,
+where the first column is the label and the second is the time in seconds from
+the start of the file that it occurs.
+
+Here's an example of how to run the tool:
+
+bazel run tensorflow/examples/speech_commands:test_streaming_accuracy -- \
+--wav=/tmp/streaming_test_bg.wav \
+--graph=/tmp/conv_frozen.pb \
+--labels=/tmp/speech_commands_train/conv_labels.txt \
+--ground_truth=/tmp/streaming_test_labels.txt --verbose \
+--clip_duration_ms=1000 --detection_threshold=0.70 --average_window_ms=500 \
+--suppression_ms=500 --time_tolerance_ms=1500
+
+ */
+
+#include <algorithm>
+#include <cstdint>
+#include <fstream>
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "xla/tsl/platform/env.h"
+#include "xla/tsl/platform/types.h"
+#include "xla/tsl/util/command_line_flags.h"
+#include "tensorflow/core/framework/graph.pb.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/wav/wav_io.h"
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/init_main.h"
+#include "tensorflow/core/platform/logging.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/public/session.h"
+#include "tensorflow/core/public/session_options.h"
+#include "tensorflow/core/util/command_line_flags.h"
+#include "tensorflow/examples/speech_commands/accuracy_utils.h"
+#include "tensorflow/examples/speech_commands/recognize_commands.h"
+
+// These are all common classes it's handy to reference with no namespace.
+using ::int64_t;
+using tensorflow::Flag;
+using tensorflow::int32;
+using tensorflow::Status;
+using tensorflow::string;
+using tensorflow::Tensor;
+using tensorflow::uint16;
+using tensorflow::uint32;
+
+namespace {
+
+// Reads a model graph definition from disk, and creates a session object you
+// can use to run it.
+Status LoadGraph(const string& graph_file_name,
+                 std::unique_ptr<tensorflow::Session>* session) {
+  tensorflow::GraphDef graph_def;
+  Status load_graph_status =
+      ReadBinaryProto(tensorflow::Env::Default(), graph_file_name, &graph_def);
+  if (!load_graph_status.ok()) {
+    return absl::NotFoundError(absl::StrCat("Failed to load compute graph at '",
+                                            graph_file_name, "'"));
+  }
+  session->reset(tensorflow::NewSession(tensorflow::SessionOptions()));
+  Status session_create_status = (*session)->Create(graph_def);
+  if (!session_create_status.ok()) {
+    return session_create_status;
+  }
+  return absl::OkStatus();
+}
+
+// Takes a file name, and loads a list of labels from it, one per line, and
+// returns a vector of the strings.
+Status ReadLabelsFile(const string& file_name, std::vector<string>* result) {
+  std::ifstream file(file_name);
+  if (!file) {
+    return absl::NotFoundError(
+        absl::StrCat("Labels file '", file_name, "' not found."));
+  }
+  result->clear();
+  string line;
+  while (std::getline(file, line)) {
+    result->push_back(line);
+  }
+  return absl::OkStatus();
+}
+
+}  // namespace
+
+int main(int argc, char* argv[]) {
+  string wav = "";
+  string graph = "";
+  string labels = "";
+  string ground_truth = "";
+  string input_data_name = "decoded_sample_data:0";
+  string input_rate_name = "decoded_sample_data:1";
+  string output_name = "labels_softmax";
+  int32_t clip_duration_ms = 1000;
+  int32_t clip_stride_ms = 30;
+  int32_t average_window_ms = 500;
+  int32_t time_tolerance_ms = 750;
+  int32_t suppression_ms = 1500;
+  float detection_threshold = 0.7f;
+  bool verbose = false;
+  std::vector<Flag> flag_list = {
+      Flag("wav", &wav, "audio file to be identified"),
+      Flag("graph", &graph, "model to be executed"),
+      Flag("labels", &labels, "path to file containing labels"),
+      Flag("ground_truth", &ground_truth,
+           "path to file containing correct times and labels of words in the "
+           "audio as <word>,<timestamp in ms> lines"),
+      Flag("input_data_name", &input_data_name,
+           "name of input data node in model"),
+      Flag("input_rate_name", &input_rate_name,
+           "name of input sample rate node in model"),
+      Flag("output_name", &output_name, "name of output node in model"),
+      Flag("clip_duration_ms", &clip_duration_ms,
+           "length of recognition window"),
+      Flag("average_window_ms", &average_window_ms,
+           "length of window to smooth results over"),
+      Flag("time_tolerance_ms", &time_tolerance_ms,
+           "maximum gap allowed between a recognition and ground truth"),
+      Flag("suppression_ms", &suppression_ms,
+           "how long to ignore others for after a recognition"),
+      Flag("clip_stride_ms", &clip_stride_ms, "how often to run recognition"),
+      Flag("detection_threshold", &detection_threshold,
+           "what score is required to trigger detection of a word"),
+      Flag("verbose", &verbose, "whether to log extra debugging information"),
+  };
+  string usage = tensorflow::Flags::Usage(argv[0], flag_list);
+  const bool parse_result = tensorflow::Flags::Parse(&argc, argv, flag_list);
+  if (!parse_result) {
+    LOG(ERROR) << usage;
+    return -1;
+  }
+
+  // We need to call this to set up global state for TensorFlow.
+  tensorflow::port::InitMain(argv[0], &argc, &argv);
+  if (argc > 1) {
+    LOG(ERROR) << "Unknown argument " << argv[1] << "\n" << usage;
+    return -1;
+  }
+
+  // First we load and initialize the model.
+  std::unique_ptr<tensorflow::Session> session;
+  Status load_graph_status = LoadGraph(graph, &session);
+  if (!load_graph_status.ok()) {
+    LOG(ERROR) << load_graph_status;
+    return -1;
+  }
+
+  std::vector<string> labels_list;
+  Status read_labels_status = ReadLabelsFile(labels, &labels_list);
+  if (!read_labels_status.ok()) {
+    LOG(ERROR) << read_labels_status;
+    return -1;
+  }
+
+  std::vector<std::pair<string, int64_t>> ground_truth_list;
+  Status read_ground_truth_status =
+      tensorflow::ReadGroundTruthFile(ground_truth, &ground_truth_list);
+  if (!read_ground_truth_status.ok()) {
+    LOG(ERROR) << read_ground_truth_status;
+    return -1;
+  }
+
+  string wav_string;
+  Status read_wav_status = tensorflow::ReadFileToString(
+      tensorflow::Env::Default(), wav, &wav_string);
+  if (!read_wav_status.ok()) {
+    LOG(ERROR) << read_wav_status;
+    return -1;
+  }
+  std::vector<float> audio_data;
+  uint32 sample_count;
+  uint16 channel_count;
+  uint32 sample_rate;
+  Status decode_wav_status = tensorflow::wav::DecodeLin16WaveAsFloatVector(
+      wav_string, &audio_data, &sample_count, &channel_count, &sample_rate);
+  if (!decode_wav_status.ok()) {
+    LOG(ERROR) << decode_wav_status;
+    return -1;
+  }
+  if (channel_count != 1) {
+    LOG(ERROR) << "Only mono .wav files can be used, but input has "
+               << channel_count << " channels.";
+    return -1;
+  }
+
+  const int64_t clip_duration_samples = (clip_duration_ms * sample_rate) / 1000;
+  const int64_t clip_stride_samples = (clip_stride_ms * sample_rate) / 1000;
+  Tensor audio_data_tensor(tensorflow::DT_FLOAT,
+                           tensorflow::TensorShape({clip_duration_samples, 1}));
+
+  Tensor sample_rate_tensor(tensorflow::DT_INT32, tensorflow::TensorShape({}));
+  sample_rate_tensor.scalar<int32>()() = sample_rate;
+
+  tensorflow::RecognizeCommands recognize_commands(
+      labels_list, average_window_ms, detection_threshold, suppression_ms);
+
+  std::vector<std::pair<string, int64_t>> all_found_words;
+  tensorflow::StreamingAccuracyStats previous_stats;
+
+  const int64_t audio_data_end = (sample_count - clip_duration_samples);
+  for (int64_t audio_data_offset = 0; audio_data_offset < audio_data_end;
+       audio_data_offset += clip_stride_samples) {
+    const float* input_start = &(audio_data[audio_data_offset]);
+    const float* input_end = input_start + clip_duration_samples;
+    std::copy(input_start, input_end, audio_data_tensor.flat<float>().data());
+
+    // Actually run the audio through the model.
+    std::vector<Tensor> outputs;
+    Status run_status = session->Run({{input_data_name, audio_data_tensor},
+                                      {input_rate_name, sample_rate_tensor}},
+                                     {output_name}, {}, &outputs);
+    if (!run_status.ok()) {
+      LOG(ERROR) << "Running model failed: " << run_status;
+      return -1;
+    }
+
+    const int64_t current_time_ms = (audio_data_offset * 1000) / sample_rate;
+    string found_command;
+    float score;
+    bool is_new_command;
+    Status recognize_status = recognize_commands.ProcessLatestResults(
+        outputs[0], current_time_ms, &found_command, &score, &is_new_command);
+    if (!recognize_status.ok()) {
+      LOG(ERROR) << "Recognition processing failed: " << recognize_status;
+      return -1;
+    }
+
+    if (is_new_command && (found_command != "_silence_")) {
+      all_found_words.push_back({found_command, current_time_ms});
+      if (verbose) {
+        tensorflow::StreamingAccuracyStats stats;
+        tensorflow::CalculateAccuracyStats(ground_truth_list, all_found_words,
+                                           current_time_ms, time_tolerance_ms,
+                                           &stats);
+        int32_t false_positive_delta = stats.how_many_false_positives -
+                                       previous_stats.how_many_false_positives;
+        int32_t correct_delta = stats.how_many_correct_words -
+                                previous_stats.how_many_correct_words;
+        int32_t wrong_delta =
+            stats.how_many_wrong_words - previous_stats.how_many_wrong_words;
+        string recognition_state;
+        if (false_positive_delta == 1) {
+          recognition_state = " (False Positive)";
+        } else if (correct_delta == 1) {
+          recognition_state = " (Correct)";
+        } else if (wrong_delta == 1) {
+          recognition_state = " (Wrong)";
+        } else {
+          LOG(ERROR) << "Unexpected state in statistics";
+        }
+        LOG(INFO) << current_time_ms << "ms: " << found_command << ": " << score
+                  << recognition_state;
+        previous_stats = stats;
+        tensorflow::PrintAccuracyStats(stats);
+      }
+    }
+  }
+
+  tensorflow::StreamingAccuracyStats stats;
+  tensorflow::CalculateAccuracyStats(ground_truth_list, all_found_words, -1,
+                                     time_tolerance_ms, &stats);
+  tensorflow::PrintAccuracyStats(stats);
+
+  return 0;
+}

tensorflow/examples/speech_commands/test_streaming_accuracy.py

@@ -0,0 +1,245 @@
+# 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.
+# ==============================================================================
+r"""Tool to create accuracy statistics on a continuous stream of samples.
+
+This is designed to be an environment for running experiments on new models and
+settings to understand the effects they will have in a real application. You
+need to supply it with a long audio file containing sounds you want to recognize
+and a text file listing the labels of each sound along with the time they occur.
+With this information, and a frozen model, the tool will process the audio
+stream, apply the model, and keep track of how many mistakes and successes the
+model achieved.
+
+The matched percentage is the number of sounds that were correctly classified,
+as a percentage of the total number of sounds listed in the ground truth file.
+A correct classification is when the right label is chosen within a short time
+of the expected ground truth, where the time tolerance is controlled by the
+'time_tolerance_ms' command line flag.
+
+The wrong percentage is how many sounds triggered a detection (the classifier
+figured out it wasn't silence or background noise), but the detected class was
+wrong. This is also a percentage of the total number of ground truth sounds.
+
+The false positive percentage is how many sounds were detected when there was
+only silence or background noise. This is also expressed as a percentage of the
+total number of ground truth sounds, though since it can be large it may go
+above 100%.
+
+The easiest way to get an audio file and labels to test with is by using the
+'generate_streaming_test_wav' script. This will synthesize a test file with
+randomly placed sounds and background noise, and output a text file with the
+ground truth.
+
+If you want to test natural data, you need to use a .wav with the same sample
+rate as your model (often 16,000 samples per second), and note down where the
+sounds occur in time. Save this information out as a comma-separated text file,
+where the first column is the label and the second is the time in seconds from
+the start of the file that it occurs.
+
+Here's an example of how to run the tool:
+
+bazel run tensorflow/examples/speech_commands:test_streaming_accuracy_py -- \
+--wav=/tmp/streaming_test_bg.wav \
+--ground-truth=/tmp/streaming_test_labels.txt --verbose \
+--model=/tmp/conv_frozen.pb \
+--labels=/tmp/speech_commands_train/conv_labels.txt \
+--clip_duration_ms=1000 --detection_threshold=0.70 --average_window_ms=500 \
+--suppression_ms=500 --time_tolerance_ms=1500
+"""
+
+import argparse
+import sys
+
+import numpy
+import tensorflow as tf
+
+from accuracy_utils import StreamingAccuracyStats
+from recognize_commands import RecognizeCommands
+from recognize_commands import RecognizeResult
+from tensorflow.python.ops import io_ops
+
+FLAGS = None
+
+
+def load_graph(mode_file):
+  """Read a tensorflow model, and creates a default graph object."""
+  graph = tf.Graph()
+  with graph.as_default():
+    od_graph_def = tf.compat.v1.GraphDef()
+    with tf.io.gfile.GFile(mode_file, 'rb') as fid:
+      serialized_graph = fid.read()
+      od_graph_def.ParseFromString(serialized_graph)
+      tf.import_graph_def(od_graph_def, name='')
+  return graph
+
+
+def read_label_file(file_name):
+  """Load a list of label."""
+  label_list = []
+  with open(file_name, 'r') as f:
+    for line in f:
+      label_list.append(line.strip())
+  return label_list
+
+
+def read_wav_file(filename):
+  """Load a wav file and return sample_rate and numpy data of float64 type."""
+  with tf.compat.v1.Session(graph=tf.Graph()) as sess:
+    wav_filename_placeholder = tf.compat.v1.placeholder(tf.string, [])
+    wav_loader = io_ops.read_file(wav_filename_placeholder)
+    wav_decoder = tf.audio.decode_wav(wav_loader, desired_channels=1)
+    res = sess.run(wav_decoder, feed_dict={wav_filename_placeholder: filename})
+  return res.sample_rate, res.audio.flatten()
+
+
+def main(_):
+  label_list = read_label_file(FLAGS.labels)
+  sample_rate, data = read_wav_file(FLAGS.wav)
+  # Init instance of RecognizeCommands with given parameters.
+  recognize_commands = RecognizeCommands(
+      labels=label_list,
+      average_window_duration_ms=FLAGS.average_window_duration_ms,
+      detection_threshold=FLAGS.detection_threshold,
+      suppression_ms=FLAGS.suppression_ms,
+      minimum_count=4)
+
+  # Init instance of StreamingAccuracyStats and load ground truth.
+  stats = StreamingAccuracyStats()
+  stats.read_ground_truth_file(FLAGS.ground_truth)
+  recognize_element = RecognizeResult()
+  all_found_words = []
+  data_samples = data.shape[0]
+  clip_duration_samples = int(FLAGS.clip_duration_ms * sample_rate / 1000)
+  clip_stride_samples = int(FLAGS.clip_stride_ms * sample_rate / 1000)
+  audio_data_end = data_samples - clip_duration_samples
+
+  # Load model and create a tf session to process audio pieces
+  recognize_graph = load_graph(FLAGS.model)
+  with recognize_graph.as_default():
+    with tf.compat.v1.Session() as sess:
+
+      # Get input and output tensor
+      data_tensor = sess.graph.get_tensor_by_name(FLAGS.input_names[0])
+      sample_rate_tensor = sess.graph.get_tensor_by_name(FLAGS.input_names[1])
+      output_softmax_tensor = sess.graph.get_tensor_by_name(FLAGS.output_name)
+
+      # Inference along audio stream.
+      for audio_data_offset in range(0, audio_data_end, clip_stride_samples):
+        input_start = audio_data_offset
+        input_end = audio_data_offset + clip_duration_samples
+        outputs = sess.run(
+            output_softmax_tensor,
+            feed_dict={
+                data_tensor:
+                    numpy.expand_dims(data[input_start:input_end], axis=-1),
+                sample_rate_tensor:
+                    sample_rate
+            })
+        outputs = numpy.squeeze(outputs)
+        current_time_ms = int(audio_data_offset * 1000 / sample_rate)
+        try:
+          recognize_commands.process_latest_result(outputs, current_time_ms,
+                                                   recognize_element)
+        except ValueError as e:
+          tf.compat.v1.logging.error('Recognition processing failed: {}' % e)
+          return
+        if (recognize_element.is_new_command and
+            recognize_element.founded_command != '_silence_'):
+          all_found_words.append(
+              [recognize_element.founded_command, current_time_ms])
+          if FLAGS.verbose:
+            stats.calculate_accuracy_stats(all_found_words, current_time_ms,
+                                           FLAGS.time_tolerance_ms)
+            try:
+              recognition_state = stats.delta()
+            except ValueError as e:
+              tf.compat.v1.logging.error(
+                  'Statistics delta computing failed: {}'.format(e))
+            else:
+              tf.compat.v1.logging.info('{}ms {}:{}{}'.format(
+                  current_time_ms, recognize_element.founded_command,
+                  recognize_element.score, recognition_state))
+              stats.print_accuracy_stats()
+  stats.calculate_accuracy_stats(all_found_words, -1, FLAGS.time_tolerance_ms)
+  stats.print_accuracy_stats()
+
+
+if __name__ == '__main__':
+  parser = argparse.ArgumentParser(description='test_streaming_accuracy')
+  parser.add_argument(
+      '--wav', type=str, default='', help='The wave file path to evaluate.')
+  parser.add_argument(
+      '--ground-truth',
+      type=str,
+      default='',
+      help='The ground truth file path corresponding to wav file.')
+  parser.add_argument(
+      '--labels',
+      type=str,
+      default='',
+      help='The label file path containing all possible classes.')
+  parser.add_argument(
+      '--model', type=str, default='', help='The model used for inference')
+  parser.add_argument(
+      '--input-names',
+      type=str,
+      nargs='+',
+      default=['decoded_sample_data:0', 'decoded_sample_data:1'],
+      help='Input name list involved in model graph.')
+  parser.add_argument(
+      '--output-name',
+      type=str,
+      default='labels_softmax:0',
+      help='Output name involved in model graph.')
+  parser.add_argument(
+      '--clip-duration-ms',
+      type=int,
+      default=1000,
+      help='Length of each audio clip fed into model.')
+  parser.add_argument(
+      '--clip-stride-ms',
+      type=int,
+      default=30,
+      help='Length of audio clip stride over main trap.')
+  parser.add_argument(
+      '--average_window_duration_ms',
+      type=int,
+      default=500,
+      help='Length of average window used for smoothing results.')
+  parser.add_argument(
+      '--detection-threshold',
+      type=float,
+      default=0.7,
+      help='The confidence for filtering unreliable commands')
+  parser.add_argument(
+      '--suppression_ms',
+      type=int,
+      default=500,
+      help='The time interval between every two adjacent commands')
+  parser.add_argument(
+      '--time-tolerance-ms',
+      type=int,
+      default=1500,
+      help='Time tolerance before and after the timestamp of this audio clip '
+      'to match ground truth')
+  parser.add_argument(
+      '--verbose',
+      action='store_true',
+      default=False,
+      help='Whether to print streaming accuracy on stdout.')
+
+  FLAGS, unparsed = parser.parse_known_args()
+  tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.INFO)
+  tf.compat.v1.app.run(main=main, argv=[sys.argv[0]] + unparsed)