Merge pull request torch#373 from bartvm/ormqr

LAPACK ormqr routine

Author: soumith

TensorMath.lua

@@ -1156,6 +1156,22 @@ static void THTensor_random1__(THTensor *self, THGenerator *gen, long b)
                       {name=Tensor},
                       {name=Tensor}}
                   )
+      interface:wrap("ormqr",
+                     cname("ormqr"),
+                     {{name=Tensor, returned=true},
+                      {name=Tensor},
+                      {name=Tensor},
+                      {name=Tensor},
+                      {name='charoption', values={'L', 'R'}, default='L'},
+                      {name='charoption', values={'N', 'T'}, default='N'}},
+                     cname("ormqr"),
+                     {{name=Tensor, default=true, returned=true, invisible=true},
+                      {name=Tensor},
+                      {name=Tensor},
+                      {name=Tensor},
+                      {name='charoption', values={'L', 'R'}, default='L'},
+                      {name='charoption', values={'N', 'T'}, default='N'}}
+                  )
    end
 
    method:register(string.format("m_torch_%sMath__", Tensor))

doc/maths.md

@@ -1953,7 +1953,8 @@ z = x * y
 ### torch.qr([q, r], x) ###
 
 Compute a QR decomposition of the matrix `x`: matrices `q` and `r` such that
-`x = q * r`, with `q` orthogonal and `r` upper triangular.
+`x = q * r`, with `q` orthogonal and `r` upper triangular. This returns
+the thin (reduced) QR factorization.
 
 `=torch.qr(x)` returns the Q and R components as new matrices.
 
