tensorflow einsum vs. matmul vs. tensordot

En tensorflow, las funciones tf.einsum , tf.matmul y tf.tensordot pueden usarse para las mismas tareas. (Me doy cuenta de que tf.einsum y tf.tensordot tienen definiciones más generales; también me doy cuenta de que tf.matmul tiene funcionalidad por lotes). En una situación en la que se podría usar cualquiera de los tres, ¿una de las funciones tiende a ser más rápida? ¿Hay otras reglas de recomendación?

Por ejemplo, suponga que A es un tensor de rango 2, y b es un tensor de rango 1 y desea calcular el producto c_j = A_ij b_j . De las tres opciones:

c = tf.einsum('ij,j->i', A, b)

c = tf.matmul(A, tf.expand_dims(b,1))

c = tf.tensordot(A, b, 1)

¿Alguno es generalmente preferible a los demás?

Tanto tf.tensordot() como tf.einsum() son azúcares sintácticos que envuelven una o más invocaciones de tf.matmul() (aunque en algunos casos especiales tf.einsum() puede reducir al elemento simple más simple tf.multiply() ) .

En el límite, esperaría que las tres funciones tengan un rendimiento equivalente para el mismo cálculo. Sin embargo, para matrices más pequeñas puede ser más eficiente usar tf.matmul() directamente, ya que produciría un gráfico TensorFlow más simple con menos operaciones y, por lo tanto, los costos de invocación por operación serán menores.