Source code for mars.tensor.linalg.qr

#!/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
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# 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)

    def method(self):
        return self._method

    def output_limit(self):
        return 2

    def _set_inputs(self, 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(
                    "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])

    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(
                    {"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))
            raise NotImplementedError("Only tsqr method supported for now")

    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. 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: Examples -------- >>> import mars.tensor as mt >>> a = mt.random.randn(9, 6) >>> q, r = mt.linalg.qr(a) >>> mt.allclose(a,, r)).execute() # a does equal qr True """ op = TensorQR(method=method) return op(a)