intel / intel/torch-xpu-ops

Register linalg_lstsq.out for XPU backend

Open
#4,154 0 comments 0 reactions 1 assignee Claimed by @jkosnox View on GitHub
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.