TF_SetAttrType(desc, "T", TF_INT32);
  // Set device to CPU since there is no version of split for int32 on GPU
  // TODO(iga): Convert all these helpers and tests to use floats because
  // they are usually available on GPUs. After doing this, remove TF_SetDevice
  // call in c_api_function_test.cc
  TF_SetDevice(desc, "/cpu:0");
  *op = TF_FinishOperation(desc, s);
  ASSERT_EQ(TF_OK, TF_GetCode(s)) << TF_Message(s);

      TFE_NewContext(opts.get(), status.get()), TFE_DeleteContext);
  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());

  // Skip the test if no TPU is available.
  std::unique_ptr<TF_DeviceList, decltype(&TF_DeleteDeviceList)> devices(
      TFE_ContextListDevices(context.get(), status.get()), TF_DeleteDeviceList);
  ASSERT_EQ(TF_GetCode(status.get()), TF_OK) << TF_Message(status.get());

  ~GraphOperation() override {}

 private:
  friend class GraphContext;  // For access to op_.
  TF_Graph* g_;
  std::unique_ptr<TF_OperationDescription> op_;
  // Hold `op_type` and `op_name` till both are available since we need both
  // to build a graph operation.
  string op_type_;
  const char* op_name_ = nullptr;
  // TODO(srbs): Use this.
  string device_name_;
};