Repository: pytorch
Issue: Convolution and norm with vmap jacobian failing
Labels: module: convolution,triaged,module: vmap,module: norms and normalization,module: forward ad,module: functorch

Hi,

I am trying to calculate batch Jacobians of a model with `vmap` and `jacfwd` and I want to map over the last dimension of the data. 

On the first call `psi = model(train_x)` the model behaves as expected. 

However, in the batch Jacobian call; `grad_psi = vmap(jacfwd(model), in_dims=-1)(train_x)` I get the error. `RuntimeError: Expected 2D (unbatched) or 3D (batched) input to conv1d, but got input of size: [3].`

I have gotten around this temporarily by unsqueeze-ing few dimensions to increase the dimensionality of the tensor but then I run into issues about normalizing a single sample. I have also tried some other norm functions than `LayerNorm` like `InstanceNorm`, `BatchNorm`, `functional.normalise`, but I have had no success with any of these.

I have also tried running the model with transpose input data and changing `in_dims` to be `0`, but no improvement there.

If anyone wants more context I am implementing a pointnet (https://arxiv.org/abs/1612.00593) style model, where my channel dimension (in code length 3) is x-y-z coordinates and the image dimension (in code length 100) are the samples.

This code was roughly developed from these sources:

https://towardsdatascience.com/deep-learning-on-point-clouds-implementing-pointnet-in-google-colab-1fd65cd3a263
https://medium.com/@itberrios6/point-net-from-scratch-78935690e496

However, they never attempt to calculate the Jacobian of this model and do not run into this issue.

Thanks in advance!

```

import torch
import math
from torch.func import jacfwd, vmap

train_x = torch.rand((3,100))

class NeuralNet(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = torch.nn.Conv1d(3, 30, 1)
        self.bn1 = torch.nn.LayerNorm(100)

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        return x

model = NeuralNet()

psi = model(train_x)
grad_psi = vmap(jacfwd(model), in_dims=-1)(train_x)


```

Full traceback:

```
Traceback (most recent call last):
  File "/home/user/dkl/minimum_bug.py", line 25, in <module>
    grad_psi = vmap(jacfwd(model), in_dims=-1)(train_x)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniconda3/envs/dkl/lib/python3.11/site-packages/torch/_functorch/apis.py", line 188, in wrapped
    return vmap_impl(func, in_dims, out_dims, randomness, chunk_size, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniconda3/envs/dkl/lib/python3.11/site-packages/torch/_functorch/vmap.py", line 278, in vmap_impl
    return _flat_vmap(
           ^^^^^^^^^^^
  File "/home/user/miniconda3/envs/dkl/lib/python3.11/site-packages/torch/_functorch/vmap.py", line 44, in fn
    return f(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^
  File "/home/user/miniconda3/envs/dkl/lib/python3.11/site-packages/torch/_functorch/vmap.py", line 391, in _flat_vmap
    batched_outputs = func(*batched_inputs, **kwargs)
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniconda3/envs/dkl/lib/python3.11/site-packages/torch/_functorch/eager_transforms.py", line 1139, in wrapper_fn
    results = vmap(push_jvp, randomness=randomness)(basis)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniconda3/envs/dkl/lib/python3.11/site-packages/torch/_functorch/apis.py", line 188, in wrapped
    return vmap_impl(func, in_dims, out_dims, randomness, chunk_size, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniconda3/envs/dkl/lib/python3.11/site-packages/torch/_functorch/vmap.py", line 278, in vmap_impl
    return _flat_vmap(
           ^^^^^^^^^^^
  File "/home/user/miniconda3/envs/dkl/lib/python3.11/site-packages/torch/_functorch/vmap.py", line 44, in fn
    return f(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^
  File "/home/user/miniconda3/envs/dkl/lib/python3.11/site-packages/torch/_functorch/vmap.py", line 391, in _flat_vmap
    batched_outputs = func(*batched_inputs, **kwargs)
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniconda3/envs/dkl/lib/python3.11/site-packages/torch/_functorch/eager_transforms.py", line 1130, in push_jvp
    output = _jvp_with_argnums(func, args, basis, argnums=argnums, has_aux=has_aux)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniconda3/envs/dkl/lib/python3.11/site-packages/torch/_functorch/vmap.py", line 44, in fn
    return f(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^
  File "/home/user/miniconda3/envs/dkl/lib/python3.11/site-packages/torch/_functorch/eager_transforms.py", line 976, in _jvp_with_argnums
    result_duals = func(*duals)
                   ^^^^^^^^^^^^
  File "/home/user/miniconda3/envs/dkl/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniconda3/envs/dkl/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/dkl/minimum_bug.py", line 18, in forward
    x = self.conv1(x)
        ^^^^^^^^^^^^^
  File "/home/user/miniconda3/envs/dkl/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniconda3/envs/dkl/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniconda3/envs/dkl/lib/python3.11/site-packages/torch/nn/modules/conv.py", line 310, in forward
    return self._conv_forward(input, self.weight, self.bias)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniconda3/envs/dkl/lib/python3.11/site-packages/torch/nn/modules/conv.py", line 306, in _conv_forward
    return F.conv1d(input, weight, bias, self.stride,
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: Expected 2D (unbatched) or 3D (batched) input to conv1d, but got input of size: [3]
```

edit: some format and wording changes


cc @zou3519 @Chillee @samdow @kshitij12345 @janeyx99