QR decomposition

[anagaev]

Pull Request Template

Checklist

• Confirmed that cargo run-checks command has been executed.
• Made sure the book is up to date with changes in this PR.

Related Issues/PRs

Linear algebra operations: #1538

Changes

• Add QR decomposition (Householder reflection) with reduced form support in crates/burn-tensor/src/tensor/linalg/.
• Add numerical and generic parameters (copied from LU decomposition tests) tests.
• Removed QR decomposition (Gram–Schmidt process) from crates/burn-core/src/module/initializer.rs

Testing

Describe how these changes have been tested.

• cargo run-checks
• Multiple numerical tests with predefined and random matrices. The tests include square, tall, wide, singular matrices, batches of matrices, and several special cases. The tests check the tolerance between the original tensor and the product of its QR decomposition.
• Tests for panic situations.
• Dimensions/shape tests for reduced=true

P.S. It may be worth adding tests for the orthonormality of Q and the upper-triangular property of R.

[anagaev]

It seems there are some problems with fusion on the metal and wgpu backends. I was able to reproduce the issue locally using the wgpu backend.

[anagaev]

I reproduced the issue locally with wgpu pipeline. It seems to me, there are 2 parts of the issue.

1. fusion feature sometimes causes computational errors: some elements of the matrix multiplication result are identical to elements from a previous matrix multiplication result.
   For example, the code below and the code from the pr have computational problem with some tests:

        ...
        let ar = sub_tensor.clone().transpose().matmul(w.clone());
        let ar1 = ar.mul(tau.clone()); // no clone of ar here
        let a = w.clone().matmul(ar1.transpose());
        r = r.slice_assign(&slices, sub_tensor.clone() - a);
        slices[D - 2] = Slice::full();
        let q_sub_tensor = q.clone().slice(&slices);
        let qr = q_sub_tensor.clone().matmul(w.clone());
        let qr1 = qr.mul(tau.clone()); // no clone of qr here
        let b = w.clone().matmul(qr1.transpose());
        q = q.slice_assign(&slices, q_sub_tensor.clone() - b.transpose());

But this one works correctly

        ...
        let ar = sub_tensor.clone().transpose().matmul(w.clone());
        let ar1 = ar.clone().mul(tau.clone()); // clone of ar here
        let a = w.clone().matmul(ar1.transpose());
        r = r.slice_assign(&slices, sub_tensor.clone() - a);
        slices[D - 2] = Slice::full();
        let q_sub_tensor = q.clone().slice(&slices);
        let qr = q_sub_tensor.clone().matmul(w.clone());
        let qr1 = qr.clone().mul(tau.clone());  //clone of qr here
        let b = w.clone().matmul(qr1.transpose());
        q = q.slice_assign(&slices, q_sub_tensor.clone() - b.transpose());

This problem appears only with wgpu backend with fusion feature. No problem for flex backend with and no fusion feature and wgpu backend without fusion.
Maybe I don't properly understand some concept of .clone(), but it looks like a bug in burn-fusion

2. Some tests with bigger matrices failed. For example,

test tensor::float::linalg::qr::test_qr_medium_wide ... FAILED
test tensor::float::linalg::qr::test_qr_300x500 ... FAILED
test tensor::float::linalg::qr::test_qr_500x300 ... FAILED
test tensor::float::linalg::qr::test_qr_500x500 ... ok

The part of the logs below:

thread 'DSU-4-0' (341314) panicked at .../.cargo/git/checkouts/cubecl-058c47895211d464/3beb9af/crates/cubecl-runtime/src/tune/local.rs:152:18:
Should run when selected by autotune.: gemm_None: An unknown error happened.
Unable to launch matmul because the config is invalid: "Dimension k=112 must be a multiple of plane size 64"

So, it's related to cubecl