Add complex tensor with subclassing

[pierreguilmin]

Pair-programming with @ezyang at the PyTorch Conference 2023 for a WIP implementation of complex tensors working with torch.compile.

The implementation is inspired from https://github.com/pytorch/pytorch/blob/main/torch/testing/_internal/two_tensor.py.

A few todos left, notably a custom autograd for the constructor.

This was tested with the nightly build 2.2.0.dev20231016.

> TORCH_LOGS=+aot python complex_tensor.py
FakeTensor(..., size=(1, 1), dtype=torch.int64) FakeTensor(..., size=(1, 1), dtype=torch.int64)
FunctionalTensor(_to_functional_tensor(FakeTensor(..., size=(1, 1), dtype=torch.int64))) FunctionalTensor(_to_functional_tensor(FakeTensor(..., size=(1, 1), dtype=torch.int64)))
FunctionalTensor(_to_functional_tensor(FakeTensor(..., size=(1, 1), dtype=torch.int64))) FunctionalTensor(_to_functional_tensor(FakeTensor(..., size=(1, 1), dtype=torch.int64)))
FunctionalTensor(_to_functional_tensor(FakeTensor(..., size=(1, 1), dtype=torch.int64))) FunctionalTensor(_to_functional_tensor(FakeTensor(..., size=(1, 1), dtype=torch.int64)))
[2023-10-17 11:32:12,381] [0/0] torch._functorch.aot_autograd.__aot_graphs: [INFO] TRACED GRAPH
[2023-10-17 11:32:12,381] [0/0] torch._functorch.aot_autograd.__aot_graphs: [INFO]  ===== Forward graph 0 =====
[2023-10-17 11:32:12,381] [0/0] torch._functorch.aot_autograd.__aot_graphs: [INFO]  <eval_with_key>.4 from /Users/pierreguilmin/.pyenv/versions/live-compile-complex/lib/python3.9/site-packages/torch/fx/experimental/proxy_tensor.py:506 in wrapped class <lambda>(torch.nn.Module):
[2023-10-17 11:32:12,381] [0/0] torch._functorch.aot_autograd.__aot_graphs: [INFO]     def forward(self, arg0_1: i64[1, 1], arg1_1: i64[1, 1], arg2_1: i64[1, 1], arg3_1: i64[1, 1]):
[2023-10-17 11:32:12,381] [0/0] torch._functorch.aot_autograd.__aot_graphs: [INFO]         # File: /Users/pierreguilmin/Desktop/live-compile-complex/foo.py:53, code: return x @ y
[2023-10-17 11:32:12,381] [0/0] torch._functorch.aot_autograd.__aot_graphs: [INFO]         mul: i64[1, 1] = torch.ops.aten.mul.Tensor(arg0_1, arg2_1)
[2023-10-17 11:32:12,381] [0/0] torch._functorch.aot_autograd.__aot_graphs: [INFO]         mul_1: i64[1, 1] = torch.ops.aten.mul.Tensor(arg1_1, arg3_1)
[2023-10-17 11:32:12,381] [0/0] torch._functorch.aot_autograd.__aot_graphs: [INFO]         sub: i64[1, 1] = torch.ops.aten.sub.Tensor(mul, mul_1);  mul = mul_1 = None
[2023-10-17 11:32:12,381] [0/0] torch._functorch.aot_autograd.__aot_graphs: [INFO]         mul_2: i64[1, 1] = torch.ops.aten.mul.Tensor(arg0_1, arg3_1);  arg0_1 = arg3_1 = None
[2023-10-17 11:32:12,381] [0/0] torch._functorch.aot_autograd.__aot_graphs: [INFO]         mul_3: i64[1, 1] = torch.ops.aten.mul.Tensor(arg1_1, arg2_1);  arg1_1 = arg2_1 = None
[2023-10-17 11:32:12,381] [0/0] torch._functorch.aot_autograd.__aot_graphs: [INFO]         add: i64[1, 1] = torch.ops.aten.add.Tensor(mul_2, mul_3);  mul_2 = mul_3 = None
[2023-10-17 11:32:12,381] [0/0] torch._functorch.aot_autograd.__aot_graphs: [INFO]         return [sub, add]
[2023-10-17 11:32:12,381] [0/0] torch._functorch.aot_autograd.__aot_graphs: [INFO]
[2023-10-17 11:32:12,381] [0/0] torch._functorch.aot_autograd.__aot_graphs: [INFO]
ComplexTensor(real=tensor([[-5]]), imag=tensor([[10]]))

[albanD]

Looks really cool!
Do you want to fix the autograd support before I merge this?

[albanD]

These should be @ right?

[pierreguilmin]

Yes! 🙇🏼‍♂️

[albanD]

The best way to add autograd support here is to do a parallel to Tensor (which is never differentiable.
So I would recommend that ComplexTensor(...) is never differentiable and you have a create_complex_tensor(...) which is differentiable and built with a custom autograd Function (where you create a ComplexTensor during the fw).

[gautierronan]

Suggested change

	dtype=torch.complex64, # todo: real to complex dtype
	dtype=re.dtype.to_complex(),

since v2.1

[ezyang]

@albanD let's just merge and fix it up on main

[pierreguilmin]

Sorry I had a very busy week, I'll let you take it from here. 😉 (except if you need anything from our side)

Btw, @gautierronan is one of my colleagues, he works with me on the dynamiqs library to simulate quantum systems with PyTorch.

[ezyang]

@pierreguilmin / @gautierronan if the two of you are interested in pushing this subclass forward, I'd recommend opening up a little repo with just this class and starting to chuck stuff into it. The subclass zoo here is just to so "it's possible", it's not a good permanent home for a feature that people want to use.

[pierreguilmin]

Thanks for your advice @ezyang. Could you also advise on the next steps to make progress on this? What do you mean by "starting to chuck stuff into it", just implement more operators?

[ezyang]

Yup!