Fix tensor_or_memref build error without math.h
Thanks for the fix!
https://github.com/tensorflow/tensorflow/pull/59536#issuecomment-1439039311

Since pevious PR has been rolled back, this one adopted reviewer's feedback to fix Mac build error.

Thanks @reedwm for the feedback and suggestion!