torch.einsum這是有關(guān)GPU內(nèi)部工作的查詢。我知道如何使用einsum。它是執(zhí)行所有可能的矩陣乘法，然后只選擇相關(guān)的矩陣乘法，還是僅執(zhí)行所需的計算？例如，考慮形狀 的兩個張量a和，我希望找到形狀的每個相應(yīng)張量 的點積。使用einsum，代碼為：b(N,P)ni(1,P)torch.einsum('ij,ij->i',a,b)在不使用 einsum 的情況下，獲取輸出的另一種方法是：torch.diag(a @ b.t())現(xiàn)在，第二個代碼應(yīng)該比第一個代碼執(zhí)行更多的計算（例如， if N= 2000，它執(zhí)行2000更多的計算）。然而，當我嘗試對這兩個操作進行計時時，它們完成所需的時間大致相同，這就引出了一個問題。是否einsum執(zhí)行所有組合（如第二個代碼），并挑選出相關(guān)值？要測試的示例代碼：import timeimport torchfor i in range(100):  a = torch.rand(50000, 256).cuda()  b = torch.rand(50000, 256).cuda()  t1 = time.time()  val = torch.diag(a @ b.t())  t2 = time.time()  val2 = torch.einsum('ij,ij->i',a,b)  t3 = time.time()  print(t2-t1,t3-t2, torch.allclose(val,val2))