#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Copyright 1999-2021 Alibaba Group Holding Ltd.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import numpy as np
from numpy.linalg import LinAlgError
from ... import opcodes as OperandDef
from ...serialization.serializables import KeyField, StringField
from ...core import ExecutableTuple
from ..array_utils import device, as_same_device
from ..datasource import tensor as astensor
from ..operands import TensorHasInput, TensorOperandMixin
from ..core import TensorOrder
from .core import SFQR, TSQR
class TensorQR(TensorHasInput, TensorOperandMixin):
_op_type_ = OperandDef.QR
_input = KeyField('input')
_method = StringField('method')
def __init__(self, method=None, **kw):
super().__init__(_method=method, **kw)
@property
def method(self):
return self._method
@property
def output_limit(self):
return 2
def _set_inputs(self, inputs):
super()._set_inputs(inputs)
self._input = self._inputs[0]
def __call__(self, a):
a = astensor(a)
if a.ndim != 2:
raise LinAlgError(f'{a.ndim}-dimensional tensor given. '
'Tensor must be two-dimensional')
tiny_q, tiny_r = np.linalg.qr(np.ones((1, 1), dtype=a.dtype))
x, y = a.shape
q_shape, r_shape = (a.shape, (y, y)) if x > y else ((x, x), a.shape)
q, r = self.new_tensors([a],
kws=[{'side': 'q', 'dtype': tiny_q.dtype,
'shape': q_shape, 'order': TensorOrder.C_ORDER},
{'side': 'r', 'dtype': tiny_r.dtype,
'shape': r_shape, 'order': TensorOrder.C_ORDER}])
return ExecutableTuple([q, r])
@classmethod
def tile(cls, op):
q, r = op.outputs
q_dtype, r_dtype = q.dtype, r.dtype
q_shape, r_shape = q.shape, r.shape
in_tensor = op.input
if in_tensor.chunk_shape == (1, 1):
in_chunk = in_tensor.chunks[0]
chunk_op = op.copy().reset_key()
qr_chunks = chunk_op.new_chunks([in_chunk], kws=[
{'side': 'q', 'shape': q_shape, 'index': in_chunk.index},
{'side': 'r', 'shape': r_shape, 'index': in_chunk.index}
])
q_chunk, r_chunk = qr_chunks
new_op = op.copy()
kws = [
{'chunks': [q_chunk], 'nsplits': ((q_shape[0],), (q_shape[1],)),
'dtype': q_dtype, 'shape': q_shape, 'order': q.order},
{'chunks': [r_chunk], 'nsplits': ((r_shape[0],), (r_shape[1],)),
'dtype': r_dtype, 'shape': r_shape, 'order': r.order}
]
return new_op.new_tensors(op.inputs, kws=kws)
elif op.method == 'tsqr':
return (yield from TSQR.tile(op))
elif op.method == 'sfqr':
return (yield from SFQR.tile(op))
else:
raise NotImplementedError('Only tsqr method supported for now')
@classmethod
def execute(cls, ctx, op):
(a,), device_id, xp = as_same_device(
[ctx[c.key] for c in op.inputs], device=op.device, ret_extra=True)
with device(device_id):
q, r = xp.linalg.qr(a)
qc, rc = op.outputs
ctx[qc.key] = q
ctx[rc.key] = r
[docs]def qr(a, method='tsqr'):
"""
Compute the qr factorization of a matrix.
Factor the matrix `a` as *qr*, where `q` is orthonormal and `r` is
upper-triangular.
Parameters
----------
a : array_like, shape (M, N)
Matrix to be factored.
method: {'tsqr', 'sfqr'}, optional
method to calculate qr factorization, tsqr as default
TSQR is presented in:
A. Benson, D. Gleich, and J. Demmel.
Direct QR factorizations for tall-and-skinny matrices in
MapReduce architectures.
IEEE International Conference on Big Data, 2013.
http://arxiv.org/abs/1301.1071
FSQR is a QR decomposition for fat and short matrix:
A = [A1, A2, A3, ...], A1 may be decomposed as A1 = Q1 * R1,
for A = Q * R, Q = Q1, R = [R1, R2, R3, ...] where A2 = Q1 * R2, A3 = Q1 * R3, ...
Returns
-------
q : Tensor of float or complex, optional
A matrix with orthonormal columns. When mode = 'complete' the
result is an orthogonal/unitary matrix depending on whether or not
a is real/complex. The determinant may be either +/- 1 in that
case.
r : Tensor of float or complex, optional
The upper-triangular matrix.
Raises
------
LinAlgError
If factoring fails.
Notes
-----
For more information on the qr factorization, see for example:
http://en.wikipedia.org/wiki/QR_factorization
Examples
--------
>>> import mars.tensor as mt
>>> a = mt.random.randn(9, 6)
>>> q, r = mt.linalg.qr(a)
>>> mt.allclose(a, mt.dot(q, r)).execute() # a does equal qr
True
"""
op = TensorQR(method=method)
return op(a)