chore: import upstream snapshot with attribution

tensorflow/cc/client/client_session.cc

@@ -0,0 +1,188 @@
+/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/cc/client/client_session.h"
+
+#include <unordered_map>
+#include <utility>
+#include <vector>
+
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/mutex.h"
+#include "tensorflow/core/protobuf/config.pb.h"
+#include "tensorflow/core/public/session.h"
+#include "tensorflow/core/public/session_options.h"
+
+namespace tensorflow {
+
+class ClientSession::Impl {
+ private:
+  friend class ClientSession;
+
+  Impl(Session* session, std::shared_ptr<Graph> graph)
+      : session_(session), graph_(std::move(graph)) {}
+
+  static SessionOptions MakeDefaultSessionOptions(const std::string& target);
+  absl::Status MaybeExtendGraph() const;
+
+  std::unique_ptr<Session> session_;
+  std::shared_ptr<Graph> graph_;
+
+  mutable mutex mu_;
+  mutable int last_num_graph_nodes_ TF_GUARDED_BY(mu_) = 0;
+};
+
+ClientSession::ClientSession(const Scope& scope, const std::string& target)
+    : ClientSession(scope, Impl::MakeDefaultSessionOptions(target)) {}
+
+ClientSession::ClientSession(const Scope& scope) : ClientSession(scope, "") {}
+
+ClientSession::ClientSession(const Scope& scope,
+                             const SessionOptions& session_options) {
+  Session* new_session;
+  absl::Status status = NewSession(session_options, &new_session);
+  TF_CHECK_OK(status) << status;
+  impl_.reset(new Impl(new_session, scope.graph_as_shared_ptr()));
+  CHECK_NOTNULL(impl()->session_.get());
+}
+
+// Define destructor here so we can forward declare `Impl` in client_session.h.
+// If we define a dtor in the header file or use the default dtor,
+// unique_ptr<Impl> needs the complete type.
+ClientSession::~ClientSession() {}
+
+SessionOptions ClientSession::Impl::MakeDefaultSessionOptions(
+    const std::string& target) {
+  SessionOptions options;
+  options.env = Env::Default();
+  options.target = target;
+  return options;
+}
+
+absl::Status ClientSession::Run(const std::vector<Output>& fetch_outputs,
+                                std::vector<Tensor>* outputs) const {
+  return Run(FeedType{}, fetch_outputs, {}, outputs);
+}
+
+absl::Status ClientSession::Run(const FeedType& inputs,
+                                const std::vector<Output>& fetch_outputs,
+                                std::vector<Tensor>* outputs) const {
+  return Run(inputs, fetch_outputs, {}, outputs);
+}
+
+absl::Status ClientSession::Run(const FeedType& inputs,
+                                const std::vector<Output>& fetch_outputs,
+                                const std::vector<Operation>& run_outputs,
+                                std::vector<Tensor>* outputs) const {
+  return Run(RunOptions(), inputs, fetch_outputs, run_outputs, outputs,
+             nullptr);
+}
+
+absl::Status ClientSession::Impl::MaybeExtendGraph() const {
+  mutex_lock l(mu_);
+  int num_nodes = graph_->num_node_ids();
+  if (num_nodes > last_num_graph_nodes_) {
+    GraphDef graph_def;
+    graph_->ToGraphDefSubRange(&graph_def, last_num_graph_nodes_);
+    last_num_graph_nodes_ = num_nodes;
+    return session_->Extend(graph_def);
+  }
+  return absl::OkStatus();
+}
+
+absl::Status ClientSession::Run(const RunOptions& run_options,
+                                const FeedType& inputs,
+                                const std::vector<Output>& fetch_outputs,
+                                const std::vector<Operation>& run_outputs,
+                                std::vector<Tensor>* outputs,
+                                RunMetadata* run_metadata) const {
+  std::vector<std::pair<std::string, Tensor>> feeds;
+  feeds.reserve(inputs.size());
+  for (auto const& feed : inputs) {
+    TF_RETURN_IF_ERROR(feed.second.status);
+    feeds.emplace_back(std::piecewise_construct,
+                       std::forward_as_tuple(feed.first.name()),
+                       std::forward_as_tuple(feed.second.tensor));
+  }
+
+  std::vector<std::string> output_tensor_names;
+  output_tensor_names.reserve(fetch_outputs.size());
+  for (auto const& output : fetch_outputs) {
+    output_tensor_names.push_back(output.name());
+  }
+  std::vector<std::string> target_node_names;
+  target_node_names.reserve(run_outputs.size());
+  for (auto const& output : run_outputs) {
+    target_node_names.push_back(output.node()->name());
+  }
+  TF_RETURN_IF_ERROR(impl()->MaybeExtendGraph());
+  return impl()->session_->Run(run_options, feeds, output_tensor_names,
+                               target_node_names, outputs, run_metadata);
+}
+
+absl::Status ClientSession::Run(
+    const RunOptions& run_options, const FeedType& inputs,
+    const std::vector<Output>& fetch_outputs,
+    const std::vector<Operation>& run_outputs, std::vector<Tensor>* outputs,
+    RunMetadata* run_metadata,
+    const thread::ThreadPoolOptions& threadpool_options) const {
+  std::vector<std::pair<std::string, Tensor>> feeds;
+  for (auto const& feed : inputs) {
+    TF_RETURN_IF_ERROR(feed.second.status);
+    feeds.emplace_back(feed.first.name(), feed.second.tensor);
+  }
+  std::vector<std::string> output_tensor_names;
+  output_tensor_names.reserve(fetch_outputs.size());
+  for (auto const& output : fetch_outputs) {
+    output_tensor_names.push_back(output.name());
+  }
+  std::vector<std::string> target_node_names;
+  target_node_names.reserve(run_outputs.size());
+  for (auto const& output : run_outputs) {
+    target_node_names.push_back(output.node()->name());
+  }
+  TF_RETURN_IF_ERROR(impl()->MaybeExtendGraph());
+  return impl()->session_->Run(run_options, feeds, output_tensor_names,
+                               target_node_names, outputs, run_metadata,
+                               threadpool_options);
+}
+
+absl::Status ClientSession::MakeCallable(
+    const CallableOptions& callable_options, CallableHandle* out_handle) {
+  TF_RETURN_IF_ERROR(impl()->MaybeExtendGraph());
+  return impl()->session_->MakeCallable(callable_options, out_handle);
+}
+
+absl::Status ClientSession::RunCallable(CallableHandle handle,
+                                        const std::vector<Tensor>& feed_tensors,
+                                        std::vector<Tensor>* fetch_tensors,
+                                        RunMetadata* run_metadata) {
+  return impl()->session_->RunCallable(handle, feed_tensors, fetch_tensors,
+                                       run_metadata);
+}
+
+absl::Status ClientSession::RunCallable(
+    CallableHandle handle, const std::vector<Tensor>& feed_tensors,
+    std::vector<Tensor>* fetch_tensors, RunMetadata* run_metadata,
+    const thread::ThreadPoolOptions& options) {
+  return impl()->session_->RunCallable(handle, feed_tensors, fetch_tensors,
+                                       run_metadata, options);
+}
+
+absl::Status ClientSession::ReleaseCallable(CallableHandle handle) {
+  return impl()->session_->ReleaseCallable(handle);
+}
+
+}  // end namespace tensorflow

