Repository: pytorch
Issue: torch.func.grad with respect to intermediate outputs
Labels: triaged,module: functorch

### 🚀 The feature, motivation and pitch

I'm working on vectorizing gradient computations with torch.func for vision explainability. I was able to use `torch.func.grad` and `torch.func.vmap` to compute saliency maps, which boils down to dy/dX, where y is my model output and X is my model input

At a higher level, I do: 

```python
 # Set up gradient operator
  compute_single_saliency = grad(
      pure_function, argnums=INPUT_IDX, has_aux=False
  )

  # Set up vmap operator for entire batch
  compute_batch_saliency = vmap(
      compute_single_saliency, in_dims=(0, 0, None)
  )

  # Execute the transformed function
  # vmap will automatically unbatch the arguments
  per_sample_grad = compute_batch_saliency(inputs, targets, model)
```

where my `pure_function` is just: 

```python
def compute_score(
    x: torch.Tensor,
    y: torch.Tensor,
    model: torch.nn.Module,
) -> torch.Tensor:
    # Add batch dim for the model
    x = x.unsqueeze(0)
    y = y.unsqueeze(0)

    # Forward
    output = model(x)

    # Get scalar output
    score = torch.sum(output * y)
    return score
```

*However*, this becomes a problem for a different type of explainability technique that requires dy/dA (different from dy/dX). In this case, A is a feature map, or an intermediate output from my model.

**I believe I'm stuck because `torch.func.grad` requires that gradients be taken with respect to function arguments. In this case, I cannot give A as a function argument because it is an intermediate output of the model.**

So, how do I compute dy/dA while taking advantage of vmap?

 

### Alternatives

I'm unfortunately not familiar enough with the design of torch.func to understand how this could be resolved. It just seems that as long as I am locked to taking a gradient with respect to an argument, I will not be able to take a gradient with respect to an intermediate output.

I will outline a couple attempts and the problems associated with them. I won't show the entire function, just the relevant parts. 

Attempt 1: This silently fails. The code executes and the output has the expected shape but everything is zeroes. Even though I can guarantee that activations is part of the computational graph for score, `grad` fails to compute the correct gradient and instead just returns a tensor of zeroes with the right shape. I suspect this is because `activation` did not directly play a role in computing the score while inside the `compute_score` scope.

```python
def compute_grad_cam(
    model: torch.nn.Module,
    inputs: torch.Tensor,
    targets: torch.Tensor,
    target_layer_name: str,
) -> torch.Tensor:

    # Grab target layer
    target_layer = get_layers_from_model(model, target_layer_name)

    # Set up hook
    act_container = {}
    hook = target_layer.register_forward_hook(
        get_activation(act_container, target_layer_name)
    )

    # Make a forward pass to get outputs
    outputs = model(inputs)

    # Grab activations off hook and remove hook
    activations = act_container[target_layer_name]
    hook.remove()

    # Define simple pure function to compute score
    def compute_score(output, y, activation):
        score = torch.sum(output * y)
        return score

    # Define grad
    compute_single_cam = grad(compute_score, argnums=2, has_aux=False)
   
   # apply vmap
   per_sample_grad = vmap(compute_single_cam, in_dims=(0, 0, 0))(
          outputs, targets, activations
      )
```

Okay, so I tried to swap out `func.grad` for `autograd.grad`

Attempt 2: This explicitly errors out with `RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn`. I am not sure why that is happening, when I check `requires_grad` for `outputs` and `activations` just after they are computed, I get `True`. Weirdly, inside the `compute_grad` function they return `False` and I am not allowed to set them to `True`.

```python
def compute_grad_cam(
    model: torch.nn.Module,
    inputs: torch.Tensor,
    targets: torch.Tensor,
    target_layer_name: str,
) -> torch.Tensor:

    # Grab target layer
    target_layer = get_layers_from_model(model, target_layer_name)

    # Set up hook
    act_container = {}
    hook = target_layer.register_forward_hook(
        get_activation(act_container, target_layer_name)
    )

    # Make a forward pass to get outputs
    outputs = model(inputs)

    # Grab activations off hook and remove hook
    activations = act_container[target_layer_name]
    hook.remove()

    def compute_grad(output, y, activation):
            score = torch.sum(output * y)
            return torch.autograd.grad(score, activation)

    # apply vmap
    per_sample_grad = vmap(compute_grad, in_dims=(0, 0, 0))(
        outputs, targets, activations
    )
```

cc @zou3519 @Chillee @samdow @kshitij12345 @janeyx99 

### Additional context

_No response_