auto tf_dlm_type = GetDlDataType(data_type, status);
  if (!status->status.ok()) {
    return nullptr;
  }

  TensorReference tensor_ref(*tensor);  // This will call buf_->Ref()
  auto* tf_dlm_tensor_ctx = new TfDlManagedTensorCtx(tensor_ref);
  tf_dlm_tensor_ctx->reference = tensor_ref;

  DLManagedTensor* dlm_tensor = &tf_dlm_tensor_ctx->tensor;
  dlm_tensor->manager_ctx = tf_dlm_tensor_ctx;

  auto output_ty = op.getType().dyn_cast<RankedTensorType>();
  if (!output_ty) return success();
  auto tensor_ty = op.getTensor().getType().cast<TensorType>();
  if (output_ty.hasStaticShape() && tensor_ty.hasStaticShape()) {
    const int64_t output_ty_size = output_ty.getNumElements();
    const int64_t tensor_ty_size = tensor_ty.getNumElements();
    if (tensor_ty_size != output_ty_size)

                                PatternRewriter &rewriter) const override {
    auto tensor_ty = mlir::dyn_cast<RankedTensorType>(op.getTensor().getType());
    auto indices_ty =
        mlir::dyn_cast<RankedTensorType>(op.getIndices().getType());
    auto updates_ty =
        mlir::dyn_cast<RankedTensorType>(op.getUpdates().getType());

    if (!tensor_ty || !indices_ty || !updates_ty) return failure();