Repository: pytorch
Issue: ExpandedWeights can't handle modules with tied weights
Labels: module: nn,triaged

### 🐛 Describe the bug

`call_for_per_sample_grads` throws a (seemingly irrelevant) error when processing a module with tied parameters (i.e. one parameter tensor is shared between two layers)

See this [forum thread](https://discuss.pytorch.org/t/how-to-share-weights-between-two-layers/55541) for an example where such a technique can be useful

Code:
```python
import torch
import torch.nn as nn
from torch.nn.utils._per_sample_grad import call_for_per_sample_grads

class SampleTiedWeights(nn.Module):
    def __init__(self):
        super().__init__()
        self.emb = nn.Embedding(100, 8)
        self.fc2 = nn.Linear(8, 100)

        w = torch.empty(100, 8)
        nn.init.uniform_(w, -100, 100)

        p = nn.Parameter(w)
        self.emb.weight = p
        self.fc2.weight = p

    def forward(self, x):
        x = x.unsqueeze(1)
        x = self.emb(x)
        x = self.fc2(x)
        x = x.squeeze(1)

        return x

m = SampleTiedWeights()
x = torch.randint(low=0, high=100, size=(512,))
out = call_for_per_sample_grads(m, x.shape[0], x)
out.sum().backward()
```

Error:
```
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
/var/folders/_r/hzffvgfd23sc3w0gr577_qmc0000gn/T/ipykernel_71302/2108605639.py in <module>
      2 x = torch.randint(low=0, high=100, size=(512,))
      3 out = call_for_per_sample_grads(m, x.shape[0], x)
----> 4 out.sum().backward()

~/Documents/opacus/venv/lib/python3.7/site-packages/torch/_tensor.py in backward(self, gradient, retain_graph, create_graph, inputs)
    394                 create_graph=create_graph,
    395                 inputs=inputs)
--> 396         torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)
    397 
    398     def register_hook(self, hook):

~/Documents/opacus/venv/lib/python3.7/site-packages/torch/autograd/__init__.py in backward(tensors, grad_tensors, retain_graph, create_graph, grad_variables, inputs)
    173     Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
    174         tensors, grad_tensors_, retain_graph, create_graph, inputs,
--> 175         allow_unreachable=True, accumulate_grad=True)  # Calls into the C++ engine to run the backward pass
    176 
    177 def grad(

~/Documents/opacus/venv/lib/python3.7/site-packages/torch/autograd/function.py in apply(self, *args)
    251                                "of them.")
    252         user_fn = vjp_fn if vjp_fn is not Function.vjp else backward_fn
--> 253         return user_fn(self, *args)
    254 
    255     def apply_jvp(self, *args):

~/Documents/opacus/venv/lib/python3.7/site-packages/torch/nn/utils/_expanded_weights/linear_expanded_weights.py in backward(ctx, grad_output)
     29 
     30         if input.requires_grad:
---> 31             results.append(grad_output.matmul(unpack_expanded_weight_or_tensor(weight)))
     32         else:
     33             results.append(None)

~/Documents/opacus/venv/lib/python3.7/site-packages/torch/nn/utils/_expanded_weights/expanded_weights_utils.py in unpack_expanded_weight_or_tensor(maybe_expanded_weight, func)
     73         return func(maybe_expanded_weight)
     74     elif isinstance(maybe_expanded_weight, torch.Tensor):
---> 75         raise RuntimeError("ExpandedWeights currently does not support a mixture of ExpandedWeight parameters "
     76                            "and normal Parameters. Please file and issue with pytorch/pytorch")
     77 

RuntimeError: ExpandedWeights currently does not support a mixture of ExpandedWeight parameters and normal Parameters. Please file and issue with pytorch/pytorch
```

### Versions

PyTorch version: 1.12.0
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 12.4 (x86_64)
GCC version: Could not collect
Clang version: 12.0.5 (clang-1205.0.22.9)
CMake version: Could not collect
Libc version: N/A

Python version: 3.7.5 (default, Oct 22 2019, 10:35:10)  [Clang 10.0.1 (clang-1001.0.46.4)] (64-bit runtime)
Python platform: Darwin-21.5.0-x86_64-i386-64bit
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
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] mypy==0.910
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.21.4
[pip3] pytorch-lightning==1.5.3
[pip3] torch==1.12.0
[pip3] torchaudio==0.12.0
[pip3] torchmetrics==0.6.0
[pip3] torchvision==0.13.0
[conda] Could not collect

cc @albanD @mruberry @jbschlosser @walterddr @kshitij12345 @saketh-are