Tensorflow: Performance bottleneck due to tf.function retracing

[lazyoracle]

Describe the bug
Possible tensorflow performance bottleneck due to repeated tracing of tf.function decorated modules

To Reproduce
In Simulated_Calibration.ipynb:

C3:STATUS:Saving as: /tmp/tmpaz0djjuu/c3logs/ORBIT_cal/2021_03_16_T_16_57_04/calibration.log
(5_w,10)-aCMA-ES (mu_w=3.2,w_1=45%) in dimension 4 (seed=1004483, Tue Mar 16 16:57:04 2021)
C3:STATUS:Adding initial point to CMA sample.
WARNING:tensorflow:5 out of the last 14 calls to <function tf_matmul_left at 0x7f97602a8dd0> triggered tf.function retracing. Tracing is expensive and the excessive number of tracings could be due to (1) creating @tf.function repeatedly in a loop, (2) passing tensors with different shapes, (3) passing Python objects instead of tensors. For (1), please define your @tf.function outside of the loop. For (2), @tf.function has experimental_relax_shapes=True option that relaxes argument shapes that can avoid unnecessary retracing. For (3), please refer to https://www.tensorflow.org/guide/function#controlling_retracing and https://www.tensorflow.org/api_docs/python/tf/function for  more details.

Suggested workarounds

Possible causes and solutions:

(1) creating @tf.function repeatedly in a loop, -- please define your @tf.function outside of the loop.
(2) passing tensors with different shapes, -- @tf.function has experimental_relax_shapes=True option that relaxes argument shapes that can avoid unnecessary retracing.
(3) passing Python objects instead of tensors. -- please refer to https://www.tensorflow.org/guide/function#controlling_retracing and https://www.tensorflow.org/api_docs/python/tf/function for more details.