Le vecteur TensorFlow multiplié par la multiplication vectorielle

J'ai un vecteur de volatilités et une matrice de corrélations

volatilities = tf.constant([0.2, 0.4, 0.6], dtype=tf.float32)
correlations = tf.constant([[1, 0.25, 0.5], [0.25, 1, 0.3], [0.5, 0.3, 1]], dtype=tf.float32)

Je veux créer une matrice de covariance en utilisant ces valeurs

covariance = tf.tensordot(volatilities, volatilities, axes=0) * correlations

Jusqu'ici tout va bien. Maintenant, j'utilise cela dans un processus de formation par lots, et les formes de mes volatilités et corrélations sont (batch_size, 3)et (batch_size, 3, 3), respectivement. J'ai passé beaucoup de temps à essayer de comprendre comment obtenir la matrice de covariance. J'ai essayé d'abord d'obtenir une matrice de volatilité de forme (batch_size, 3, 3), puis scalaire multipliant avec les corrélations, mais je n'ai pas réussi à obtenir la matrice de volatilité. Le plus proche que j'obtiens est un (3, 3)tenseur en faisant ceci:

volatility_matrix = tf.tensordot(volatilities, volatilities, axes=[[0], [0]])

Comment puis-je obtenir de manière efficace ma matrice de corrélation avec la batch_sizedimension ?

Solution du problème

Comme le mentionne le post cité par @Lescurel, ce problème peut être résolu en utilisant tf.einsum:

volatility_matrix = tf.einsum("ij,ik->ijk", volas, volas)

résout le problème d'obtention de la matrice de volatilité. À partir de là, c'est juste une multiplication scalaire le long du deuxième axe. Plus précisément, pour obtenir ensuite la matrice de covariance, j'ai fait

covariance_matrix = tf.linalg.matmul(volatility_matrix, correlation_matrix)