opt_einsum с Tensorflow 2

Question

Я хочу вычислить множественные сложные тензорные сокращения с тензорным потоком. У меня есть сокращения в «чередующейся» форме ввода для opt_einsum. Когда я пытаюсь провести различие с использованием тензорного потока, я получаю предупреждение

WARNING:tensorflow:AutoGraph could not transform <function _gcd_import at 0x000001ED736ABEA0> and will run it as-is.
Please report this to the TensorFlow team. When filing the bug, set the verbosity to 10 (on Linux, `export AUTOGRAPH_VERBOSITY=10`) and attach the full output.
Cause: Unable to locate the source code of <function _gcd_import at 0x000001ED736ABEA0>. Note that functions defined in certain environments, like the interactive Python shell do not expose their source code. If that is the case, you should to define them in a .py source file. If you are certain the code is graph-compatible, wrap the call using @tf.autograph.do_not_convert. Original error: could not get source code

Минимальным примером, иллюстрирующим проблему, будет

import numpy
import tensorflow as tf
import tensorflow.keras.optimizers as opt
import opt_einsum

a = tf.Variable(numpy.zeros((2, 2)))
b = tf.Variable(numpy.zeros((2, 2)))

opt = opt.Adam()        
opt.minimize(
    tf.function(
        lambda : opt_einsum.contract(a, (1, 2), b, (1, 2), backend = "tensorflow")
    ), var_list = [a, b]
)

Я пытался придерживаться инструкций из пакета opt_einsum, но они, кажется, сделаны для Tensorflow 1. Есть ли простой способ заставить это работать.