Repository: pytorch
Issue: Problem with combination of jacref, vjp, tanh in torchscript
Labels: oncall: jit

### 🐛 Describe the bug

Hello,

I am trying to trace the following specific function of a trained model as torchscript to load and use in C++:

```
jacrev(lambda p, t: vjp(model, p)[1](t)[0])
```

My approach to creating torchscript from this is slightly complicated (I am very open to suggestions for more straightforward approaches):

```
p_d = torch.zeros((1, 3))
t_d = torch.zeros((1, 1))
f = jacrev(lambda p, t: vjp(model, p)[1](t)[0])
fx = make_fx(functionalize(f, remove='mutations_and_views'))(p_d, t_d)
script_f = functorch.compile.ts_compile(fx, (p_d, t_d))
```

Depending on the activation function used in `model`, this will work or fail. Specifically when more than one tanh activation is used in model this will fail with:

```
RuntimeError: 
The size of tensor a (3) must match the size of tensor b (512) at non-singleton dimension 0
The above operation failed shape propagation in this context:
  File "<eval_with_key>.5 from /Users/TimSalzmann/miniconda3/envs/l4casadi/lib/python3.8/site-packages/torch/fx/experimental/proxy_tensor.py:551 in wrapped", line 60
    tanh_backward_2 = torch.ops.aten.tanh_backward(_unsafe_view_1, tanh)
    mul = torch.ops.aten.mul(tanh, -2)
    mul_1 = torch.ops.aten.mul(mul, mm_1);  mul = mm_1 = None
            ~~~~~~~~~~~~~~~~~~ <--- HERE
    mul_2 = torch.ops.aten.mul(_unsafe_view_1, mul_1);  _unsafe_view_1 = mul_1 = None
    t_10 = torch.ops.aten.t(t_5);  t_5 = None
```

Swapping `tanh` for `relu` or `jacrev` for `jacfwd` does not show the same behavior.

Appreciate any help!

Minimal example:
```
import functorch.compile
import torch

from torch.func import jacrev, jacfwd, functionalize, vjp
from torch.fx.experimental.proxy_tensor import make_fx


class Model(torch.nn.Module):
    def __init__(self, dim_in, dim_out, act):
        super().__init__()

        self.input_layer = torch.nn.Linear(dim_in, 512)

        hidden_layers = []
        for i in range(2): # If only one layer with one tanh activation is used this will work too.
            hidden_layers.append(torch.nn.Linear(512, 512))

        self.act = act
        self.hidden_layer = torch.nn.ModuleList(hidden_layers)
        self.out_layer = torch.nn.Linear(512, dim_out)

    def forward(self, x):
        x = self.input_layer(x)
        for layer in self.hidden_layer:
           x = self.act(layer(x))
        x = self.out_layer(x)
        return x

##########
# Using activation function RELU -> this will work.
##########
model = Model(3, 1, torch.relu)


def _vjp(p, x):
    return vjp(model, p)[1](x)[0]


p_d = torch.zeros((1, 3))
t_d = torch.zeros((1, 1))
f = jacrev(lambda p, t: vjp(model, p)[1](t)[0])
fx = make_fx(functionalize(f, remove='mutations_and_views'))(p_d, t_d)
script_f = functorch.compile.ts_compile(fx, (p_d, t_d))

print(script_f((p_d, t_d)))

##########
# Using activation function TANH and jacfwd -> this will work.
##########
model = Model(3, 1, torch.tanh)


p_d = torch.zeros((1, 3))
t_d = torch.zeros((1, 1))
f = jacfwd(lambda p, t: vjp(model, p)[1](t)[0])
fx = make_fx(functionalize(f, remove='mutations_and_views'))(p_d, t_d)

script_f = functorch.compile.ts_compile(fx, (p_d, t_d))

print(script_f((p_d, t_d)))

##########
# Using activation function TANH -> this will fail.
##########
model = Model(3, 1, torch.tanh)

p_d = torch.zeros((1, 3))
t_d = torch.zeros((1, 1))
f = jacrev(lambda p, t: vjp(model, p)[1](t)[0])
fx = make_fx(functionalize(f, remove='mutations_and_views'))(p_d, t_d)

script_f = functorch.compile.ts_compile(fx, (p_d, t_d)) # This fails

print(script_f((p_d, t_d)))
```

### Versions

PyTorch version: 2.3.0
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 14.0 (arm64)
GCC version: Could not collect
Clang version: 15.0.0 (clang-1500.0.40.1)
CMake version: version 3.29.3
Libc version: N/A

Python version: 3.8.18 (default, Sep 11 2023, 08:17:16)  [Clang 14.0.6 ] (64-bit runtime)
Python platform: macOS-14.0-arm64-arm-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:
Apple M1

Versions of relevant libraries:
[pip3] numpy==1.24.4
[pip3] torch==2.3.0
[conda] numpy                     1.24.4                   pypi_0    pypi
[conda] torch                     2.3.0                    pypi_0    pypi

cc @EikanWang @jgong5 @wenzhe-nrv @sanchitintel