| 284 | } |
| 285 | |
| 286 | void Split3Helper(TF_Operation* input, TF_Graph* graph, TF_Status* s, |
| 287 | const char* name, TF_Operation** op) { |
| 288 | TF_Operation* zero = ScalarConst( |
| 289 | 0, graph, s, ::tensorflow::strings::StrCat(name, "_const0").c_str()); |
| 290 | TF_OperationDescription* desc = TF_NewOperation(graph, "Split", name); |
| 291 | TF_AddInput(desc, {zero, 0}); |
| 292 | TF_AddInput(desc, {input, 0}); |
| 293 | TF_SetAttrInt(desc, "num_split", 3); |
| 294 | TF_SetAttrType(desc, "T", TF_INT32); |
| 295 | // Set device to CPU since there is no version of split for int32 on GPU |
| 296 | // TODO(iga): Convert all these helpers and tests to use floats because |
| 297 | // they are usually available on GPUs. After doing this, remove TF_SetDevice |
| 298 | // call in c_api_function_test.cc |
| 299 | TF_SetDevice(desc, "/cpu:0"); |
| 300 | *op = TF_FinishOperation(desc, s); |
| 301 | ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s); |
| 302 | ASSERT_NE(*op, nullptr); |
| 303 | } |
| 304 | |
| 305 | TF_Operation* Split3(TF_Operation* input, TF_Graph* graph, TF_Status* s, |
| 306 | const char* name) { |
no test coverage detected