tensorflow/cc/client/client_session.h

@@ -0,0 +1,164 @@
+/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CC_CLIENT_CLIENT_SESSION_H_
+#define TENSORFLOW_CC_CLIENT_CLIENT_SESSION_H_
+
+#include <memory>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+#include "tensorflow/cc/framework/ops.h"
+#include "tensorflow/cc/framework/scope.h"
+#include "tensorflow/core/public/session_options.h"
+
+namespace tsl {
+namespace thread {
+struct ThreadPoolOptions;
+}
+}  // namespace tsl
+
+namespace tensorflow {
+
+namespace thread {
+using tsl::thread::ThreadPoolOptions;
+}
+
+/// @addtogroup core
+/// @{
+
+/// A `ClientSession` object lets the caller drive the evaluation of the
+/// TensorFlow graph constructed with the C++ API.
+///
+/// Example:
+///
+///     Scope root = Scope::NewRootScope();
+///     auto a = Placeholder(root, DT_INT32);
+///     auto c = Add(root, a, {41});
+///
+///     ClientSession session(root);
+///     std::vector<Tensor> outputs;
+///
+///     Status s = session.Run({ {a, {1}} }, {c}, &outputs);
+///     if (!s.ok()) { ... }
+class ClientSession {
+ public:
+  /// A data type to represent feeds to a Run call.
+  ///
+  /// This is a map of `Output` objects returned by op-constructors to the value
+  /// to feed them with. See `Input::Initializer` for details on what can be
+  /// used as feed values.
+  typedef std::unordered_map<Output, Input::Initializer, OutputHash> FeedType;
+
+  /// Create a new session to evaluate the graph contained in `scope` by
+  /// connecting to the TensorFlow runtime specified by `target`.
+  ClientSession(const Scope& scope, const std::string& target);
+
+  /// Same as above, but use the empty string ("") as the target specification.
+  explicit ClientSession(const Scope& scope);
+
+  /// Create a new session, configuring it with `session_options`.
+  ClientSession(const Scope& scope, const SessionOptions& session_options);
+
+  ~ClientSession();
+
+  /// Evaluate the tensors in `fetch_outputs`. The values are returned as
+  /// `Tensor` objects in `outputs`. The number and order of `outputs` will
+  /// match `fetch_outputs`.
+  absl::Status Run(const std::vector<Output>& fetch_outputs,
+                   std::vector<Tensor>* outputs) const;
+
+  /// Same as above, but use the mapping in `inputs` as feeds.
+  absl::Status Run(const FeedType& inputs,
+                   const std::vector<Output>& fetch_outputs,
+                   std::vector<Tensor>* outputs) const;
+
+  /// Same as above. Additionally runs the operations ins `run_outputs`.
+  absl::Status Run(const FeedType& inputs,
+                   const std::vector<Output>& fetch_outputs,
+                   const std::vector<Operation>& run_outputs,
+                   std::vector<Tensor>* outputs) const;
+
+  /// Use `run_options` to turn on performance profiling. `run_metadata`, if not
+  /// null, is filled in with the profiling results.
+  absl::Status Run(const RunOptions& run_options, const FeedType& inputs,
+                   const std::vector<Output>& fetch_outputs,
+                   const std::vector<Operation>& run_outputs,
+                   std::vector<Tensor>* outputs,
+                   RunMetadata* run_metadata) const;
+
+  /// Same as above. Additionally allows user to provide custom threadpool
+  /// implementation via ThreadPoolOptions.
+  absl::Status Run(const RunOptions& run_options, const FeedType& inputs,
+                   const std::vector<Output>& fetch_outputs,
+                   const std::vector<Operation>& run_outputs,
+                   std::vector<Tensor>* outputs, RunMetadata* run_metadata,
+                   const thread::ThreadPoolOptions& threadpool_options) const;
+
+  /// \brief A handle to a subgraph, created with
+  /// `ClientSession::MakeCallable()`.
+  typedef int64_t CallableHandle;
+
+  /// \brief Creates a `handle` for invoking the subgraph defined by
+  /// `callable_options`.
+  /// NOTE: This API is still experimental and may change.
+  absl::Status MakeCallable(const CallableOptions& callable_options,
+                            CallableHandle* out_handle);
+
+  /// \brief Invokes the subgraph named by `handle` with the given options and
+  /// input tensors.
+  ///
+  /// The order of tensors in `feed_tensors` must match the order of names in
+  /// `CallableOptions::feed()` and the order of tensors in `fetch_tensors` will
+  /// match the order of names in `CallableOptions::fetch()` when this subgraph
+  /// was created.
+  /// NOTE: This API is still experimental and may change.
+  absl::Status RunCallable(CallableHandle handle,
+                           const std::vector<Tensor>& feed_tensors,
+                           std::vector<Tensor>* fetch_tensors,
+                           RunMetadata* run_metadata);
+
+  /// \brief Invokes the subgraph named by `handle` with the given options and
+  /// input tensors.
+  ///
+  /// The order of tensors in `feed_tensors` must match the order of names in
+  /// `CallableOptions::feed()` and the order of tensors in `fetch_tensors` will
+  /// match the order of names in `CallableOptions::fetch()` when this subgraph
+  /// was created.
+  /// NOTE: This API is still experimental and may change.
+  absl::Status RunCallable(CallableHandle handle,
+                           const std::vector<Tensor>& feed_tensors,
+                           std::vector<Tensor>* fetch_tensors,
+                           RunMetadata* run_metadata,
+                           const thread::ThreadPoolOptions& options);
+
+  /// \brief Releases resources associated with the given `handle` in this
+  /// session.
+  /// NOTE: This API is still experimental and may change.
+  absl::Status ReleaseCallable(CallableHandle handle);
+
+ private:
+  class Impl;
+  std::unique_ptr<Impl> impl_;
+  Impl* impl() { return impl_.get(); }
+  const Impl* impl() const { return impl_.get(); }
+};
+
+/// @}
+
+}  // end namespace tensorflow
+
+#endif  // TENSORFLOW_CC_CLIENT_CLIENT_SESSION_H_

