Implement Bayesian optimization with TF or Jax instead of using sklearn

[Anselmoo]

Is your feature request related to a problem? Please describe.

Optimisation via Bayesian might cause some performance issues due to the evaluation of the Einstein Summation.

https://github.com/keras-team/keras-tuner/blob/a7a361f9521cb1033a05aba865c86eb30784d907/keras_tuner/tuners/bayesian.py#L124-L127

Describe the solution you'd like

Let performing jax.numpy.einsum can help.

Partially support can look like:

try:
    from jax.numpy import einsum
except ImportError:
    from numpy import einsum

or system dependent:

if sys.platform != "win32":
    import jax.numpy as np
else:
    import numpy as np

[haifeng-jin]

Thanks for the issue! We try to reduce the amount of dependencies. What about use TF (https://www.tensorflow.org/api_docs/python/tf/experimental/numpy/einsum)?

How much of the performance difference would you expect?

[Anselmoo]

@haifeng-jin, thx for the feedback; that's an excellent idea. I will setup a repo for benchmark!

[Anselmoo]

Dear @haifeng-jin

To be honest, I don't see a performance difference at all based on my manual test.

It's tested with the same test as bayesian_test.py, and I tried to use pytest-benchmark, but it didn't show any useful data.

So, I set up a local time based and got the following results for numpy vs. tf.experimental:

• Total time for Numpy Oracle: 727.1199822425842 sec for three runs
• Total time for TensorFlow Oracle: 780.9310626983643 sec for three runs

The log files are here

It's noticeable that after 60 trials, the total elapsed time is a little bit fluctuating between seconds and minutes. Maybe it's related to how the new hyper parameters will be added?

About the implementation

  y_var = np.ones(len(x), dtype=np.float)
  # y_var -= tf.experimental.numpy.einsum
  y_var -= tf.experimental.numpy.einsum(
            "ij,ij->i",
            tf.experimental.numpy.dot(kernel_trans, kernel_inv),
            kernel_trans,
        )

  # Convert y_var[y_var < 0] = 0.0 into a tensorflow conform way.
  y_var = tf.where(y_var < 0, tf.zeros_like(y_var), y_var)

Further references

https://github.com/Anselmoo/keras-tuner/blob/Anselmoo/issue766/keras_tuner/tuners/bayesian_tf.py

https://github.com/Anselmoo/keras-tuner/blob/f8de7077c2ee51eb33a21e295b319fbd3214daaf/keras_tuner/tuners/bayesian_tf.py#L124-L131

[haifeng-jin]

This issue is blocked. Please refer to #798 for the conditions when to start working on this issue.