TF_CAPI_EXPORT extern bool TFE_CancellationManagerRegisterCallback(
    TFE_CancellationManager*, TFE_CancellationToken token,
    const TFE_CancelCallback* c_callback, const char* callback_name);
TF_CAPI_EXPORT extern bool TFE_CancellationManagerDeregisterCallback(
    TFE_CancellationManager*, TFE_CancellationToken token);
TF_CAPI_EXPORT extern bool TFE_CancellationManagerTryDeregisterCallback(
    TFE_CancellationManager*, TFE_CancellationToken token);

    callback(context);
  };
  return tensorflow::unwrap(cancellation_manager)
      ->RegisterCallbackWithErrorLogging(token, callback, callback_name);
}

bool TFE_CancellationManagerDeregisterCallback(
    TFE_CancellationManager* cancellation_manager,
    TFE_CancellationToken token) {
  return tensorflow::unwrap(cancellation_manager)->DeregisterCallback(token);
}

  EXPECT_TRUE(TFE_CancellationManagerRegisterCallback(c_mgr, token3, &callback1,
                                                      "callback3"));

  EXPECT_TRUE(TFE_CancellationManagerDeregisterCallback(c_mgr, token1));
  EXPECT_TRUE(TFE_CancellationManagerTryDeregisterCallback(c_mgr, token3));

  TFE_CancellationManagerStartCancel(c_mgr);
  EXPECT_TRUE(TFE_CancellationManagerIsCancelled(c_mgr));

Ilya Tikhonovskiy <******@****.***> 1771913400 -0800