Add diagflat complex support

[rjkilpatrick]

Adds complex numbers support for torch.diag

>>> import torch
>>> a = torch.ones(2, dtype=torch.complex128)
>>> torch.diagflat(a)
tensor([[1.+0.j, 0.+0.j],
        [0.+0.j, 1.+0.j]], dtype=torch.complex128)
>>> b = a.cuda()
>>> torch.diagflat(b)
tensor([[1.+0.j, 0.+0.j],
        [0.+0.j, 1.+0.j]], device='cuda:0', dtype=torch.complex128)

Note that automatic differentiation isn't implemented:

>>> d = torch.ones(1, dtype=torch.complex128, requires_grad=True)
>>> torch.diagflat(d)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
RuntimeError: diag does not support automatic differentiation for outputs with complex dtype.

Fixes #47499

[dr-ci]

💊 CI failures summary and remediations

As of commit 5917c83 (more details on the Dr. CI page):

---

✅ None of the CI failures appear to be your fault 💚

---

• 1/1 broken upstream at merge base f23a2a1 since Nov 06

---

🚧 1 ongoing upstream failure:

These were probably caused by upstream breakages that are not fixed yet:

• binary_linux_libtorch_3_7m_cpu_devtoolset7_shared-with-deps_build since Nov 06
  • 🔁 rerun

---

This comment was automatically generated by Dr. CI (expand for details).

Follow this link to opt-out of these comments for your Pull Requests.

Please report bugs/suggestions on the GitHub issue tracker or post in the (internal) Dr. CI Users group.

See how this bot performed.

This comment has been revised 1 time.

[anjali411]

@rjkilpatrick is this PR ready for review?

[rjkilpatrick]

@anjali411 this issue is ready to be reviewed.

The complex autograd should be a separate issue.

[heitorschueroff]

@albanD can you please advise whether we can leave autograd support for a later PR?

[anjali411]

torch.diagflat calls torch.diag so it's ok to not add autograd support for torch.diagflat in this PR. let's add that in a follow-up PR

[anjali411]

lgtm

[facebook-github-bot]

@anjali411 has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

[facebook-github-bot]

@anjali411 merged this pull request in 8aca85d.

[albanD]

@anjali411 the autograd should be disabled anyways as it is not in the whitelist right?

[anjali411]

@albanD diagflat calls diag whose backward uses diagonal which is supported for complex