Multi-output forward grad codegen is wrong if output values are used

[albanD]

The current codegen allows the user to re-use the output's value in their formula. But when they are multiple outputs, this might already have been changed by the previous formula.

For example eigh is:

- name: linalg_eigh(Tensor self, str UPLO="L") -> (Tensor eigenvalues, Tensor eigenvectors)
  self: eigh_backward(grads, self, /*eigenvectors=*/true, eigenvalues, eigenvectors)
  eigenvalues: eigh_jvp_eigenvalues(self_t, eigenvalues, eigenvectors)
  eigenvectors: eigh_jvp_eigenvectors(self_t, eigenvalues, eigenvectors)

And so the schematic codegen for it is (with some small edits):

  if (_any_has_forward_grad_eigenvalues) {
      // stuff
      auto eigenvalues_new_fw_grad = eigh_jvp_eigenvalues(self_t, eigenvalues, eigenvectors);
      eigenvalues._set_fw_grad(eigenvalues_new_fw_grad, /* level */ 0, /* is_inplace_op */ false);
  }
  if (_any_has_forward_grad_eigenvectors) {
      // stuff
      auto eigenvectors_new_fw_grad = eigh_jvp_eigenvectors(self_t, eigenvalues, eigenvectors);
      eigenvectors._set_fw_grad(eigenvectors_new_fw_grad, /* level */ 0, /* is_inplace_op */ false);
  }

You can see that by the time the eigenvectors formula uses eigenvalues, it already has its fw_grad populated. And so we're going to apply forward AD to eigh_jvp_eigenvectorsv and the output will be a weird Tensor whose forward grad also has a forward grad (which should not be possible!).

Note that having such nested forward grads actually lead to a deadlock on destruction of the level (depending on the order in which they are stored in the ordered set, so this is a non-consistent deadlock).

We should:

• Properly enforce that a forward grad cannot have a forward grad at the same level itself. You should not be able to do make_dual(a, make_dual(b, c)) (note that this one does not usually deadlock due to the order in which they are registered on the level)
• We should fix the codegen to compute all the gradients first and then set all the gradients. This will ensure that we don't compute forward grad within the formula.

cc @ezyang @gchanan @zou3519 @bdhirsh @jbschlosser @albanD @gqchen @pearu @nikitaved @soulitzer @lezcano @Varal7

[albanD]

cc @IvanYashchuk @lezcano this should be fixed at the same time/before we add support for multiple entries in one line for forward mode AD.

[lezcano]

Both @nikitaved and @IvanYashchuk had reported that the CI hangs up when testing AD for some functions. I guess that now we know the why to this!

[albanD]

Yes, some were fixed by #67360
This is going to fix the remaining ones.

[ejguan]

This can be related as @soulitzer suggests: #67463