Manually register einsum xla

[pgmoka]

Do manual registration of XLANativeFunctions::einsum for XLA.

This is necessary because currently PyTorch overwrites the key AutogradXLA registration with a its XLA key registration. While ideally we would be able to resolve this problem, this work around resolves the issue from our end. It is also not possible to use full code generation due to #8739.

This manual registration relies on the XLANativeFunctions::einsum function from xla/torch_xla/csrc/aten_xla_type.cpp

[pgmoka]

As the overwrite is written, there is no meaningful unit test we can add as we rely on the generated XLANativeFunctions function. I could do something like what nms_kernel is doing, and refer to tensor_methods::einsum, but this will require rewriting the conditions from

xla/torch_xla/csrc/aten_xla_type.cpp

Line 1531 in 17270e2

at::Tensor XLANativeFunctions::einsum(std::string_view equation,

This cleanest way to do something like that would be to create a utility function shared by both XLANativeFunctions::einsum and our overwrite, and then that might let us test the overwrite. This can be the next step here, or if we think it is unnecessary, we can just do the implementation from this PR.

nms_kernel is itself not being tested directly as far as I can tell. Perhaps this is a separate larger issue we can track separately from the PR.

@tengyifei @ysiraichi: Do y'all have any opinions on this?

[tengyifei]

@pgmoka are you looking for a unit test? A good test IMO is what we wrote in the https://github.com/tengyifei/playground/blob/master/aot-einsum-3.ipynb notebook. We could verify the lowering of einsum in an custom op.

Another test is we should remove the two workarounds referenced in https://github.com/search?q=repo%3Apytorch%2Fxla+8713&type=code, and then the unit test for XLAPatchedLinear should still pass. Because we also check its lowering there.

[pgmoka]

CC: @lsy323

[tengyifei]

Let's add a unit test in Python and also remove the _xla_einsum workaround in this PR (which will also test that this registration worked).

[tengyifei]

Not sure how all these commits got into this branch. Usually I rebase the branch on top of the latest master and then force push. This way I only have a single commit in the PR.

[ysiraichi]

One thing you can do is to check for xla::einsum in the XLA counters. I believe that, before your PR, it wouldn't be in there, since the CompositeImplicitAutograd kernel was called.

[pgmoka]

Not sure how all these commits got into this branch. Usually I rebase the branch on top of the latest master and then force push. This way I only have a single commit in the PR.

I honestly don't know how this happened either. I think I messed something up while fetching the current master to rebase to branch with. I needed to do this to get the latest changes related to _einsum. The final state is what I wanted, but it creates this unfortunate commit map on the PR

[tengyifei]

I honestly don't know how this happened either. I think I messed something up while fetching the current master to rebase to branch with. I needed to do this to get the latest changes related to _einsum. The final state is what I wanted, but it creates this unfortunate commit map on the PR

Gotcha. In that case could you squash the commits from git and reset the commit message so that it hopefully doesn't confuse future readers? Thanks!

[pgmoka]

Too many conflicts. I accidentally merged from master rather than rebasing, and it caused a bunch of issues. My changes are small enough that I will just carry on in a separate PR. I apologize to the reviewers for the noise