@@ -2028,6 +2029,23 @@ given by `torch.geqrf`. See
  [LAPACK documentation](http://www.netlib.org/netlib/lapack/double/dorgqr.f) for
 further details.
 
+<a name="torch.ormqr"></a>
+### torch.ormqr([res], m, tau, mat [, 'L' or 'R'] [, 'N' or 'T']) ###
+
+Multiply a matrix with `Q` as defined by the elementary reflectors and
+scalar factors returned by `geqrf`. This is a low-level function for
+calling LAPACK directly. You'll generally want to use `torch.qr()`
+instead.
+
+* `side` (`'L'` or `'R'`) specifies whether `mat` should be
+  left-multiplied, `Q * mat`, or right-multiplied, `mat * Q`.
+* `trans` (`'N'` or `'T`') specifies whether `Q` should be transposed
+  before being multiplied.
+
+See [LAPACK
+documentation](http://www.netlib.org/netlib/lapack/double/dormqr.f) for
+further details.
+
 <a name="torch.logical.dok"></a>
 ## Logical Operations on Tensors ##
 

lib/TH/generic/THLapack.c

@@ -29,6 +29,8 @@ TH_EXTERNC void sgeqrf_(int *m, int *n, float *a, int *lda, float *tau, float *w
 TH_EXTERNC void dgeqrf_(int *m, int *n, double *a, int *lda, double *tau, double *work, int *lwork, int *info);
 TH_EXTERNC void sorgqr_(int *m, int *n, int *k, float *a, int *lda, float *tau, float *work, int *lwork, int *info);
 TH_EXTERNC void dorgqr_(int *m, int *n, int *k, double *a, int *lda, double *tau, double *work, int *lwork, int *info);
+TH_EXTERNC void sormqr_(char *side, char *trans, int *m, int *n, int *k, float *a, int *lda, float *tau, float *c, int *ldc, float *work, int *lwork, int *info);
+TH_EXTERNC void dormqr_(char *side, char *trans, int *m, int *n, int *k, double *a, int *lda, double *tau, double *c, int *ldc, double *work, int *lwork, int *info);
 
 
 /* Compute the solution to a real system of linear equations  A * X = B */
@@ -218,5 +220,19 @@ void THLapack_(orgqr)(int m, int n, int k, real *a, int lda, real *tau, real *wo
 #endif
 }
 
+/* Multiply Q with a matrix using the output of geqrf */
+void THLapack_(ormqr)(char side, char trans, int m, int n, int k, real *a, int lda, real *tau, real *c, int ldc, real *work, int lwork, int *info)
+{
+#ifdef  USE_LAPACK
+#if defined(TH_REAL_IS_DOUBLE)
+  dormqr_(&side, &trans, &m, &n, &k, a, &lda, tau, c, &ldc, work, &lwork, info);
+#else
+  sormqr_(&side, &trans, &m, &n, &k, a, &lda, tau, c, &ldc, work, &lwork, info);
+#endif
+#else
+  THError("ormqr: Lapack library not found in compile time\n");
+#endif
+}
+
 
 #endif

lib/TH/generic/THLapack.h

@@ -31,5 +31,7 @@ void THLapack_(potrs)(char uplo, int n, int nrhs, real *a, int lda, real *b, int
 void THLapack_(geqrf)(int m, int n, real *a, int lda, real *tau, real *work, int lwork, int *info);
 /* Build Q from output of geqrf */
 void THLapack_(orgqr)(int m, int n, int k, real *a, int lda, real *tau, real *work, int lwork, int *info);
+/* Multiply Q with a matrix from output of geqrf */
+void THLapack_(ormqr)(char side, char trans, int m, int n, int k, real *a, int lda, real *tau, real *c, int ldc, real *work, int lwork, int *info);
 
 #endif

lib/TH/generic/THTensorLapack.c

@@ -637,4 +637,63 @@ void THTensor_(orgqr)(THTensor *ra_, THTensor *a, THTensor *tau)
   THTensor_(free)(work);
 }
 
+/*
+  The ormqr function multiplies Q with another matrix from a sequence of
+  elementary reflectors, such as is produced by the geqrf function.
+
+  Args:
+  * `ra_`   - result Tensor, which will contain the matrix Q' c.
+  * `a`     - input Tensor, which should be a matrix with the directions of the
+              elementary reflectors below the diagonal. If NULL, `ra_` is used as
+              input.
+  * `tau`   - input Tensor, containing the magnitudes of the elementary
+              reflectors.
+  * `c`     - input Tensor, containing the matrix to be multiplied.
+  * `side`  - char, determining whether c is left- or right-multiplied with Q.
+  * `trans` - char, determining whether to transpose Q before multiplying.
+
+  For further details, please see the LAPACK documentation.
+
+*/
+void THTensor_(ormqr)(THTensor *ra_, THTensor *a, THTensor *tau, THTensor *c, const char *side, const char *trans)
+{
+  if (a == NULL) a = ra_;
+  THArgCheck(a->nDimension == 2, 1, "A should be 2 dimensional");
+
+  THTensor *ra__ = NULL;
+  ra__ = THTensor_(cloneColumnMajor)(ra_, c);
+
+  int m = c->size[0];
+  int n = c->size[1];
+  int k = tau->size[0];
+  int lda;
+  if (*side == 'L')
+  {
+    lda = m;
+  }
+  else
+  {
+    lda = n;
+  }
+  int ldc = m;
+
+  /* Dry-run to query the suggested size of the workspace. */
+  int info = 0;
+  real wkopt = 0;
+  THLapack_(ormqr)(side[0], trans[0], m, n, k, THTensor_(data)(a), lda,
+                   THTensor_(data)(tau), THTensor_(data)(ra__), ldc,
+                   &wkopt, -1, &info);
+
+  /* Allocate the workspace and call LAPACK to do the real work. */
+  int lwork = (int)wkopt;
+  THTensor *work = THTensor_(newWithSize1d)(lwork);
+  THLapack_(ormqr)(side[0], trans[0], m, n, k, THTensor_(data)(a), lda,
+                   THTensor_(data)(tau), THTensor_(data)(ra__), ldc,
+                   THTensor_(data)(work), lwork, &info);
+
+  THLapackCheck(" Lapack Error %s : unknown Lapack error. info = %i", "ormqr", info);
+  THTensor_(freeCopyTo)(ra__, ra_);
+  THTensor_(free)(work);
+}
+
 #endif

lib/TH/generic/THTensorLapack.h

@@ -15,5 +15,6 @@ TH_API void THTensor_(potrf)(THTensor *ra_, THTensor *a);
 TH_API void THTensor_(qr)(THTensor *rq_, THTensor *rr_, THTensor *a);
 TH_API void THTensor_(geqrf)(THTensor *ra_, THTensor *rtau_, THTensor *a);
 TH_API void THTensor_(orgqr)(THTensor *ra_, THTensor *a, THTensor *tau);
+TH_API void THTensor_(ormqr)(THTensor *ra_, THTensor *a, THTensor *tau, THTensor *c, const char *side, const char *trans);
 
 #endif

test/test_qr.lua

@@ -60,6 +60,21 @@ local function qrManual(geqrfFunc, orgqrFunc)
   end
 end
 
+-- Check that Q multiplied with a matrix with ormqr gives the correct result
+local function checkQM(testOpts, mat1, mat2)
+  local q, r = torch.qr(mat1)
+  local m, tau = torch.geqrf(mat1)
+  local requiredPrecision = 1e-5
+  tester:assertTensorEq(torch.mm(q, mat2), torch.ormqr(m, tau, mat2),
+                        requiredPrecision)
+  tester:assertTensorEq(torch.mm(mat2, q), torch.ormqr(m, tau, mat2, 'R'),
+                        requiredPrecision)
+  tester:assertTensorEq(torch.mm(q:t(), mat2),
+                        torch.ormqr(m, tau, mat2, 'L', 'T'), requiredPrecision)
+  tester:assertTensorEq(torch.mm(mat2, q:t()),
+                        torch.ormqr(m, tau, mat2, 'R', 'T'), requiredPrecision)
+end
+
 -- Check that the given `q`, `r` matrices are a valid QR decomposition of `a`.
 local function checkQR(testOpts, a, q, r)
   local qrFunc = testOpts.qr
@@ -250,5 +265,10 @@ addTestVariations(tests, 'randomNonContiguous', function(testOpts)
   end
 end)
 
+function tests.testQM()
+  checkQM({}, torch.randn(10, 10), torch.randn(10, 10))
+  -- checkQM({}, torch.randn(20, 10), torch.randn(20, 20))
+end
+
 tester:add(tests)
 tester:run()