Cache graph functions and add If gradient test (#637)

Cache graph functions and add If gradient test (#637) #321