<p>Changing from tf1 to tf2 I was exposed to the same question and after a little bit of experimenting I found the solution below that shows how to establish the interface between a function decorated with tf.function and a scipy optimizer. The important changes compared to the question are:</p>

<ol>
<li>As mentioned by Ives  scipy's lbfgs 
needs to get function value and gradient, so you need to provide a function that delivers both and then set <code>jac=True</code></li>
<li>scipy's lbfgs is a Fortran function that expects the interface to provide np.float64 arrays while tensorflow tf.function uses tf.float32. 
So one has to cast input and output.</li>
</ol>

<p>I provide an example of how this can be done for a toy problem here below. </p>

<pre><code>import tensorflow as tf
import numpy as np
import scipy.optimize as sopt

def model(x):
    return tf.reduce_sum(tf.square(x-tf.constant(2, dtype=tf.float32)))

@tf.function
def val_and_grad(x):
    with tf.GradientTape() as tape:
        tape.watch(x)
        loss = model(x)
    grad = tape.gradient(loss, x)
    return loss, grad

def func(x):
    return [vv.numpy().astype(np.float64)  for vv in val_and_grad(tf.constant(x, dtype=tf.float32))]

resdd= sopt.minimize(fun=func, x0=np.ones(5),
                                      jac=True, method='L-BFGS-B')

print("info:\n",resdd)
</code></pre>

<p>displays</p>

<pre><code>info:
       fun: 7.105427357601002e-14
 hess_inv: &lt;5x5 LbfgsInvHessProduct with dtype=float64&gt;
      jac: array([-2.38418579e-07, -2.38418579e-07, -2.38418579e-07, -2.38418579e-07,
       -2.38418579e-07])
  message: b'CONVERGENCE: NORM_OF_PROJECTED_GRADIENT_&lt;=_PGTOL'
     nfev: 3
      nit: 2
   status: 0
  success: True
        x: array([1.99999988, 1.99999988, 1.99999988, 1.99999988, 1.99999988])
</code></pre>

<h1>Benchmark</h1>

<p>For comparing speed
I use 
 the lbfgs optimizer for a style transfer 
problem (see <a href="https://arxiv.org/pdf/1910.09497.pdf" rel="noreferrer">here</a> for the network). Note, that for this problem the network parameters are fixed and the input signal is adapted. As the optimized parameters (the input signal) are 1D the function factory is not needed. </p>

<p>I compare four implementations</p>

<ol>
<li>TF1.12: TF1 with with ScipyOptimizerInterface </li>
<li>TF2.0 (E): the approach above without using tf.function decorators</li>
<li>TF2.0 (G): the approach above using tf.function decorators</li>
<li>TF2.0/TFP: using the lbfgs minimizer from
<a href="https://www.tensorflow.org/probability/api_docs/python/tfp/optimizer/lbfgs_minimize" rel="noreferrer">tensorflow_probability</a></li>
</ol>

<p>For this comparison the optimization is stopped after 300 iterations (generally for convergence the problem requires 3000 iterations)</p>

<h2>Results</h2>

<pre><code>Method       runtime(300it)      final loss         
TF1.12          240s                0.045     (baseline)
TF2.0 (E)       299s                0.045
TF2.0 (G)       233s                0.045
TF2.0/TFP       226s                0.053
</code></pre>

<p>The TF2.0 eager mode (TF2.0(E)) works correctly but is about 20% slower than the TF1.12 baseline version. TF2.0(G) with tf.function works fine and is marginally faster than TF1.12, which is a good thing to know.</p>

