torch.linalg in PyTorch 1.10 tracker

[mattip]

This is a tracking issue for torch.linalg tasks for PyTorch 1.10. The goals for the 1.10 release are:

• Reviewing the Python Array AP's linalg extension and pursuing consistency with it
• Modernizing linear algebra functions still under the torch namespace to newer torch.linalg versions
• Continuing to deprecate older linear algebra operators (see Deprecations tracking issue #55953)
• Improving CUDA performance (see Linear algebra GPU backend tracking issue [magma/cusolver/cublas] #47953)
• Improving linear algebra gradient computations

Tasks

• Add the "upper" kwarg to linalg.cholesky Enable upper for torch.linalg.cholesky #62434
• Rationalize the kwargs of linalg.eigh, linalg.eigvalsh, linalg.lstsq, and linalg.pinv with each other and the Python Array API
• Add relative and absolute tolerances for matrix_rank, pinv #54151 (comment)
• Add modern versions of lu, lu_unpack, geqrf, orgqr, ormqr, lu_solve, triangular_solve, cholesky_solve, cholesky_inverse, tensordot, matrix_exp, {svd, pca}_lowrank and lobpcg to torch.linalg
• Move matrix_exp into torch.linalg #61648
• Move tensordot into torch.linalg #61649
• Move {svd,pca}_lowrank into torch.linalg #61650
• Move lobpcg into torch.linalg #61653
• Move lu, lu_solve, and lu_unpack into torch.linalg #61657
• Move torch.cholesky_solve into torch.linalg. #61658
• Investigate a common "solver API" to rationalize the various solvers

cc @ezyang @gchanan @zou3519 @bdhirsh @jbschlosser @anjali411 @jianyuh @nikitaved @pearu @mruberry @heitorschueroff @walterddr @IvanYashchuk @xwang233 @lezcano @rgommers @vincentqb @vishwakftw @ssnl

[rgommers]

xref gh-42053, the correct way to implement aliases is still under discussion there.

[mruberry]

xref gh-42053, the correct way to implement aliases is still under discussion there.

We ultimately decided on the approach taken in #42586. In the future we may implement an "alias" entry in native_functions.yaml, too.

[albanD]

I think an important thing to keep in mind for these new ops is to make sure we don't do the same mistake as the current ones and provide proper documentation about them.
In particular, it should contain enough information for people that know the field to know what happens (for example, high level functions should mention how they are implemented)
Moreover, for less experience people, it would be important to have informations like:

• What are the requirements on the inputs for this function to be valid
• What are the numerical properties of this function in general
• When it exists, propose a more numerically stable alternative for the problematic inputs

Finally, when the backward is not implemented for all cases that the function supports, this should be mentioned as well.

[rgommers]

Regarding the task list here: I'm not too sure about deprecating so many functions. For norm it makes sense, because the torch.linalg.norm function is different from (and arguably nicer) than torch.norm. But for simple aliases, I'd at least do those last, because there's only value once the whole linalg namespace is complete so we can say "torch.linalg matches numpy.linalg". And for some heavily used functions I'd consider not doing the deprecation at all.

[rgommers]

@mruberry, @ezyang and I had a chat about priorities:

1. Implement functions that don’t yet exist, most-used ones first (e.g. kron, eigh - here is some usage data)
2. Implement functions that exist in torch but need a change to match numpy when they go into torch.linalg
3. Add aliases to linalg namespace for functions that exist and already match numpy
4. Add the necessary deprecations to old functions in one go when the torch.linalg namespace is complete (TBD whether they should all be deprecated, see my comment above).

[mruberry]

@rgommers I updated the rollup to reflect these priorities and remove all deprecation tasks (except deprecating torch.norm, as we discussed earlier). I appreciate your point about maybe being less aggressive with deprecations than this initially was. Let's review on a case-by-case basis.

[antocuni]

picking up torch.linalg.svd

[IvanYashchuk]

Picking up kron, eigh, cond, tensorinv.

[antocuni]

I have a design question about torch.linalg.svd, but I think it's general enough that it's good to discuss it here.
numpy.lingalg.svd differs from torch.svd in two major ways:

1. the numpy version has a parameter called full_matrices, while torch.svd has one called some which is basically the opposite, so it's easy to adapt

2. the numpy signature is irregular, and the return type changes depending on the value of compute_uv. In particular, if compute_uv==True it returns a 3-tuple, else it returns a single array. OTOH, torch.svd always return a 3-tuple.

I see that the the existing torch.linalg.* functions all provide a C++ API in the torch::linalg namespace, but point 2 makes it hard/impossible to provide a C++ signature which is consistent with the Python one. I see at least two solutions:

1. implement torch.linalg.svd entirely in Python on top of torch.svd, and not offering the C++ API

2. implement a C++ version of torch::linalg::svd which always returns a 3-tuple, and write a thin Python wrapper which discards the two extra elements in the compute_uv==False case.

Moreover, what are we going to do with torch.svd? Should we:

1. keep it as is and implement torch.linalg.svd on top of it OR
2. move the current implementation to torch.linalg.svd and re-implement the old torch.svd on top of it for backwards compatibility?

[IvanYashchuk]

I'm in favor of this solution:

implement a C++ version of torch::linalg::svd which always returns a 3-tuple, and write a thin Python wrapper which discards the two extra elements in the compute_uv==False case.

I have done the same for ChainerX here.

[rgommers]

Moreover, what are we going to do with torch.svd? Should we:

1. keep it as is and implement torch.linalg.svd on top of it OR
2. move the current implementation to torch.linalg.svd and re-implement the old torch.svd on top of it for backwards compatibility?

It seems to me that (1) is a little less work, but would be much less appealing if torch.svd gets deprecated in the future. That's not decided, but seems fairly likely to happen at some point. So (2) is probably preferred. @mruberry does that make sense to you?

[antocuni]

It seems to me that (1) is a little less work, but would be much less appealing if torch.svd gets deprecated in the future. That's not decided, but seems fairly likely to happen at some point. So (2) is probably preferred. @mruberry does that make sense to you?

I opened PR #45562 to discuss this more concretely. I'll wait for some feedback on it before continuing

[mruberry]

@antocuni Sorry for the delay, working through a backlog now that the 1.7 branch is cut. I'll get to this ASAP (definitely before Monday).

[mruberry]

Moreover, what are we going to do with torch.svd? Should we:

1. keep it as is and implement torch.linalg.svd on top of it OR
2. move the current implementation to torch.linalg.svd and re-implement the old torch.svd on top of it for backwards compatibility?

It seems to me that (1) is a little less work, but would be much less appealing if torch.svd gets deprecated in the future. That's not decided, but seems fairly likely to happen at some point. So (2) is probably preferred. @mruberry does that make sense to you?

Unfortunately I'm not very familiar with the code. I would do the simplest thing to implement torch.linalg.svd(), as discussed on the PR. We can always rearrange the functions internally later.

[mruberry]

@muthuArivoli It's taken awhile, but I was finally able to review linear algebra related issues. You mentioned you were interested in working on one. Would torch.inner(), torch.kron(), fixing torch.einsum, torch.linalg.cond(), or torch.linalg.matrix_rank() be interesting? Or maybe reviewing if there's an alternative to a MAGMA-based implementation of a function like torch.svd() or torch.lstsq?

[mruberry]

Closing in favor of individual issues