Repository: pytorch
Issue: parameterizations.orthogonal does not work as intended with nn.GRU or nn.LSTM
Labels: module: rnn,triaged,module: nn.utils.parametrize

### 🐛 Describe the bug

This is a follow-up to [a post on the forum](https://discuss.pytorch.org/t/how-to-use-orthogonal-parameterization-with-gru/180487). 

I am attempting to create a GRU that uses orthogonal parameterization for the weight matrices stored in `weight_hh_l0`. From the [GRU documentation](https://pytorch.org/docs/stable/generated/torch.nn.GRU.html), I know that `weight_hh_l0` contains three weight matrices `W_hr`, `W_hz`, and `W_hn` concatenated together. The problem that I’m facing is that because the GRU module stores the matrices in a concatenated format, applying the orthogonal parameterization via `torch.nn.utils.parametrizations.orthogonal` yields a matrix that is entirely 0 for the bottom two-thirds of the rows.

```
import torch
from torch import nn
from torch.nn.utils.parametrizations import orthogonal
model_dim = 2
n_batch = 4
max_time = 20

input_tensor = torch.randn((max_time, n_batch, model_dim))

net = nn.GRU(model_dim, model_dim, num_layers=1, batch_first=False)
net = orthogonal(net, name=f"weight_hh_l0")
net.forward(input_tensor)

for name, param in net.named_parameters():
    print(name)
    print(param)

parametrizations.weight_hh_l0.original
Parameter containing:
tensor([[-1.,  0.],
        [ 0., -1.],
        [ 0.,  0.],
        [ 0.,  0.],
        [ 0.,  0.],
        [ 0.,  0.]], requires_grad=True)
```

Clearly, I do not wish to constrain two-thirds of my weights to be 0s, as this defeats the purpose of using a GRU.

I have solved this by writing a [custom parameterization](https://discuss.pytorch.org/t/how-to-use-orthogonal-parameterization-with-gru/180487/8?u=user0). However, I feel that this is not a great solution.

In my opinion, the functionality that ships with PyTorch should be compatible with other functionality that ships with PyTorch. The exceptions to this principle should be rare & documented clearly in the relevant documentation pages. There are several alternative ways that this goal can be achieved. These are the possible solutions that I see to this problem:

1. `torch.nn.utils.parametrizations.orthogonal` raises an exception when applied to a non-square matrix. This is consistent with the textbook mathematical definition of orthogonal. which stipulates that orthogonal matrices are square matrices. The goal of this exception is to alert users when they are applying orthogonal outside of its intended context. It will also make it impossible to apply `torch.nn.utils.parametrizations.orthogonal` in an inappropriate context (e.g. GRU, LSTM), which is desirable because that application is obviously not an intended usage of this function. Right now, the result is a semantic bug, so this change will protect naive users from making an error; I suspect there are many users who have committed this mistake and are unaware of it. Instead, raising an exception will alert users that to use `torch.nn.utils.parametrizations.orthogonal`, either they will have to change their model or write custom parameterizations. I do not foresee any real cost to making this change, because the way that it is implemented now implies that non-square weight matrices are constrained to have blocks of 0s, so users are not gaining any appreciable value from `torch.nn.utils.parametrizations.orthogonal` admitting non-square matrices.
2. Revise the documentation to alert users that `torch.nn.utils.parametrizations.orthogonal` may result in undesirable behavior when applied to certain PyTorch modules, including `nn.GRU` and `nn.LSTM`. This will make an effort to inform users that they'll have to write their own code to achieve certain goals, at least.
3. Change the GRU (and LSTM) classes to store weight matrices separately, instead of in the concatenated format. This would allow users to apply `torch.nn.utils.parametrizations.orthogonal` to any or all of these matrices & the results will be orthogonal, so it will "just work." Presumably there's some reason why the matrices are concatenated, but in the event that users care about using _parameterizations_ of the weight matrices, the concatenation choice requires writing & maintaining custom code to split, transform, and concatenate these matrices. It also means that applying `torch.nn.utils.parametrizations.orthogonal` creates semantic bugs (per above).
4. Extend `torch.nn.utils.parametrizations.orthogonal` in some way so that it is _not naïve_ about the kinds of weights it is being applied to. This does not seem like a great choice, because it means the function needs to inspect and make a decision about how to parameterize weights, and that seems hard to develop & maintain. (Right now, the signature only "sees" the weights matrix, so it doesn't have context about the originating class of the weights.) It may also be hard to communicate to users how it works & why it works the way it does.

### Versions

```Collecting environment information...
PyTorch version: 2.0.1
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 13.3.1 (x86_64)
GCC version: Could not collect
Clang version: 14.0.3 (clang-1403.0.22.14.1)
CMake version: Could not collect
Libc version: N/A

Python version: 3.10.6 (main, Aug 11 2022, 13:49:25) [Clang 13.1.6 (clang-1316.0.21.2.5)] (64-bit runtime)
Python platform: macOS-13.3.1-x86_64-i386-64bit
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
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

CPU:
Intel(R) Core(TM) i7-9750H CPU @ 2.60GHz

Versions of relevant libraries:
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.24.3
[pip3] torch==2.0.1
[conda] Could not collect
```

cc @zou3519