Remove references to torch_xla._XLAC._xla_einsum from xla_sharding.py

Author: pgmoka

torch_xla/distributed/spmd/xla_sharding.py

@@ -704,19 +704,19 @@ def _einsum_linear_backward(grad_output: Tensor, input: Tensor, weight: Tensor,
     grad_input = grad_weight = grad_bias = None
 
     if needs_input_grad_input:
-      grad_input = torch_xla._XLAC._xla_einsum('...m,mn->...n',
+      grad_input = torch.einsum('...m,mn->...n',
                                                (grad_output, weight))
     else:
       grad_input = None
 
     if needs_input_grad_weight:
-      grad_weight = torch_xla._XLAC._xla_einsum('...m,...n->mn',
+      grad_weight = torch.einsum('...m,...n->mn',
                                                 (grad_output, input))
     else:
       grad_weight = None
 
     if bias is not None and needs_input_grad_bias:
-      grad_bias = torch_xla._XLAC._xla_einsum('...m->m', (grad_output,))
+      grad_bias = torch.einsum('...m->m', (grad_output,))
     else:
       grad_bias = None