API for variable number of returns in linalg

[rgommers]

Feedback from @mruberry: it may be nice to support always returning the same type from functions like Cholesky, with tensors accessed by names. That way frameworks could return additional tensors and there wouldn't be BC-compatibility issues if more tensor returns are added later. This suggestion, however, comes from a place of "it'd be nice if people wanted to use this API", as opposed to the perspective of, "this API is a lowest common denominator only intended for library writers."

cholesky always return a single array, so I thought it'd be fine - but the comment may be related to pytorch/pytorch#47608 (desire to return an error code rather than raise a RuntimeError for failures). For qr, svd and perhaps some more functions there's the related issue of different returns based on keywords. Using a class to stuff all return values in is a common design pattern for, e.g., scipy.optimize. I'm not convinced it's a good idea for the linalg functions here, but worth considering at least.

[asmeurer]

I don't understand the technical details of that pytorch issue. Does it mean that any API that raises an exception is potentially problematic?

[mruberry]

I don't understand the technical details of that pytorch issue. Does it mean that any API that raises an exception is potentially problematic?

Yes, because checking for the exception will cause a cross-device synchronization, which PyTorch tries very hard to avoid. In PyTorch we also have cases where error checking occurs exclusively on the CPU to avoid this issue.

We intend to keep np.linalg-like error checking behavior (with a warning that it causes cross-device synchronization) but also offer _ex variants that avoid it. For example, torch.linalg.cholesky and torch.linalg.cholesky_ex, the latter of which performs no error checking but returns the info required for error checking as an array (tensor) on the device. This lets the user perform their own error checking (or not).

This thought was also motivated by functions that return different number of tensors depending on their input args. For example, in NumPy linalg.svd can return a variable number of tensors. There are various ways this could be designed, but I would like to suggest that functions always return the same type, and it'd be nice if they were designed to be forward-compatible. For example, to accommodate returning yet another function from linalg.svd would be easy if the object returned from svd only allowed access to its tensor elements by name. Another option would be to always return a tuple that contains only the requested tensors, but I think this is a confusing and potentially dangerous UX.

[rgommers]

Now that there are a number of PRs up for linear algebra functions, it looks to me like the issues with variable number of returned values, order of returns, and transposes within the array even for matches shapes are pretty painful:

• svd: NumPy and PyTorch order the returned tuple (u, s, v), while TF orders as (s, u, v) (see Add SVD function specification #114).
• svd:TensorFlow returns both the left and right singular vectors in columns. NumPy returns the left singular vectors in rows.
• lstsq: what gets returned is wildly different, Add specification for returning the least-squares solution to a linear matrix equation (linalg: lstsq) #119 lists six variations.
• pinv: mostly okay for array/tensor libraries, but SciPy returns rank in addition to the pseudo-inverse (see Add specification for computing the pseudo-inverse (linalg: pinv) #118)

Some of the linalg functions that are still to be done (e.g., qr) also won't be pretty.

For example, to accommodate returning yet another function from linalg.svd would be easy if the object returned from svd only allowed access to its tensor elements by name.

This would have some advantages:

1. We would not have to worry about differences in order of returns
2. It would allow adding more return values in the future
3. It could avoid issues with wanting different behaviour on accelerators, like the _ex variants that don't do error checking. This could still be the same function then, with an extra keyword - nicer than extra functions.

Disadvantages:

1. It would be a bigger deviation from existing APIs
2. It's harder to support with type annotations
3. The API would be weird for functions returning a single array.

And what it doesn't solve:

1. Content of a particular returned array, like row vs. column layout as for svd

Overall, I'd still lean towards not returning custom objects. The future extensibility is be the most compelling reason to consider it though.

Another option would be to always return a tuple that contains only the requested tensors, but I think this is a confusing and potentially dangerous UX.

I agree, not a fan of this option.

[rgommers]

Didn't think of this before, but at least there's now a natural Python way for the "return custom objects" option - using dataclasses. All you'd need is:

@dataclass
class SVDResult:  # This could inherit from a base LinalgResult for typing/consistency perhaps
    s: array
    u: array
    v: array

Usage pattern if you just want a single returned array in the rest of your code:

s = svd(x).s

and if you want them all, just carry around the object, or unpack:

x_svd = svd(x)
s, u, v = x_svd.s, x_svd.u, x_svd.v

That is more verbose, but there's an upside to that as well. For example for svd, there would be libraries where you'd otherwise have

s, u, v = main_namespace.svd(x)
u, s, v = arrayapi_namespace.svd(x)

which is probably more confusing than:

s, u, v = main_namespace.svd(x)
x_svd = arrayapi_namespace.svd(x)
s, u, v = x_svd.s, x_svd.u, x_svd.v

[rgommers]

Didn't think of this before, but at least there's now a natural Python way for the "return custom objects" option - using dataclasses.

Which I think y'all decided was a sane idea @mruberry: pytorch/pytorch#46187.

[rgommers]

Comment copied from gh-18:

Another alternative is to simply return a dictionary.

That's a bit nicer static typing wise, but clunkier API wise (result.x is nicer than result['x'], and the __repr__ is worse for dicts).

[rgommers]

I re-read pytorch/pytorch#46187, and PyTorch deviated on purpose from NumPy in the new design by always return a tuple with the max number of elements if there's a boolean keyword to control number of returns in NumPy (e.g., svd(x, compute_uv=False) return an array in NumPy and a tuple (tensor([]), S, tensor([]) in PyTorch. The rationale is that this is much easier for the JIT; doing overloads on the boolean keyword is a pain.

The conclusion from that issue was that indeed dataclasses would be desirable long-term.

[rgommers]

Still discussing, PyTorch may change the above to the NumPy-compatible behaviour.

[kgryte]

@rgommers I believe we've addressed this issue in recently merged PRs. Notably, only return namedtuples when namedtuples are necessary (i.e., multiple return arrays); otherwise, just return an array. If libraries want to return other info (e.g., error info), they can expose a different API.

So, I believe we can close this.

[rgommers]

Yep, yay for svdvals:)

There's only one function left with a return_xxx or compute_xxx kind of keyword:

def unique(x, /, *, return_counts=False, return_index=False, return_inverse=False):

but there's little we can do about that one.

Thanks for the suggestion @mruberry!