Is it a good idea to use tensorflow instead of numpy for numerical approximations?
TensorFlow is a library that supports dataflow paradigm which makes it a good candidate for numerical calculations.

In general, TensorFlow data loading performance can be optimized using tf.data module. In particular, there is a concept called GPUs starvation, where GPUs are waiting for data to be loaded. GPU starvation can be minimized in TensorFlow with Pipelining, overlapping the preprocessing and model execution of a training step.

In particular, your code can be optimized by refactoring to Expectation_Value operate on vectors instead of scalars. Right now N and i are scalars and there is a for-loop that processes elements one at a time. tf.linalg.trace can take a tensor as an input so that N and i could be vectors. The function will be much faster once it operates on vectors and no longer uses a for-loop.