[Python 3.X] Numpy: optimisation sur des opérations tensorielles

Bonjour à vous,

Je travaille présentement sur une fonction réalisant de simples calculs tensoriels, avec numpy 1.20.2 sous python 3.9.0.
Cette dernière prends donc 2 np.array (de dimensions spécifiques) en entré et doit retourner un np.array résultant des deux premiers.
voici un code représentant la-dite fonction :

Code:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
 
import numpy as np  
 
 
def process(a:np.array, b:np.array)->np.array:
    """
    Dans l'utilisation de cette fonction, trois situations sont possibles.
    Pour chaque cas, i et j > 0, a.ndim == 2 et b.ndim == 2 ou b.ndim == 3
    cas1: a.shape == b.shape == (i,j) soit ij,ij->ij
        dans ce cas, on souhaite juste une multiplication terme à terme
        des entrées, tel que c = a * b et donc c.shape == (i,j)
    
    cas 2: a.shape == (1,j) et b.shape == (j,j) soit 1j,jj->1j
        dans cette situation, c'est donc un dot product qui est attendu et
        a.shape[0] == 1 (pas de cas où a.shape[0] > 1 dans cette situation)
        c = a @ b ou encore c = np.dot(a,b) ; là encore, c.shape == (1,j)
    
    cas 3: a.shape == (i,j) et b.shape == (i,j,j), retour attendu: ij,ijj->ij
        pour cette situation le processus est particulier car le tenseur b est
        à traiter comme un empilement de matrices dans le sens où:
            pour tout "i":
                c[i] = a[i] @ b[i]
        donc pas de produit attendu sur la troisième dimension
        
    """
 
    # cas 1: a.shape == b.shape == (i,j)
    if a.shape == b.shape and a.ndim == b.ndim == 2:
        return a * b # j'attends bien ici une multiplication terme à terme
 
    # cas 2: a.shape == (1,j) et b.shape == (j,j)
    elif a.shape[1] == b.shape[0] and b.ndim == 2 and a.shape[0] == 1:
        return a @ b # ou encore return np.dot(a,b)
 
    # cas 3: a.shape == (i,j) et b.shape == (i,j,j)
    elif a.shape[0] == b.shape[0] and b.ndim == 3:
        # ici, la meilleure solution que j'ai trouvé:
        buff = np.zeros(shape=(a.shape))
        for i in range(a.shape[0]):
            buff[i] = np.dot(a[i], b[i])
        return buff
 
    # erreur si on rencontre un autre cas, ce qui ne devrais pas arriver
    else:
        raise ValueError("dimensions de a et/ou b non reconnues")
 
 
 
if __name__ == "__main__":
 
    # cas 1:
    a = np.random.randint(0,5, size=(1,4))
    b = np.random.randint(0,5, size=(1,4))
    c = process(a,b)
    print("cas 1: ", c.shape)
 
    # cas 1 bis:
    a = np.random.randint(0,5, size=(2,4))
    b = np.random.randint(0,5, size=(2,4))
    c = process(a,b)
    print("cas 1b: ", c.shape)
 
    # cas 2:
    a = a = np.random.randint(0,5, size=(1,4))
    b = b = np.random.randint(0,5, size=(4,4))
    c = process(a,b)
    print("cas 2: ", c.shape)
 
    # cas 3:
    a = np.random.randint(0,5, size=(3,4))
    b = np.random.randint(0,5, size=(3,4,4))
    c = process(a,b)
    print("cas 3: ", c.shape)

résultat:

Citation:

cas 1: (1, 4)
cas 1b: (2, 4)
cas 2: (1, 4)
cas 2: (3, 4)

ceci correspond bien à ce que je souhaite faire et les résultats numériques obtenus pour c semblent bien convenir aux résultats attendus.

Je pense que tout le monde comprends rapidement le problème: il ne me semble pas possible que le cas 1 et 2 puissent être traités avec une seule et même fonction compte tenu de la nature des calculs attendus; Cependant, peut-être que ceci existe tout de même auquel-cas, je suis preneur de la solution.
Pour le cas 3 en revanche, je suis absolument persuadé que numpy propose une solution bien plus performante que ma boucle sur la première dimension et, peut-être même que la fonction employée pour le cas 2 et 3 pourrait être la même.
Ainsi, sauriez-vous m’orienter vers une solution plus générale, plus élégante, mais surtout plus performante que celle avec laquelle je suis arrivé ?

Afin de limiter éventuellement des questions subsidiaires, sachez que j'ai déjà fait quelques recherches en vain sur le sujet dont voici les résultats :

Code:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
 
def process_bis(a:np.array, b:np.array)->np.array:
    """
    ici, l'objectif est de trouver une solution vectorisée pour le cas 3 vu
    précédement, afin de rendre le code plus lisible et surtout plus performant.
    Pour rappel, on souhaite que ij,ijj->ij
    J'ai pu tester les trois solutions suivantes, sans succes.
    avec, pour exemple: 
    a = [[4,1,2,4],
         [1,2,4,0]]
    b = [[[0,0,3,0],
          [0,1,1,4],
          [2,3,0,0],
          [3,0,4,0]],
         
         [[2,3,2,4],
          [0,3,3,1],
          [0,2,1,2],
          [1,2,1,1]]]
    la valeur attendue est :
    c = [[16,7,23,4],
         [2,17,12,14]]
    """
 
    # solution A:
    sol_a = np.dot(a,b)
    # retour : sol_a.shape == (i,i,j) != (i,j,j)
    # sol_a = [[[16  7 29  4],
    #           [12 27 17 25]],
 
    #          [[ 8 14  5  8],
    #           [ 2 17 12 14]]]
 
 
    # solution B:
    sol_b = np.matmul(a,b)
    # même comportement que ci-dessus
    # retour : sol_b.shape == (i,i,j) != (i,j,j)
    # sol_b = [[[16  7 29  4], 
    #           [ 8 14  5  8]],
 
    #          [[12 27 17 25],
    #           [ 2 17 12 14]]]
 
 
    # solution C:
    sol_c = np.einsum("ij,ijj->ij",a,b)
    # les dimensions sont bonnes mais pas les valeurs 
    # c = [[0 1 0 0],
    #      [2 6 4 0]]
 
 
 
    return sol_a, sol_b, sol_c
 
 
if __name__ == "__main__":
    # process_bis:
    a = np.array([[4,1,2,4],
         [1,2,4,0]])
    b = np.array([[[0,0,3,0],
          [0,1,1,4],
          [2,3,0,0],
          [3,0,4,0]],
 
         [[2,3,2,4],
          [0,3,3,1],
          [0,2,1,2],
          [1,2,1,1]]])
 
    c = process(a,b)
    sol_a, sol_b, sol_c = process_bis(a,b)
    print("valeur attendue: ", c)
    print("sol A: ", sol_a)
    print("sol B: ", sol_b)
    print("sol C: ", sol_c)

Un grand merci par avance pour tout indice ou solution que vous pourriez m'apporter afin d'optimiser un peu cette fonction.
Pour info, cette dernière intervient dans une boucle assez costaud dans mon programme, et les dimensions des vecteurs / matrices en jeux peuvent être relativement grandes (entre 2 et plusieurs milliers de valeurs).

Merci et bonne journée !