Repository: pytorch
Issue: Different gradients in torch.matmul depending on input shape
Labels: module: numerical-stability,triaged

## 🐛 Bug

`torch.matmul` often returns different gradients for the same matrices, when the computation is done with an additional dimension (batched version). So if A and B are 2D matrices:

```python
C = torch.matmul(A,B)
D = torch.matmul(A.unsqueeze(0), B.unsqueeze(0)).squeeze(0)
```

Computing the gradient from C and D will give different results. This doesn't happens always, however, and it seems to be an issue related to numerical precision.

## To Reproduce

Make repeated executions of random matrix multiplications with both methods, but sharing the input values. The gradients will diverge at some point. I have prepared the following code for an accurate reproduction:

```python
import torch
import random

""" Do matrix-multiplications and gradient computation
several times. When gradients differ, compute total absolute
difference wrt. reference gradient computed in plain Python. """

def compute_accumulated_difference():
	err_matmul, err_bmm = 0.0, 0.0
	for _ in range(1000):
		N = random.randint(1,100)
		M = random.randint(1,100)
		L = random.randint(1,100)
		A = new_matrix(N,M)
		B = new_matrix(M,L)
		refs = matmul_grad(A,B)
		a = torch.tensor(A, dtype=torch.float)
		b = torch.tensor(B, dtype=torch.float)
		emm, ebmm = compare_methods(a, b, refs)
		err_matmul += emm
		err_bmm += ebmm
	print(err_matmul)
	print(err_bmm)

# Matrix multiplication functions in plain Python and native double format
def new_matrix(N,M):
	return [ [ random.gauss(0,1) for _ in range(M) ] for _ in range(N) ]

def transpose(A):
	return list(zip(*A))

def matmul(A,B):
	N, M, L = len(A), len(B), len(B[0])
	res = [ [ 0.0 for _ in range(L) ] for _ in range(N) ]
	for i in range(N):
		for j in range(L):
			res[i][j] = sum([ A[i][k]*B[k][j] for k in range(M) ])
	return res

def matmul_grad(A,B):
	N, L = len(A), len(B[0])
	G = [ [ 1.0 for _ in range(L) ] for _ in range(N) ]
	dA = matmul(G, transpose(B))
	dB = matmul(transpose(A), G)
	return dA, dB

# Functions in PyTorch to compare
def compute_matmul(a,b):
	return torch.matmul(a,b)

def compute_bmm(a,b):
	return torch.matmul(a.unsqueeze(0), b.unsqueeze(0)).squeeze(0)

def compute_grads(a,b,fn):
	a = a.detach().requires_grad_()
	b = b.detach().requires_grad_()
	c = fn(a,b)
	c.sum().backward()
	return a.grad, b.grad

def compare_methods(a,b,refs):
	res_1 = compute_grads(a,b,compute_matmul)
	res_2 = compute_grads(a,b,compute_bmm)
	err_1, err_2 = 0.0, 0.0
	for g1, g2, ref in zip(res_1, res_2, refs):
		if not (g1==g2).all():
			r = torch.tensor(ref, dtype=torch.double)
			err_1 += (g1.double()-r).abs().sum().item()
			err_2 += (g2.double()-r).abs().sum().item()
	return err_1, err_2

if __name__ == '__main__':
	compute_accumulated_difference()
```

## Expected behavior

Gradients shouldn't change their values because the shape is different, so computing `torch.matmul` or its batched version (unsqueezing inputs and squeezing outputs or `torch.bmm`) should give the exact same results.

## Environment

```
PyTorch version: 1.7.1+cpu
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.5 LTS (x86_64)
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
Clang version: Could not collect
CMake version: version 3.10.2

Python version: 3.6 (64-bit runtime)
Is CUDA available: False
CUDA runtime version: No CUDA
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] adabelief-pytorch==0.2.1
[pip3] numpy==1.19.5
[pip3] pytorch-lightning==1.1.6
[pip3] torch==1.7.1+cpu
[pip3] torchfile==0.1.0
[conda] blas                      1.0                         mkl  
[conda] mkl                       2020.1                      217  
[conda] mkl-service               2.3.0            py38he904b0f_0  
[conda] mkl_fft                   1.1.0            py38h23d657b_0  
[conda] mkl_random                1.1.1            py38h0573a6f_0  
[conda] numpy                     1.18.5           py38ha1c710e_0  
[conda] numpy-base                1.18.5           py38hde5b4d6_0  
[conda] numpydoc                  1.1.0                      py_0
```