text: #!/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
  
  
  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)
  