<p>The optimizer from tensorflow_probability (TF2.0/TFP) is slightly faster than TF2.0(G) using scipy's lbfgs but does not achieve the same error reduction. In fact the decrease of the loss over time is not monotonous which seems a bad sign. Comparing the two implementations of lbfgs (scipy and tensorflow_probability=TFP) it is clear that the Fortran code in scipy is significantly more complex.
So either the simplification of the algorithm in TFP is harming here or even the fact that TFP is performing all calculations in float32 may also be a problem.</p>
Use Scipy Optimizer with Tensorflow 2.0 for Neural Network training
<p>After the introduction of Tensorflow 2.0 the scipy interface (tf.contrib.opt.ScipyOptimizerInterface) has been removed. However, I would still like to use the scipy optimizer <strong>scipy.optimize.minimize(method=’L-BFGS-B’)</strong> to train a neural network (<strong>keras model sequential</strong>). In order for the optimizer to work, it requires as input a function <strong>fun(x0)</strong> with <strong>x0</strong> being an array of shape (n,). Therefore, the first step would be to "flatten" the weights matrices to obtain a vector with the required shape. To this end, I modified the code provided by <a href="https://pychao.com/2019/11/02/optimize-tensorflow-keras-models-with-l-bfgs-from-tensorflow-probability/" rel="nofollow noreferrer">https://pychao.com/2019/11/02/optimize-tensorflow-keras-models-with-l-bfgs-from-tensorflow-probability/</a>. This provides a function factory meant to create such a function <strong>fun(x0)</strong>. However, the code does not seem to work and the loss function does not decrease. I would be really grateful if someone could help me work this out.</p>

<p>Here the piece of code I am using:</p>

<pre><code>func = function_factory(model, loss_function, x_u_train, u_train)

# convert initial model parameters to a 1D tf.Tensor
init_params = tf.dynamic_stitch(func.idx, model.trainable_variables)
init_params = tf.cast(init_params, dtype=tf.float32)

# train the model with L-BFGS solver
results = scipy.optimize.minimize(fun=func, x0=init_params, method='L-BFGS-B')


def loss_function(x_u_train, u_train, network):
    u_pred = tf.cast(network(x_u_train), dtype=tf.float32)
    loss_value = tf.reduce_mean(tf.square(u_train - u_pred))
    return tf.cast(loss_value, dtype=tf.float32)


def function_factory(model, loss_f, x_u_train, u_train):
    """A factory to create a function required by tfp.optimizer.lbfgs_minimize.

    Args:
        model [in]: an instance of `tf.keras.Model` or its subclasses.
        loss [in]: a function with signature loss_value = loss(pred_y, true_y).
        train_x [in]: the input part of training data.
        train_y [in]: the output part of training data.

    Returns:
        A function that has a signature of:
            loss_value, gradients = f(model_parameters).
    """

    # obtain the shapes of all trainable parameters in the model
    shapes = tf.shape_n(model.trainable_variables)
    n_tensors = len(shapes)

    # we'll use tf.dynamic_stitch and tf.dynamic_partition later, so we need to
    # prepare required information first
    count = 0
    idx = [] # stitch indices
    part = [] # partition indices

    for i, shape in enumerate(shapes):
        n = np.product(shape)
        idx.append(tf.reshape(tf.range(count, count+n, dtype=tf.int32), shape))
        part.extend([i]*n)
        count += n

    part = tf.constant(part)


    def assign_new_model_parameters(params_1d):
        """A function updating the model's parameters with a 1D tf.Tensor.

        Args:
            params_1d [in]: a 1D tf.Tensor representing the model's trainable parameters.
        """

        params = tf.dynamic_partition(params_1d, part, n_tensors)
        for i, (shape, param) in enumerate(zip(shapes, params)):

            model.trainable_variables[i].assign(tf.cast(tf.reshape(param, shape), dtype=tf.float32))

    # now create a function that will be returned by this factory

    def f(params_1d):
        """
        This function is created by function_factory.
        Args:
            params_1d [in]: a 1D tf.Tensor.

        Returns:
            A scalar loss.
        """

        # update the parameters in the model
        assign_new_model_parameters(params_1d)
        # calculate the loss
        loss_value = loss_f(x_u_train, u_train, model)

        # print out iteration &amp; loss
        f.iter.assign_add(1)
        tf.print("Iter:", f.iter, "loss:", loss_value)

        return loss_value

    # store these information as members so we can use them outside the scope
    f.iter = tf.Variable(0)
    f.idx = idx
    f.part = part
    f.shapes = shapes
    f.assign_new_model_parameters = assign_new_model_parameters

    return f
</code></pre>

<p>Here <strong>model</strong> is an object tf.keras.Sequential.</p>

<p>Thank you in advance for any help!</p>