pytorchの内積計算の結果が異なる

einsum

1import numpy as np 2import torch 3 4def p(x,M1): 5 v = x.shape[0] 6 I0 = torch.tensor([1, 0], dtype=torch.float32) 7 I1 = torch.tensor([0, 1], dtype=torch.float32) 8 I = I0 if x[0] == 0 else I1 9 _M1 = torch.tensor(M1[0][:], dtype=torch.float32) 10 Z = torch.einsum("ia,i->a",_M1,I) 11 for i in range(1,v-1): 12 M = torch.tensor(M1[i][:], dtype=torch.float32) 13 I = I0 if x[i] == 0 else I1 14 Z = torch.einsum("a,aib,i->b",Z,M,I) 15 M = torch.tensor(M1[v-1][:], dtype=torch.float32) 16 I = I0 if x[v-1] == 0 else I1 17 psi = torch.einsum("a,ai,i->...",Z,M,I) 18 return psi 19 20def p_batch(x,M1): 21 v = x.shape[1] 22 I0 = torch.tensor([[1, 0]], dtype=torch.float32) 23 I1 = torch.tensor([[0, 1]], dtype=torch.float32) 24 I = torch.where((x[:, 0] == 0)[:, None], I0, I1) 25 _M1 = torch.tensor(M1[0][:], dtype=torch.float32) 26 p = torch.einsum("ia,xi->xa",_M1,I) 27 for i in range(1,v-1): 28 M = torch.tensor(M1[i][:], dtype=torch.float32) 29 I = torch.where((x[:, i-1] == 0)[:, None], I0, I1) 30 p = torch.einsum("xa,aib,xi->xb",p,M,I) 31 M = torch.tensor(M1[v-1][:], dtype=torch.float32) 32 I = torch.where((x[:, v-1] == 0)[:, None], I0, I1) 33 p = torch.einsum("xa,ai,xi->x",p,M,I) 34 return p