Added C-API based unit tests for GPU and XLA GPU testing.

Also refined the API comment for TF_NewSession().

PiperOrigin-RevId: 185739196

1 files changed, 14 insertions, 4 deletions

diff --git a/tensorflow/c/c_test_util.cc b/tensorflow/c/c_test_util.cc
index 2c5f08d672..a55af46ae2 100644
--- a/tensorflow/c/c_test_util.cc
+++ b/tensorflow/c/c_test_util.cc
@@ -163,10 +163,14 @@ TF_Operation* AddWithCtrlDependency(TF_Operation* l, TF_Operation* r,
   return TF_FinishOperation(desc, s);
 }
 
+// If `op_device` is non-empty, set the created op on that device.
 void BinaryOpHelper(const char* op_name, TF_Operation* l, TF_Operation* r,
                     TF_Graph* graph, TF_Status* s, const char* name,
-                    TF_Operation** op, bool check) {
+                    TF_Operation** op, const string& op_device, bool check) {
   TF_OperationDescription* desc = TF_NewOperation(graph, op_name, name);
+  if (!op_device.empty()) {
+    TF_SetDevice(desc, op_device.c_str());
+  }
   TF_AddInput(desc, {l, 0});
   TF_AddInput(desc, {r, 0});
   *op = TF_FinishOperation(desc, s);
@@ -176,13 +180,19 @@ void BinaryOpHelper(const char* op_name, TF_Operation* l, TF_Operation* r,
   }
 }
 
-TF_Operation* Min(TF_Operation* l, TF_Operation* r, TF_Graph* graph,
-                  TF_Status* s, const char* name) {
+TF_Operation* MinWithDevice(TF_Operation* l, TF_Operation* r, TF_Graph* graph,
+                            const string& op_device, TF_Status* s,
+                            const char* name) {
   TF_Operation* op;
-  BinaryOpHelper("Min", l, r, graph, s, name, &op, true);
+  BinaryOpHelper("Min", l, r, graph, s, name, &op, op_device, true);
   return op;
 }
 
+TF_Operation* Min(TF_Operation* l, TF_Operation* r, TF_Graph* graph,
+                  TF_Status* s, const char* name) {
+  return MinWithDevice(l, r, graph, /*op_device=*/"", s, name);
+}
+
 TF_Operation* Add(TF_Output l, TF_Output r, TF_Graph* graph, TF_Status* s,
                   const char* name) {
   TF_OperationDescription* desc = TF_NewOperation(graph, "AddN", name);