Soz, version 2 just to avoid doing a for loop over each dimension:
import torch
# x.shape => (B,N,3)
# e_x.shape => (B,3) {also e_y and e_x}
B=2
N=4
x=torch.rand(B,N,3)
x_custom=x.unsqueeze(3).expand(B,N,3,3)
e_x=torch.rand(B,3)
e_y=torch.rand(B,3)
e_z=torch.rand(B,3)

x_new = torch.zeros((B,N,3))

for b in range(B): # loop over batch
   for j in range(N): # loop over atoms
      x_new[b,j,0]=torch.dot(x[b,j,:], e_x[b,:]) # x component of atom j in the new reference frame
      x_new[b,j,1]=torch.dot(x[b,j,:], e_y[b,:]) # y component
      x_new[b,j,2]=torch.dot(x[b,j,:], e_z[b,:]) # z component


einsum = torch.einsum('bijk,bjk->bik',x_custom,torch.stack([e_x,e_y,e_z],dim=2))

print(torch.norm((x_new-einsum)))