}

// TODO(b/158596585): Replace this with a cost model analysis.
bool IsTrivialUnaryOperation(Operation* op) {
  return llvm::isa<TF::CastOp, TF::IdentityOp>(op);
}

// Adds outside compilation attributes to unary ops such as Identity/Cast ops
// at the head of TPU computation that is used only by other outside compiled
// ops. Identity ops and Cast ops is commonly added to the start of TPU

        return success();
      }
    }
    return op->emitOpError("Block has no predecessor");
  }

  if (v.getDefiningOp()->getNumOperands() == 1) {
    // If the value is originating from an unary op, assume it's something
    // simple like "cast" and keep tracing.
    return traceUpwardsToArgument(v.getDefiningOp()->getOperand(0), seen, out);
  } else {
    // Typically a tf.VarHandle op.