tensorflow/cc/client/client_session_test.cc

@@ -0,0 +1,270 @@
+/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#define EIGEN_USE_THREADS
+
+#include "tensorflow/cc/client/client_session.h"
+
+#include <utility>
+#include <vector>
+
+#include "absl/synchronization/barrier.h"
+#include "unsupported/Eigen/CXX11/Tensor"  // from @eigen_archive
+#include "tensorflow/cc/ops/standard_ops.h"
+#include "tensorflow/core/framework/tensor_testutil.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/lib/core/threadpool.h"
+#include "tensorflow/core/lib/core/threadpool_options.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/util/work_sharder.h"
+
+namespace tensorflow {
+namespace {
+
+using ops::Add;
+using ops::BatchMatMul;
+using ops::Const;
+using ops::Mul;
+using ops::Placeholder;
+using ops::Sub;
+
+tensorflow::SessionOptions GetSessionOptions() {
+  tensorflow::SessionOptions options;
+  // Disable optimizations for static graph to allow calls to Session::Extend.
+  options.config.mutable_experimental()->set_disable_optimize_for_static_graph(
+      true);
+  return options;
+}
+
+class CustomThreadPoolImpl : public thread::ThreadPoolInterface {
+ public:
+  explicit CustomThreadPoolImpl(int numThreads) {
+    underlying_threadpool_.reset(new thread::ThreadPool(
+        tensorflow::Env::Default(), "custom_threadpool", numThreads));
+    num_schedule_called_ = 0;
+  }
+
+  void Schedule(std::function<void()> fn) override {
+    num_schedule_called_ += 1;
+    underlying_threadpool_->Schedule(std::move(fn));
+  }
+
+  void ScheduleWithHint(std::function<void()> fn, int start, int end) override {
+    num_schedule_called_ += 1;
+    underlying_threadpool_->ScheduleWithHint(std::move(fn), start, end);
+  }
+
+  void Cancel() override {}
+
+  int NumThreads() const override {
+    return underlying_threadpool_->NumThreads();
+  }
+
+  int CurrentThreadId() const override {
+    return underlying_threadpool_->CurrentThreadId();
+  }
+
+  int GetNumScheduleCalled() { return num_schedule_called_; }
+
+ private:
+  int num_schedule_called_;
+  std::unique_ptr<tensorflow::thread::ThreadPool> underlying_threadpool_;
+};
+
+TEST(ClientSessionTest, Basic) {
+  Scope root = Scope::NewRootScope();
+  auto c = Const(root, {{1, 1}});
+  ClientSession session(root);
+  std::vector<Tensor> outputs;
+
+  TF_EXPECT_OK(session.Run({c}, &outputs));
+  test::ExpectTensorEqual<int>(outputs[0], test::AsTensor<int>({1, 1}, {1, 2}));
+}
+
+TEST(ClientSessionTest, Feed) {
+  Scope root = Scope::NewRootScope();
+  auto a = Placeholder(root, DT_INT32);
+  auto b = Placeholder(root, DT_INT32);
+  auto c = Add(root, a, b);
+  ClientSession session(root);
+  std::vector<Tensor> outputs;
+
+  TF_EXPECT_OK(session.Run({{a, 1}, {b, 41}}, {c}, &outputs));
+  test::ExpectTensorEqual<int>(outputs[0], test::AsTensor<int>({42}, {}));
+}
+
+TEST(ClientSessionTest, Extend) {
+  Scope root = Scope::NewRootScope();
+  auto a = Placeholder(root, DT_INT32, Placeholder::Shape({2}));
+  auto c = Add(root, a, {2, 2});
+  ClientSession session(root, GetSessionOptions());
+  std::vector<Tensor> outputs;
+
+  TF_EXPECT_OK(session.Run({{a, {1, 1}}}, {c}, &outputs));
+  test::ExpectTensorEqual<int>(outputs[0], test::AsTensor<int>({3, 3}, {2}));
+
+  auto d = Add(root, c, {39, 39});
+  outputs.clear();
+  TF_EXPECT_OK(session.Run({{a, {-10, 1}}}, {d}, &outputs));
+  test::ExpectTensorEqual<int>(outputs[0], test::AsTensor<int>({31, 42}, {2}));
+}
+
+TEST(ClientSessionTest, MultiThreadedWithDefaultThreadpool) {
+  Scope root = Scope::NewRootScope();
+  auto a = Add(root, {1, 2}, {3, 4});
+  auto b = Mul(root, {1, 2}, {3, 4});
+  ClientSession session(root, GetSessionOptions());
+  {
+    thread::ThreadPool thread_pool(Env::Default(), "pool", 2);
+    thread_pool.Schedule([&session, a]() {
+      std::vector<Tensor> outputs;
+      TF_EXPECT_OK(session.Run({a}, &outputs));
+      test::ExpectTensorEqual<int>(outputs[0],
+                                   test::AsTensor<int>({4, 6}, {2}));
+    });
+    thread_pool.Schedule([&session, b]() {
+      std::vector<Tensor> outputs;
+      TF_EXPECT_OK(session.Run({b}, &outputs));
+      test::ExpectTensorEqual<int>(outputs[0],
+                                   test::AsTensor<int>({3, 8}, {2}));
+    });
+  }
+  auto c = Sub(root, b, a);
+  std::vector<Tensor> outputs;
+  TF_EXPECT_OK(session.Run({c}, &outputs));
+  test::ExpectTensorEqual<int>(outputs[0], test::AsTensor<int>({-1, 2}, {2}));
+}
+
+TEST(ClientSessionTest, MultiThreadedWithCustomThreadpool) {
+  Scope root = Scope::NewRootScope();
+  int num_threads = 3;
+  auto a = Add(root, {1, 2}, {3, 4});
+  auto b = Mul(root, {1, 2}, {3, 4});
+  ClientSession session(root, GetSessionOptions());
+
+  auto inter_op_threadpool =
+      absl::make_unique<CustomThreadPoolImpl>(num_threads);
+  ASSERT_EQ(inter_op_threadpool->GetNumScheduleCalled(), 0);
+
+  auto intra_op_threadpool =
+      absl::make_unique<CustomThreadPoolImpl>(num_threads);
+  ASSERT_EQ(intra_op_threadpool->GetNumScheduleCalled(), 0);
+
+  tensorflow::thread::ThreadPoolOptions threadPoolOptions;
+  threadPoolOptions.inter_op_threadpool = inter_op_threadpool.get();
+  threadPoolOptions.intra_op_threadpool = intra_op_threadpool.get();
+
+  {
+    thread::ThreadPool thread_pool(Env::Default(), "pool", 2);
+    thread_pool.Schedule([&session, a]() {
+      std::vector<Tensor> outputs;
+      TF_EXPECT_OK(session.Run(RunOptions(), ClientSession::FeedType{}, {a}, {},
+                               &outputs, nullptr, thread::ThreadPoolOptions()));
+      test::ExpectTensorEqual<int>(outputs[0],
+                                   test::AsTensor<int>({4, 6}, {2}));
+    });
+    thread_pool.Schedule([&session, b]() {
+      std::vector<Tensor> outputs;
+      TF_EXPECT_OK(session.Run(RunOptions(), ClientSession::FeedType{}, {b}, {},
+                               &outputs, nullptr, thread::ThreadPoolOptions()));
+      test::ExpectTensorEqual<int>(outputs[0],
+                                   test::AsTensor<int>({3, 8}, {2}));
+    });
+  }
+  auto c = Sub(root, b, a);
+  std::vector<Tensor> outputs;
+  TF_EXPECT_OK(session.Run(RunOptions(), ClientSession::FeedType{}, {c}, {},
+                           &outputs, nullptr, thread::ThreadPoolOptions()));
+  test::ExpectTensorEqual<int>(outputs[0], test::AsTensor<int>({-1, 2}, {2}));
+}
+
+TEST(ClientSessionTest, CallableWithDefaultThreadPool) {
+  Scope root = Scope::NewRootScope();
+  auto a = Placeholder(root, DT_INT32);
+  auto b = Placeholder(root, DT_INT32);
+  auto c = Add(root, a, b);
+  ClientSession session(root);
+  std::vector<Tensor> outputs;
+
+  CallableOptions options;
+  options.add_feed(a.node()->name());
+  options.add_feed(b.node()->name());
+  options.add_fetch(c.node()->name());
+  ClientSession::CallableHandle callable;
+  TF_CHECK_OK(session.MakeCallable(options, &callable));
+  TF_EXPECT_OK(session.RunCallable(
+      callable, {test::AsTensor<int>({1}, {}), test::AsTensor<int>({41}, {})},
+      &outputs, nullptr));
+  test::ExpectTensorEqual<int>(outputs[0], test::AsTensor<int>({42}, {}));
+  TF_EXPECT_OK(session.ReleaseCallable(callable));
+}
+
+TEST(ClientSessionTest, CallableWithCustomThreadPool) {
+  Scope root = Scope::NewRootScope();
+  int num_threads = 3;
+
+  TensorShape data_shape({1, 1});
+  auto a = Placeholder(root, DT_INT32, Placeholder::Shape(data_shape));
+  auto b = Placeholder(root, DT_INT32, Placeholder::Shape(data_shape));
+  auto c = BatchMatMul(root, a, b);
+  ClientSession session(root);
+  std::vector<Tensor> outputs;
+
+  auto inter_op_threadpool =
+      absl::make_unique<CustomThreadPoolImpl>(num_threads);
+  ASSERT_EQ(inter_op_threadpool->GetNumScheduleCalled(), 0);
+
+  auto intra_op_threadpool =
+      absl::make_unique<CustomThreadPoolImpl>(num_threads);
+  ASSERT_EQ(intra_op_threadpool->GetNumScheduleCalled(), 0);
+
+  tensorflow::thread::ThreadPoolOptions threadPoolOptions;
+  threadPoolOptions.inter_op_threadpool = inter_op_threadpool.get();
+  threadPoolOptions.intra_op_threadpool = intra_op_threadpool.get();
+
+  CallableOptions options;
+  options.add_feed(a.node()->name());
+  options.add_feed(b.node()->name());
+  options.add_fetch(c.node()->name());
+  ClientSession::CallableHandle callable;
+  TF_CHECK_OK(session.MakeCallable(options, &callable));
+
+  // This is needed to have BatchMatMul computation be scheduled in the
+  // intra_op_threadpool.
+  absl::Barrier barrier(num_threads + 1);
+  for (int i = 0; i < num_threads; i++) {
+    intra_op_threadpool->Schedule([&barrier, num_threads]() {
+      tensorflow::SetPerThreadMaxParallelism(num_threads - 1);
+      barrier.Block();
+    });
+  }
+  barrier.Block();
+
+  TF_EXPECT_OK(session.RunCallable(
+      callable,
+      {test::AsTensor<int>({2}, {1, 1}), test::AsTensor<int>({10}, {1, 1})},
+      &outputs, nullptr, threadPoolOptions));
+  test::ExpectTensorEqual<int>(outputs[0], test::AsTensor<int>({20}, {1, 1}));
+  TF_EXPECT_OK(session.ReleaseCallable(callable));
+  ASSERT_GT(inter_op_threadpool->GetNumScheduleCalled(), 0);
+  ASSERT_GT(intra_op_threadpool->GetNumScheduleCalled(), 0);
+
+  // Free intra_op_threadpool and wait for its threads to exit before freeing
+  // other objects (e.g. barrier). This is needed to avoid data race.
+  intra_op_threadpool.reset();
+}
+
+}  // namespace
+}  // namespace tensorflow