Repository: pytorch
Issue: Consider making torch.cond return zero rather than None for the gradients of tensors that are in the not-taken branch of the if-else.
Labels: high priority,triaged,oncall: pt2,module: higher order operators,module: pt2-dispatcher

### 🐛 Describe the bug

While working on #140979, I noticed that torch.cond()'s current backward pass will return None for the gradients of inputs that are not used in the branch selected by `pred`.

There are a few challenges with this:

1. This makes it basically impossible to capture torch.cond with a cuda graph. By specifying that a tensor's gradient is None rather than a tensor of zeros, we cannot pass this gradient into a kernel. Therefore, we cannot capture a single static cuda graph that can handle both sides of a torch.cond, unfortunately, since the arguments to later kernels would be dependent upon whether pred is true or false. We would be forced to fall back to cuda graph trees rather than using a single cuda graph for workloads using torch.cond that need to compute gradients, which is disappointing. This argument likely also applies to other accelerators like TPU's.

Here is an example that demonstrates this point:

```
import torch

torch.set_default_device("cuda")

def cond_branch(x, w1, w2):
    return torch.cond(x > 0, lambda x: w1 * x, lambda x: w2 * x, [x])

def imperative_branch(x, w1, w2):
    if x > 0:
        return w1 * x
    else:
        return w2 * x

x = torch.ones((), requires_grad=False)
neg_x = -torch.ones((), requires_grad=False)
w1 = torch.zeros((), requires_grad=True)
w2 = torch.zeros((), requires_grad=True)

y2 = cond_branch(x, w1, w2)

grad_w = torch.autograd.grad(y2, [w1, w2], allow_unused=True)

print("GALVEZ:grads=", grad_w)
```

Prints:

```
GALVEZ:grads= (tensor(1., device='cuda:0'), None)
```

Rather than 

```
GALVEZ:grads= (tensor(1., device='cuda:0'), tensor(0., device='cuda:0'))
```

2. This does not match the behavior of torch.vmap(torch.cond()), which does return a tensor of 0's. There are a few optimizers that are not no-ops when the gradient is 0. One obvious example is ADAM. Therefore, I think that we may encounter subtle differences under a transformation that *should* retain parity in terms of semantics.

Here is an example.

```
import torch

torch.set_default_device("cuda")

def cond_branch(x, w1, w2):
    return torch.cond(x > 0, lambda x: w1 * x, lambda x: w2 * x, [x])

def imperative_branch(x, w1, w2):
    if x > 0:
        return w1 * x
    else:
        return w2 * x

x = torch.ones((), requires_grad=False)
neg_x = -torch.ones((), requires_grad=False)
w1 = torch.zeros((), requires_grad=True)
w2 = torch.zeros((), requires_grad=True)

y2 = torch.vmap(cond_branch)(torch.unsqueeze(x, 0), torch.unsqueeze(w1, 0), torch.unsqueeze(w2, 0))

grad_w = torch.autograd.grad(y2, [w1, w2], allow_unused=True)

print("GALVEZ:grads=", grad_w)
```

Prints:

```
GALVEZ:grads= (tensor(1., device='cuda:0'), tensor(0., device='cuda:0'))
```

3. Some may not care about this, but it doesn't match the behavior of jax.lax.cond. I surmise that the reason for jax.lax.cond's behavior is that it TPU is a target, which cannot handle gradients being set to None rather than being set to 0, for similar to cuda graphs, because then it won't be able to create a single static compute graph.

```python
import jax
import jax.numpy as jnp

# Define the conditional branch function using jax.lax.cond                                                                                                                                                 
def cond_branch(x, w1, w2):
    return jax.lax.cond(
        x > 0,
        lambda x: w1 * x,
        lambda x: w2 * x,
        x,
    )

# Inputs                                                                                                                                                                                                    
x = jnp.ones(())
neg_x = -jnp.ones(())
w1 = jnp.zeros(())
w2 = jnp.zeros(())

w_grads = jax.grad(cond_branch, [1, 2])(x, w1, w2)

print("GALVEZ:", w_grads[0], w_grads[1])
```

Prints:

```
GALVEZ: 1.0 0.0
```

(Rather than 1.0 None)


Now, there is a "problem" with my proposal, which is that currently torch.cond does match the behavior of imperative if-else in python. See here:

```

torch.set_default_device("cuda")

def cond_branch(x, w1, w2):
    return torch.cond(x > 0, lambda x: w1 * x, lambda x: w2 * x, [x])

def imperative_branch(x, w1, w2):
    if x > 0:
        return w1 * x
    else:
        return w2 * x

x = torch.ones((), requires_grad=False)
neg_x = -torch.ones((), requires_grad=False)
w1 = torch.zeros((), requires_grad=True)
w2 = torch.zeros((), requires_grad=True)

y2 = imperative_branch(x, w1, w2)
y2.backward()
print(w1.grad)
print(w2.grad)
```

```
tensor(1., device='cuda:0')
None
```

I hope that this isn't consider to be particularly problematic, though. I think the benefits of being able to build a single static cuda graph for your workload outweigh the benefits of matching.

Another option, of course, is to add a new flag to torch.cond() to make it output torch.zeros() rather than None in the situations I have described, if backwards compatibility is considered important.

@ydwu4 @eellison @bohnstingl 

### Versions

Not important in this case, but same as #141259 

cc @ezyang @gchanan @zou3519 @kadeng @msaroufim @chauhang @penguinwu @ydwu4 @bdhirsh @yf225