When using tfp.optimizer.lbfgs_minimize, encountered OperatorNotAllowedInGraphError: Iterating over a symbolic `tf.Tensor` is not allowed:

code

datatime = np.array([datetime.fromtimestamp(i / 1000000) for i in datatime])[20:]
depthlist = np.linspace(0, 22, data.shape[0])
data = data[3948:3970, 20:] / 1e5
data = data[:, :300]
datatime = datatime[:300]

obs_strain = data * 1e5

well = Well.Well()
well.set_well_path(1, 1, np.linspace(-11, 11, 22))
n = len(datatime)
taxis = np.linspace(0, n, n, dtype=np.float32)
length = np.linspace(0, 100, n, dtype=np.float32)

initial_width_array = np.linspace(0.1, 5, n, dtype=np.float32)
initial_height_array = np.ones_like(initial_width_array) * 10
initial_S1_array = np.zeros_like(initial_width_array)
initial_S2_array = np.zeros_like(initial_width_array)

initial_params = np.concatenate((initial_width_array, initial_height_array, initial_S1_array, initial_S2_array)).astype(np.float32)
prams = [taxis, well, length, n]
variables = tf.Variable([initial_params], dtype=tf.float32)

def green_function_model(taxis, length, heigth, width, well, S1=0, S2=0):
    frac = DynamicFracture.RectangularFracture()
    frac.length_growth(taxis, length, heigth, width, S1=S1, S2=S2)
    frac.set_monitor_wells(well)
    frac.calculate()
    strain = frac.gather_strain_data()[0]
    return strain

@tf.function
def objective_function_variable_width(sim_strain_params):
    global prams
    taxis, well, length, n = prams
    widths = tf.cast(sim_strain_params[:n], dtype=tf.float32)
    heights = tf.cast(sim_strain_params[n:2 * n], dtype=tf.float32)
    S1s = tf.cast(sim_strain_params[2 * n:3 * n], dtype=tf.float32)
    S2s = tf.cast(sim_strain_params[3 * n:], dtype=tf.float32)
    
    params = tf.stack([widths, heights, S1s, S2s], axis=1)

    def compute_strain(p):
        width, height, S1, S2 = p[0], p[1], p[2], p[3]
        return green_function_model(taxis, length, height, width, well, S1, S2)
    
    sim_strain_array = tf.vectorized_map(compute_strain, params)
    
    sim_strain_matrix = tf.stack(sim_strain_array)
    error = tf.reduce_sum(tf.square(sim_strain_matrix - obs_strain))
    
    return error

@tf.function
def train_model():
    return tfp.optimizer.lbfgs_minimize(
        objective_function_variable_width,
        initial_position=initial_params,
        num_correction_pairs=10,
        tolerance=1e-8,
        max_iterations=50,
        parallel_iterations=1
    )

res = train_model()

get
Traceback (most recent call last):
File "d:/me/wordspace/DDM_Process/demo2.py", line 115, in
res = train_model()
File "D:\app\envs\tensorflow_gpu\lib\site-packages\tensorflow\python\util\traceback_utils.py", line 153, in error_handler
raise e.with_traceback(filtered_tb) from None
File "D:\app\envs\tensorflow_gpu\lib\site-packages\tensorflow\python\framework\func_graph.py", line 1233, in autograph_handler
raise e.ag_error_metadata.to_exception(e)
tensorflow.python.framework.errors_impl.OperatorNotAllowedInGraphError: in user code:

File "d:/me/wordspace/DDM_Process/demo2.py", line 107, in train_model  *
    objective_function_variable_width,
File "D:\app\envs\tensorflow_gpu\lib\site-packages\tensorflow_probability\python\optimizer\lbfgs.py", line 278, in minimize  **
    initial_state = _get_initial_state(value_and_gradients_function,
File "D:\app\envs\tensorflow_gpu\lib\site-packages\tensorflow_probability\python\optimizer\lbfgs.py", line 297, in _get_initial_state   
    init_args = bfgs_utils.get_initial_state_args(
File "D:\app\envs\tensorflow_gpu\lib\site-packages\tensorflow_probability\python\optimizer\bfgs_utils.py", line 77, in get_initial_state_args
    f0, df0 = value_and_gradients_function(initial_position)

OperatorNotAllowedInGraphError: Iterating over a symbolic `tf.Tensor` is not allowed: AutoGraph did convert this function. This might indicate you are trying to use an unsupported feature.

Reproduce the error with the supplied demo2.py script, starting at train_model and the objective_function_variable_width callback. Trace the call into tensorflow_probability/python/optimizer/lbfgs.py and bfgs_utils.py, then determine a minimal supported reproduction and verify that the reported graph-mode failure is fixed or clearly documented.