chore: import upstream snapshot with attribution

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()