Register linalg_lstsq.out for XPU backend
- Dominant language
- Python
- Stars
- 113
- Forks
- 128
- Avg merge
- 5d 13h
- Merged PRs (30d)
- 107
Description
### 🚀 The feature, motivation and pitch
Register and implement **linalg_lstsq.out** on XPU.
### Native function fragment (from native_functions.yaml)
```yaml
- func: linalg_lstsq(Tensor self, Tensor b, float? rcond=None, *, str? driver=None) -> (Tensor solution, Tensor residuals, Tensor rank, Tensor singular_values)
python_module: linalg
variants: function
dispatch:
CompositeExplicitAutograd: linalg_lstsq
tags: dynamic_output_shape
- func: linalg_lstsq.out(Tensor self, Tensor b, float? rcond=None, *, str? driver=None, Tensor(a!) solution, Tensor(b!) residuals, Tensor(c!) rank, Tensor(d!) singular_values) -> (Tensor(a!) solution, Tensor(b!) residuals, Tensor(c!) rank, Tensor(d!) singular_values)
python_module: linalg
variants: function
dispatch:
CPU, CUDA, MPS: linalg_lstsq_out
tags: dynamic_output_shape
```
### Goals
- Register the above native signatures for the XPU backend and add proper dispatch mapping
- Implement the XPU kernel for `linalg_lstsq.out` (or provide an efficient device-side fallback if native kernel is not feasible)
- Ensure unit tests covering
- Driver parameter: Test with different driver options (None, "gels", "gelsy", "gelsd", "gelss") to ensure proper dispatch or graceful error handling on XPU.
- rcond parameter: Validate behavior with rcond=None (default) and explicit rcond values for rank determination.
- Output tensor reuse: Confirm that pre-allocated output tensors (solution, residuals, rank, singular_values) are correctly written in-place without memory issues.
- Ensure performance has no regression compared to CPU
### Alternatives
_No response_
### Additional context
Refer to `oneapi::mkl::lapack::gels` and `oneapi::mkl::lapack::gels_batch`
Contributor guide
Assessment
This issue has not been assessed yet.