// cudamatrix/cu-matrix.h
  
  // Copyright 2009-2012  Karel Vesely
  //                2013  Johns Hopkins University (author: Daniel Povey)
  //                2013  Hainan Xu
  //                2013  Xiaohui Zhang
  //           2013-2015  Guoguo Chen
  //                2017  Shiyin Kang
  //                2019  Yiwen Shao
  
  // See ../../COPYING for clarification regarding multiple authors
  //
  // 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
  //
  // THIS CODE IS PROVIDED *AS IS* BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
  // KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WITHOUT LIMITATION ANY IMPLIED
  // WARRANTIES OR CONDITIONS OF TITLE, FITNESS FOR A PARTICULAR PURPOSE,
  // MERCHANTABLITY OR NON-INFRINGEMENT.
  // See the Apache 2 License for the specific language governing permissions and
  // limitations under the License.
  
  
  
  #ifndef KALDI_CUDAMATRIX_CU_MATRIX_H_
  #define KALDI_CUDAMATRIX_CU_MATRIX_H_
  
  #include <sstream>
  #include <vector>
  
  #include "cudamatrix/cu-matrixdim.h"
  #include "cudamatrix/cu-common.h"
  #include "cudamatrix/cu-value.h"
  #include "matrix/matrix-common.h"
  #include "matrix/kaldi-matrix.h"
  #include "cudamatrix/cu-array.h"
  #include "cudamatrix/cu-math.h"
  #include "cudamatrix/cu-rand.h"
  #include "cudamatrix/cu-sparse-matrix.h"
  
  namespace kaldi {
  
  template<typename Real>
  Real TraceMatMat(const CuMatrixBase<Real> &A, const CuMatrixBase<Real> &B,
                   MatrixTransposeType trans = kNoTrans);
  
  /// Does multiple matrix multiplications, executing them in parallel using
  /// cuBLAS's gemmBatched if we are using a GPU. Vectors A, B and C must have
  /// the same length; for each i, this function executes the matrix operation
  /// C[i] = alpha *  A[i](^T)*B[i](^T) + beta * C[i].
  template<typename Real>
  void AddMatMatBatched(const Real alpha, std::vector<CuSubMatrix<Real>* > &C,
                        const std::vector<CuSubMatrix<Real>* > &A,
                        MatrixTransposeType transA,
                        const std::vector<CuSubMatrix<Real>* > &B,
                        MatrixTransposeType transB,
                        const Real beta);
  
  /**
   * Matrix for CUDA computing.
   * Does the computation on the CUDA card when CUDA is compiled in and
   * we have a suitable GPU (CuDevice::Instantiate().Enabled() == true);
   * otherwise, does it on the CPU.
   */
  
  /*
  template<typename Real>
  struct MatrixElement {
    int row;
    int column;
    Real weight;
  };
  // */
  
  template<typename Real>
  class CuMatrixBase {
   public:
    friend class CuMatrixBase<float>;
    friend class CuMatrixBase<double>;
    friend class CuVectorBase<float>;
    friend class CuVectorBase<double>;
    friend class VectorBase<Real>;
    friend class CuSpMatrix<Real>;
    friend class CuTpMatrix<float>;
    friend class CuTpMatrix<double>;
    friend class CuVectorBase<Real>;
    friend class CuSubMatrix<Real>;
    friend class CuRand<Real>;
    friend class CuSubVector<Real>;
    friend class CuBlockMatrix<Real>;
    friend class CuSparseMatrix<float>;
    friend class CuSparseMatrix<double>;
    friend class CuSparseMatrix<Real>;
  
    /// Copies column r from column indexes[r] of src.
    /// As a special case, if indexes[i] == -1, sets column i to zero
    /// indexes.size() must equal this->NumCols(),
    /// and src.NumRows() must equal this.NumRows()
    void CopyCols(const CuMatrixBase<Real> &src,
                  const CuArrayBase<MatrixIndexT> &indexes);
  
  
    /// Add column indices[r] of src to column r.
    /// As a special case, if indexes[i] == -1, skip column i
    /// indices.size() must equal this->NumCols(),
    /// and src.NumRows() must equal this.NumRows()
    void AddCols(const CuMatrixBase<Real> &src,
                 const CuArrayBase<MatrixIndexT> &indices);
  
    /// Copies row r from row indexes[r] of src.
    /// As a special case, if indexes[i] < 0, sets row i to zero.
    /// src.NumCols() must equal this.NumCols()
    void CopyRows(const CuMatrixBase<Real> &src,
                  const CuArrayBase<MatrixIndexT> &indexes);
  
    /// Copies row r of this matrix from an array of floats at the location given
    /// by src[r], where src[r] is assumed to be obtained from the RowData()
    /// function of another CuMatrix, or from CuVector::Data() (the point is: the
    /// data it points to should be on the GPU if we're using a GPU, and on a CPU
    /// otherwise).  src.size() must equal this.NumRows(), and if any src[r] is
    /// NULL then this.Row(r) will be set to zero.
    void CopyRows(const CuArrayBase<const Real*> &src);
  
    /// For each row r of this matrix, copies it to the array of floats at the
    /// location given by dst[r], where dst[r] is assumed to be obtained from the
    /// RowData() function of another CuMatrix, or from CuVector::Data() (i.e. it
    /// should point to memory on the GPU if we're using a GPU, or on the CPU
    /// otherwise).  If dst[r] is NULL, does not copy anywhere.  Requires that
    /// none of the memory regions pointed to by the pointers in "dst" overlap
    /// (e.g. none of the pointers should be the same).
    void CopyToRows(const CuArrayBase<Real*> &dst) const;
  
    /// Does for each row r, this.Row(r) += alpha * src.row(indexes[r]).
    /// If indexes[r] < 0, does not add anything.
    /// src.NumCols() must equal this.NumCols()
    void AddRows(Real alpha,
                 const CuMatrixBase<Real> &src,
                 const CuArrayBase<MatrixIndexT> &indexes);
  
  
    /// Does for each row r, this.Row(r) *= alpha * src.row(indexes[r]),
    /// where '*=' is elementwise multiplication.
    /// If indexes[r] < 0, does not add anything.
    /// src.NumCols() must equal this.NumCols()
    void MulRows(const CuMatrixBase<Real> &src,
                 const CuArrayBase<MatrixIndexT> &indexes);
  
  
    /// Does for each row r, this.Row(r) += alpha * src[r],
    /// treating src[r] as the beginning of a region of memory representing
    /// a vector of floats, of the same length as this.NumCols().
    void AddRows(Real alpha,
                 const CuArrayBase<const Real*> &src);
  
  
    /// For each row i of *this, adds this->Row(i) to
    /// dst->Row(indexes(i)) if indexes(i) >= 0, else do nothing.
    /// Requires that all the indexes[i] that are >= 0
    /// be distinct, otherwise the behavior is undefined.
    void AddToRows(Real alpha,
                   const CuArrayBase<MatrixIndexT> &indexes,
                   CuMatrixBase<Real> *dst) const;
  
  
    /// For each row r of this matrix, adds it (times alpha) to the array of
    /// floats at the location given by dst[r], where dst[r] is assumed to be
    /// obtained from the RowData() function of another CuMatrix, or from
    /// CuVector::Data() (i.e. it should point to memory on the GPU if we're using
    /// a GPU, or on the CPU otherwise).  If dst[r] is NULL, does not do anything
    /// for that row.  Requires that none of the memory regions pointed to by the
    /// pointers in "dst" overlap (e.g. none of the pointers should be the same).
    void AddToRows(Real alpha, const CuArrayBase<Real*> &dst) const;
  
  
    /// For each row r of this and for each column c, sets (*this)(r, c) to the
    /// sum \sum_j src(r, j), where j ranges from indexes[c].first through
    /// indexes[c].second - 1.
    void SumColumnRanges(const CuMatrixBase<Real> &src,
                         const CuArrayBase<Int32Pair> &indexes);
  
  
    /// For each row r of this and for each column c, do
    /// (*this)(r, c) += \sum_j src(j, c),
    /// where j ranges from indexes[r].first through indexes[r].second - 1.
    /// In general indexes must be >= 0 and < src.NumRows(); but to represent an empty range
    /// you may use the pair (-1, -1) or any pair of numbers (i, j) such that i >= j.
    void AddRowRanges(const CuMatrixBase<Real> &src,
                      const CuArrayBase<Int32Pair> &indexes);
  
  
    friend Real TraceMatMat<Real>(const CuMatrixBase<Real> &A,
                                  const CuMatrixBase<Real> &B,
                                  MatrixTransposeType trans);
  
    friend Real TraceMatSmat<Real>(const CuMatrixBase<Real> &A,
                                   const CuSparseMatrix<Real> &B,
                                   MatrixTransposeType trans);
  
    friend void AddMatMatBatched<Real>(const Real alpha,
                                       std::vector<CuSubMatrix<Real>* > &C,
                                       const std::vector<CuSubMatrix<Real>* > &A,
                                       MatrixTransposeType transA,
                                       const std::vector<CuSubMatrix<Real>* > &B,
                                       MatrixTransposeType transB,
                                       const Real beta);
  
    /// Adds "value" to the diagonal elements of the matrix.  The matrix
    /// *this does not have to be square.
    void AddToDiag(Real value);
  
    /// Dimensions
    MatrixIndexT NumRows() const { return num_rows_;  }
    MatrixIndexT NumCols() const { return num_cols_;  }
    MatrixIndexT Stride() const { return stride_; }
  
    // MatrixDim is a struct containing "rows", "cols" and "stride",
    // that is an argument of most CUDA kernels.
    ::MatrixDim Dim() const {
      ::MatrixDim d = { num_rows_, num_cols_, stride_ };
      return d;
    }
  
    Real FrobeniusNorm() const { return sqrt(TraceMatMat(*this, *this, kTrans)); }
  
    bool IsUnit(Real tol = 0.001) const;
  
    /// True if ((*this)-other).FrobeniusNorm() <= tol * this->FrobeniusNorm()
    bool ApproxEqual(const CuMatrixBase<Real> &other, float tol = 0.01) const;
  
    /// Get size of matrix in bytes
    MatrixIndexT SizeInBytes() const { return num_rows_*stride_*sizeof(Real); }
  
    // Copy functions.  These do not resize.
    template<typename OtherReal>
    void CopyFromMat(const MatrixBase<OtherReal> &src,
                     MatrixTransposeType trans = kNoTrans);
  
  
    void CopyFromGeneralMat(const GeneralMatrix &src,
                            MatrixTransposeType trans = kNoTrans);
  
    void CopyFromMat(const MatrixBase<Real> &src,
                     MatrixTransposeType trans = kNoTrans);
  
    void CopyFromSp(const CuSpMatrix<Real> &M);
  
    template<typename OtherReal>
    void CopyFromTp(const CuTpMatrix<OtherReal> &M,
                    MatrixTransposeType trans = kNoTrans);
    
    // This function will copy from source rows (start_range, end_range]
    // if the range is outside of the clamped region then the clamped
    // row will be replicated across the out of range areas
    void CopyRangeFromMatClamped(const CuMatrixBase<Real> & src,
        int32_t start_range, int32_t end_range,
        int32_t clamp_low, int32_t clamp_high);
  
    template<typename OtherReal>
    void CopyFromMat(const CuMatrixBase<OtherReal> &M,
                     MatrixTransposeType trans = kNoTrans);
  
    template<typename OtherReal>
    void CopyToMat(MatrixBase<OtherReal> *dst,
                   MatrixTransposeType trans = kNoTrans) const;
  
    /// This function has two modes of operation.  If v.Dim() == NumRows() *
    /// NumCols(), then treats the vector as a row-by-row concatenation of a
    /// matrix and copies to *this.
    /// if v.Dim() == NumCols(), it sets each row of *this to a copy of v.
    void CopyRowsFromVec(const CuVectorBase<Real> &v);
  
    /// Version of CopyRowsFromVec() that takes a CPU-based vector.
    void CopyRowsFromVec(const VectorBase<Real> &v);
  
    /// Copies vector into matrix, column-by-column.
    /// Note that rv.Dim() must either equal NumRows()*NumCols() or NumRows();
    /// this has two modes of operation.
    void CopyColsFromVec(const CuVectorBase<Real> &v);
  
    /// Copy vector into specific column of matrix.
    void CopyColFromVec(const CuVectorBase<Real> &v, const MatrixIndexT col);
  
    /// Set each element to the sigmoid of the corresponding element of "src":
    /// element by element, x = 1 / (1 + exp(-x))
    void Sigmoid(const CuMatrixBase<Real> &src);
  
    /// Set each element to the Heaviside function of the corresponding element
    /// of "src", which we define as the function (x > 0 ? 1.0 : 0.0) [note:
    /// in general, there are different ways to deal with the situation when x==0.]
    void Heaviside(const CuMatrixBase<Real> &src);
  
    void Exp(const CuMatrixBase<Real> &src);
  
    void Log(const CuMatrixBase<Real> &src);
  
    void Pow(const CuMatrixBase<Real> &src, Real power);
  
    /// Apply power to the absolute value of each element.
    /// If include_sign is true, the result will be multiplied with
    /// the sign of the input value.
    /// If the power is negative and the input to the power is zero,
    /// The output will be set zero. If include_sign is true, it will
    /// multiply the result by the sign of the input.
    void PowAbs(const CuMatrixBase<Real> &src, Real power, bool include_sign=false);
  
    void Floor(const CuMatrixBase<Real> &src, Real floor_val);
    
    void Ceiling(const CuMatrixBase<Real> &src, Real ceiling_val);
    
    /// This is equivalent to running:
    /// Floor(src, lower_limit);
    /// Ceiling(src, upper_limit);
    /// Exp(src)
    void ExpLimited(const CuMatrixBase<Real> &src, Real lower_limit, Real upper_limit);
  
    /// For each element x of the matrix, set it to
    /// (x < 0 ? exp(x) : x + 1).  This function is used
    /// in our RNNLM training.
    void ExpSpecial(const CuMatrixBase<Real> &src);
    
    /// Softmax nonlinearity
    /// Y = Softmax(X) : Yij = e^Xij / sum_k(e^Xik), done to each row,
    /// with attention to avoiding  overflow or underflow.
    /// Supports in-place operation (i.e. this == &src).
    void SoftMaxPerRow(const CuMatrixBase<Real> &src);
  
    /// LogSoftmax nonlinearity
    /// Y = LogSoftmax(X) : Yij = Xij - log(sum_k(e^Xik)), done to each row,
    /// with attention to avoiding  overflow or underflow.
    /// Supports in-place operation (i.e. this == &src).
    void LogSoftMaxPerRow(const CuMatrixBase<Real> &src);
  
    
    /// Apply the function y = log(1 + exp(x)), to each element.
    /// Note: the derivative of this function is the sigmoid function.
    /// This is like a soft ReLU.
    void SoftHinge(const CuMatrixBase<Real> &src);
  
    /// Apply the function y(i) = (sum_{j = i*G}^{(i+1)*G-1} x_j ^ (power)) ^ (1 / p)
    /// where G = x.NumCols() / y.NumCols() must be an integer.
    /// [note: y corresponds to *this and x to src, so
    ///  src.NumCols() / this->NumCols() must be an integer.
    void GroupPnorm(const CuMatrixBase<Real> &src, Real pow);
  
    /// Differentiate backward through the GroupPnorm function.
    /// It is a combination of GroupPnormDeriv and MulRowsGroupMat.
    void DiffGroupPnorm(const CuMatrixBase<Real> &in_value,
                        const CuMatrixBase<Real> &out_value,
                        const CuMatrixBase<Real> &out_deriv, Real power);
  
    /// Apply the function y(i) = (max_{j = i*G}^{(i+1)*G-1} x_j
    /// where G = x.NumCols() / y.NumCols() must be an integer.
    /// [note: y corresponds to *this and x to src, so
    ///  src.NumCols() / this->NumCols() must be an integer.
    void GroupMax(const CuMatrixBase<Real> &src);
  
    /// Calculate derivatives for the GroupMax function above, where
    /// "input" is the input to the GroupMax function above (i.e. the "src" variable),
    /// and "output" is the result of the computation (i.e. the "this" of that function
    /// call), and *this must have the same dimension as "input". Each element
    /// of *this will be set to 1 if the corresponding input equals the output of
    /// the group, and 0 otherwise. The equals the function derivative where it is
    /// defined (it's not defined where multiple inputs in the group are equal to the output).
    void GroupMaxDeriv(const CuMatrixBase<Real> &input,
                       const CuMatrixBase<Real> &output);
  
    /// Compute the parametric rectified linear unit function;
    /// element by element, *this = src * (src > 0 ? alpha : beta)
    void ParametricRelu(const CuMatrixBase<Real> &src,
                        const CuVectorBase<Real> &alpha,
                        const CuVectorBase<Real> &beta);
  
    /// Differentiate backward through the parametric relu function.
    /// Here the "value" is the Relu input. Does, element-by-element.
    /// *this = diff * (value > 0 ? alpha : beta)
    void DiffParametricRelu(const CuMatrixBase<Real> &value,
                            const CuMatrixBase<Real> &diff,
                            const CuVectorBase<Real> &alpha,
                            const CuVectorBase<Real> &beta);
  
    /// Compute the hyperbolic tangent (tanh) function; element by element,
    /// *this = tanh(src).
    void Tanh(const CuMatrixBase<Real> &src);
  
    /// Differentiate backward through the sigmoid function.  Here, "value" is the
    /// sigmoid output.  Does, element-by-element, *this = diff * value * (1 - value).
    void DiffSigmoid(const CuMatrixBase<Real> &value,
                     const CuMatrixBase<Real> &diff);
  
    /// Differentiate backward through the tanh function.  Here, "value" is the
    /// tanh output.  Does, element-by-element, *this = diff * (1 - value^2).
    void DiffTanh(const CuMatrixBase<Real> &value,
                  const CuMatrixBase<Real> &diff);
  
    /// Differentiate backward through the softmax function.  Here, "value" is the
    /// softmax output. Does, for each row i,
    /// *this(i) =  diff(i) * diag(value(i)) - diff(i) * (value(i)^T * value(i))
    /// xxxx(i) is row-vector; '*' and '-' are matrix operations.
    /// Supports in-place operation, this  == &diff.
    void DiffSoftmaxPerRow(const CuMatrixBase<Real> &value,
                           const CuMatrixBase<Real> &diff);
  
    /// Differentiate backward through the log softmax function.
    /// Here, "out_value" is the log softmax output. Does, for each row i,
    /// *this(i) =  out_deriv(i) - sum(out_deriv(i)) .* exp(out_value(i))
    /// xxxx(i) is row-vector.
    /// Supports in-place operation, this == &out_deriv.
    void DiffLogSoftmaxPerRow(const CuMatrixBase<Real> &out_value,
                              const CuMatrixBase<Real> &out_deriv);
  
    /// Differentiate the block [softmax+cross-entropy] :
    /// dE/da = posterior_mat - target_mat,
    /// 'E' is error function, 'a' is activation on softmax input
    ///
    /// Interface:
    /// tgt ... index vector, encodes the matrix of targets
    /// net_out_or_diff ... before invocation net output, after diff dE/da
    /// log_post_tgt ... per-frame statistics for cross-entropy computations :
    ///                  log(sum_row(posterior_mat .* target_mat))
    void DiffXent(const CuArrayBase<int32> &tgt,
                  CuVector<Real> *log_post_tgt);
  
    /// This function does sets *this to the Cholesky factor of *this (i.e.  the C
    /// satisfying *this = C C^T), and sets "inv_cholesky" (if supplied) to its
    /// inverse.  *this is treated as a symmetric matrix but only the lower triangle
    /// is accessed.
    void Cholesky(CuMatrixBase<Real> *inv_cholesky = NULL);
  
  
    /// Inversion for positive definite symmetric matrices.
    /// Treats the input as symmetric but only reads the lower triangle.
    /// The output is symmetric.
    void SymInvertPosDef();
  
    inline void ApplyPow(Real power) {
      this -> Pow(*this, power);
    };
  
    
    inline void ApplyPowAbs(Real power, bool include_sign=false) {
      this -> PowAbs(*this, power, include_sign);
    };
    
    inline void ApplyHeaviside() {
      this -> Heaviside(*this);
    };
    
    inline void ApplyFloor(Real floor_val) {
      this -> Floor(*this, floor_val);
    };
    
    inline void ApplyCeiling(Real ceiling_val) {
      this -> Ceiling(*this, ceiling_val);
    };
    
    inline void ApplyExp() {
      this -> Exp(*this);
    };
  
  
    inline void ApplyExpLimited(Real lower_limit, Real upper_limit) {
      this -> ExpLimited(*this, lower_limit, upper_limit);
    };
  
    inline void ApplyExpSpecial() {
      this -> ExpSpecial(*this);
    };
  
    inline void ApplySoftMaxPerRow() {
      this -> SoftMaxPerRow(*this);
    };
  
    inline void ApplyLogSoftMaxPerRow() {
      this -> LogSoftMaxPerRow(*this);
    };
  
    inline void ApplyLog() {
      this -> Log(*this);
    };
  
    /// Find the id of the maximal element for each row (resizes the 'id'
    /// array to the appropriate size).
    void FindRowMaxId(CuArray<int32> *id) const;
  
    /// Math operations, some calling kernels
    void SetZero();
    void Set(Real value);
    void Add(Real value);
    /// Zeroes all elements for which col > row.
    void SetZeroAboveDiag();
    void Scale(Real value);
  
    /// Multiply two matrices elementwise: C = C .* A
    void MulElements(const CuMatrixBase<Real> &A);
    /// Divide two matrices elementwise: C = A ./ A
    void DivElements(const CuMatrixBase<Real> &A);
    /// Do, elementwise, *this = max(*this, A).
    void Max(const CuMatrixBase<Real> &A);
    /// Do, elementwise, *this = min(*this, A).
    void Min(const CuMatrixBase<Real> &A);
    /// scale i'th column by scale[i]
    void MulColsVec(const CuVectorBase<Real> &scale);
    /// scale i'th row by scale[i]
    void MulRowsVec(const CuVectorBase<Real> &scale);
    /// divide each row into src.NumCols() groups, and then scale i'th row's jth group of elements by src[i, j].
    void MulRowsGroupMat(const CuMatrixBase<Real> &src);
    /// divide i'th row by scale[i]
    void DivRowsVec(const CuVectorBase<Real> &div);
    /// invert the matrix by elements.
    void InvertElements();
    /// *this += alpha * A
    void AddMat(Real alpha, const CuMatrixBase<Real> &A,
                MatrixTransposeType trans = kNoTrans);
  
    /// *this += alpha * A.
    void AddSmat(Real alpha, const CuSparseMatrix<Real> &A,
                MatrixTransposeType trans = kNoTrans);
  
    /// (*this) = alpha * op(A) * B + beta * (*this), where A is sparse.
    /// Multiplication of sparse with dense matrix.  See also AddMatSmat.
    /// Note: we recommend, for greatest efficiency, that transA be kNoTrans.
    /// Use AddMatSmat() for better efficiency, as 2 dense mat transpose ops
    /// are called in this API.
    void AddSmatMat(Real alpha, const CuSparseMatrix<Real> &A,
                    MatrixTransposeType transA, const CuMatrixBase<Real> &B,
                    Real beta);
  
    /// (*this) = alpha * A * op(B) + beta * (*this), where B is sparse
    /// and op(B) is either B or trans(B) depending on the 'transB' argument.
    /// This is multiplication of a dense by a sparse matrix.  See also
    /// AddSmatMat.
    void AddMatSmat(Real alpha, const CuMatrixBase<Real> &A,
                    const CuSparseMatrix<Real> &B, MatrixTransposeType transB,
                    Real beta);
  
  
    /// This is a rather special purpose function; we might
    /// generalize it later by adding a transpose-type option.
    /// It expects 'elements.Dim()' to equal NumRows(), and
    /// for each elements[i] to be either -1, or
    /// 0 <= element[i] < NumCols().
    /// It adds alpha to each element (*this)(i, elements[i])
    /// for 0 <= i < NumRows().
    void AddToElements(Real alpha, const CuArrayBase<int32> &elements);
  
  
    /// This function is like AddMat (it does *this += alpha * src),
    /// except that it supports cases where *this and src have
    /// different dimension.  There are two allowed cases:
    ///
    ///  (1) *this is larger than src; we do a broadcasting operation.  *this must
    ///       have NumRows() == a * src.NumRows() and NumCols() == b *
    ///       src.NumCols() for integer a >= 1, b >= 1.  *this will be treated as
    ///       a being made up of of blocks with the same size as src, and to each
    ///       block we'll add alpha * src.  This case does not support trans ==
    ///       kTrans.
    ///
    ///  (2) *this is smaller than src; we sum.  src.NumRows() must == a *
    ///      this->NumRows(), and src.NumCols() must == b * this->NumCols(), for a
    ///      >= 1, b >= 1.  In this case, src will be treated as being made up of
    ///      blocks with the same size as *this, and to *this we will add the
    ///      summation of all of those blocks.
    void AddMatBlocks(Real alpha, const CuMatrixBase<Real> &A,
                      MatrixTransposeType trans = kNoTrans);
  
    /// (for each column c of *this), c = alpha * col + beta * c
    void AddVecToCols(Real alpha, const CuVectorBase<Real> &col, Real beta = 1.0);
    /// (for each row r of *this), r = alpha * row + beta * r
    void AddVecToRows(Real alpha, const CuVectorBase<Real> &row, Real beta = 1.0);
    /// C = alpha * A(^T)*B(^T) + beta * C
    void AddMatMat(Real alpha, const CuMatrixBase<Real> &A, MatrixTransposeType transA,
                   const CuMatrixBase<Real> &B, MatrixTransposeType transB, Real beta);
    /// A = alpha * x * y^T + A .
    void AddVecVec(Real alpha, const CuVectorBase<Real> &x, const CuVectorBase<Real> &y);
    /// *this = a * b / c (by element; when c = 0, *this = a)
    /// *this can be an alias of a, b or c safely and get expected result.
    void SetMatMatDivMat(const CuMatrixBase<Real> &A, const CuMatrixBase<Real> &B, const CuMatrixBase<Real> &C);
  
    /// *this = beta * *this + alpha * M M^T, for symmetric matrices.  It only
    /// updates the lower triangle of *this.  It will leave the matrix asymmetric;
    /// if you need it symmetric as a regular matrix, do CopyLowerToUpper().
    void SymAddMat2(const Real alpha, const CuMatrixBase<Real> &M,
                    MatrixTransposeType transA, Real beta);
  
  
    /// This function is like AddMatMat but for where the second argument is of
    /// type CuBlockMatrix (a block-diagonal matrix of blocks).
    void AddMatBlock(Real alpha, const CuMatrixBase<Real> &A, MatrixTransposeType transA,
                     const CuBlockMatrix<Real> &B, MatrixTransposeType transB, Real beta);
  
    /// *this = beta * *this + alpha * diag(v) * M [or M^T].
    /// The same as adding M but scaling each row M_i by v(i).
    void AddDiagVecMat(const Real alpha, const CuVectorBase<Real> &v,
                       const CuMatrixBase<Real> &M, MatrixTransposeType transM,
                       Real beta = 1.0);
  
    // *this = beta * *this + alpha * M  * diag(v) [or M^T].
    // The same as adding M but scaling each column M_j by v(j).
    void AddMatDiagVec(const Real alpha,
                       const CuMatrixBase<Real> &M, MatrixTransposeType transM,
                       CuVectorBase<Real> &v,
                       Real beta = 1.0);
  
    /// *this = beta * *this + alpha * A .* B (.* element by element multiplication)
    void AddMatMatElements(const Real alpha,
                           const CuMatrixBase<Real>& A,
                           const CuMatrixBase<Real>& B,
                           const Real beta);
  
    /// this <-- beta*this + alpha*A*B
    void AddMatSp(const Real alpha,
                  const CuMatrixBase<Real> &A, MatrixTransposeType transA,
                  const CuSpMatrix<Real> &B,
                  const Real beta) {
      CuMatrix<Real> M(B);
      return AddMatMat(alpha, A, transA, M, kNoTrans, beta);
    }
  
    /// this <-- beta*this + alpha*SpA*B
    void AddSpMat(const Real alpha,
                  const CuSpMatrix<Real> &A,
                  const CuMatrixBase<Real> &B, MatrixTransposeType transB,
                  const Real beta) {
      CuMatrix<Real> M(A);
      return AddMatMat(alpha, M, kNoTrans, B, transB, beta);
    }
  
    /// this <-- beta*this + alpha*A*B.
    void AddTpMat(const Real alpha,
                  const CuTpMatrix<Real> &A, MatrixTransposeType transA,
                  const CuMatrixBase<Real> &B, MatrixTransposeType transB,
                  const Real beta) {
      CuMatrix<Real> M(A);
      return AddMatMat(alpha, M, transA, B, transB, beta);
    }
  
    /// this <-- beta*this + alpha*A*B.
    void AddMatTp(const Real alpha,
                  const CuMatrixBase<Real> &A, MatrixTransposeType transA,
                  const CuTpMatrix<Real> &B, MatrixTransposeType transB,
                  const Real beta) {
      CuMatrix<Real> M(B);
      return AddMatMat(alpha, A, transA, M, transB, beta);
    }
  
    void CopyFromBlock(const CuBlockMatrix<Real> &B,
                       MatrixTransposeType trans = kNoTrans);
    void CopyLowerToUpper();
    void CopyUpperToLower();
    inline CuSubMatrix<Real> Range(const MatrixIndexT row_offset,
                                   const MatrixIndexT num_rows,
                                   const MatrixIndexT col_offset,
                                   const MatrixIndexT num_cols) const {
      return CuSubMatrix<Real>(*this, row_offset, num_rows,
                               col_offset, num_cols);
    }
    inline CuSubMatrix<Real> RowRange(const MatrixIndexT row_offset,
                                      const MatrixIndexT num_rows) const {
      return CuSubMatrix<Real>(*this, row_offset, num_rows,
                               0, num_cols_);
    }
    inline CuSubMatrix<Real> ColRange(const MatrixIndexT col_offset,
                                      const MatrixIndexT num_cols) const {
      return CuSubMatrix<Real>(*this, 0, num_rows_, col_offset, num_cols);
    }
  
    inline const CuSubVector<Real> Row(MatrixIndexT i) const {
      KALDI_ASSERT(static_cast<UnsignedMatrixIndexT>(i) <
                   static_cast<UnsignedMatrixIndexT>(num_rows_));
      return CuSubVector<Real>(data_ + (i * stride_), NumCols());
    }
  
    inline CuSubVector<Real> Row(MatrixIndexT i) {
      KALDI_ASSERT(static_cast<UnsignedMatrixIndexT>(i) <
                   static_cast<UnsignedMatrixIndexT>(num_rows_));
      return CuSubVector<Real>(data_ + (i * stride_), NumCols());
    }
  
    inline CuValue<Real> operator() (MatrixIndexT r, MatrixIndexT c) {
      KALDI_PARANOID_ASSERT(static_cast<UnsignedMatrixIndexT>(r) <
                            static_cast<UnsignedMatrixIndexT>(num_rows_) &&
                            static_cast<UnsignedMatrixIndexT>(c) <
                            static_cast<UnsignedMatrixIndexT>(num_cols_));
      return CuValue<Real>(data_ + r * stride_ + c);
    }
  
    inline Real operator() (MatrixIndexT r, MatrixIndexT c) const {
      KALDI_PARANOID_ASSERT(static_cast<UnsignedMatrixIndexT>(r) <
                            static_cast<UnsignedMatrixIndexT>(num_rows_) &&
                            static_cast<UnsignedMatrixIndexT>(c) <
                            static_cast<UnsignedMatrixIndexT>(num_cols_));
      return CuValue<Real>(data_ + r * stride_ + c);  // will be casted to Real.
    }
  
    Real Sum() const;
    Real Max() const;
    Real Min() const;
  
    /// Return the trace. If check_square = true, will crash if matrix is not square.
    Real Trace(bool check_square = true) const;
  
    void SetRandn();
  
    void SetRandUniform();
  
    void Write(std::ostream &os, bool binary) const;
  
    // This function, adds a list of MatrixElements (scaled by alpha) to corresponding locations to
    // (*this).
    void AddElements(Real alpha, const std::vector<MatrixElement<Real> >& input);
  
    // For each i, with indexes[i] = (j, k), does (*this)(j, k) += input[i].
    // Requires, but does not check, that the vector of indexes does not contrain
    // repeated elements, 'input' is the start of an array of length equal to
    // indexes.Dim(), which is located on GPU memory if we are using the GPU.
    void AddElements(Real alpha, const CuArrayBase<Int32Pair> &indexes,
                     const Real *input);
  
    // This function requires that 'output' is a host array and is allocated with size
    // of indexes.size(), and for each element of 'indexes' it interprets it as
    // a (row, column) index into *this, and puts (*this)(row, column) into
    // the corresponding element of 'output'.
    void Lookup(const std::vector<Int32Pair> &indexes,
                Real *output) const;
  
    // CUDA version of Lookup, would be called internally by the above function.
    void Lookup(const CuArrayBase<Int32Pair> &indexes,
                Real *output) const;
  
    // Creates binary mask with per-element equality predicates of *this, mat.
    // Output stored to 'mask', values : 1.0 = equal, 0.0 = not-equal.
    void EqualElementMask(const CuMatrixBase<Real> &mat, CuMatrix<Real> *mask) const;
  
  
    /// Get raw row pointer (const).  Warning: may return a pointer to GPU memory.  Use at
    /// your own risk.
    inline const Real* RowData(MatrixIndexT r) const { return data_ + r * stride_; }
    /// Get raw row pointer.  Warning: may return a pointer to GPU memory.  Use at
    /// your own risk.
    inline Real* RowData(MatrixIndexT r) { return data_ + r * stride_; }
    /// Return data pointer (const).  Warning: may return a pointer to GPU memory.
    /// Use at your own risk.
    inline const Real *Data() const { return data_; }
    /// Return data pointer.  Warning: may return a pointer to GPU memory.  Use at
    /// your own risk.
    inline Real *Data() { return data_; }
  
    // The following two functions should only be called if we did not compile
    // with CUDA or could not get a CUDA card; in that case the contents are
    // interpreted the same as a regular matrix.  DON'T USE THESE UNLESS YOU KNOW
    // WHAT YOU ARE DOING!
    inline const MatrixBase<Real> &Mat() const {
      return *(reinterpret_cast<const MatrixBase<Real>* >(this));
    }
    inline MatrixBase<Real> &Mat() {
      return *(reinterpret_cast<MatrixBase<Real>* >(this));
    }
  
   protected:
  
    // The constructors are protected to prevent the user creating an instance of
    // this class (you should create a child class CuMatrix or CuSubMatrix.
  
    CuMatrixBase(): data_(NULL), num_cols_(0), num_rows_(0), stride_(0) { }
  
    /// This constructor takes the #rows, #cols and stride; it's called from
    /// the constructor of CuSubMatrix.
    CuMatrixBase(Real *data,
                 MatrixIndexT num_rows,
                 MatrixIndexT num_cols,
                 MatrixIndexT stride):
    data_(data), num_cols_(num_cols), num_rows_(num_rows), stride_(stride) { }
  
    Real *data_;       ///< GPU data pointer (or regular matrix data pointer,
    ///< if either CUDA was not compiled in or we could not
    ///< acquire the device).
    // Note: it might seem a bit backwards that we have the number of columns
    // first here; it's necessary because we need the data to be laid out the same
    // as for MatrixBase so the Mat() function call will work.  We don't want to
    // change the layout of MatrixBase at this point, or there will be crashes if
    // people don't thoroughly recompile.
    MatrixIndexT num_cols_;
    MatrixIndexT num_rows_;
    MatrixIndexT stride_;
  
   private:
    KALDI_DISALLOW_COPY_AND_ASSIGN(CuMatrixBase);
  }; // class CuMatrixBase
  
  /// This class represents a matrix that's stored on the GPU if we have one,
  /// and in memory if not.
  template<typename Real>
  class CuMatrix: public CuMatrixBase<Real> {
   public:
  
    CuMatrix() { }
  
    /// Constructor with memory initialisation
    CuMatrix(MatrixIndexT rows, MatrixIndexT cols,
             MatrixResizeType resize_type = kSetZero,
             MatrixStrideType stride_type = kDefaultStride) {
      Resize(rows, cols, resize_type, stride_type);
    }
  
    // Note: we had to remove the "explicit" keyword due
    // to problems with STL vectors of CuMatrixBase.
    CuMatrix(const CuMatrix<Real> &other,
             MatrixTransposeType trans = kNoTrans);
  
    explicit CuMatrix(const CuBlockMatrix<Real> &other,
                      MatrixTransposeType trans = kNoTrans);
  
    explicit CuMatrix(const CuMatrixBase<Real> &other,
                      MatrixTransposeType trans = kNoTrans);
  
    template<typename OtherReal>
    explicit CuMatrix(const MatrixBase<OtherReal> &other,
                      MatrixTransposeType trans = kNoTrans);
  
    /// Copy constructor taking SpMatrix...
    explicit CuMatrix(const CuSpMatrix<Real> &M) : CuMatrixBase<Real>() {
      Resize(M.NumRows(), M.NumRows(), kUndefined);
      this->CopyFromSp(M);
    }
  
    /// Copy constructor taking TpMatrix...
    template <typename OtherReal>
    explicit CuMatrix(const CuTpMatrix<OtherReal> & M,
                      MatrixTransposeType trans = kNoTrans) : CuMatrixBase<Real>() {
      Resize(M.NumCols(), M.NumRows(), kUndefined);
      this->CopyFromTp(M, trans);
    }
  
    /// Copy constructor: as above, but from another type.
    template<typename OtherReal>
    explicit CuMatrix(const CuMatrixBase<OtherReal> &M,
                      MatrixTransposeType trans = kNoTrans);
  
    CuMatrix<Real> &operator = (const CuMatrixBase<Real> &other) {
      this->Resize(other.NumRows(), other.NumCols(), kUndefined);
      this->CopyFromMat(other);
      return *this;
    }
  
    CuMatrix<Real> &operator = (const CuMatrix<Real> &other) {
      this->Resize(other.NumRows(), other.NumCols(), kUndefined);
      this->CopyFromMat(other);
      return *this;
    }
  
    CuMatrix<Real> &operator = (const MatrixBase<Real> &other) {
      this->Resize(other.NumRows(), other.NumCols(), kUndefined);
      this->CopyFromMat(other);
      return *this;
    }
  
    void Transpose();
  
    /// Allocate the memory
    void Resize(MatrixIndexT rows, MatrixIndexT cols,
                MatrixResizeType resize_type = kSetZero,
                MatrixStrideType stride_type = kDefaultStride);
  
    void Swap(Matrix<Real> *mat);
    void Swap(CuMatrix<Real> *mat);
  
    template<typename OtherReal>
    void Swap(CuMatrix<OtherReal> *mat);
  
    /// I/O functions
    void Read(std::istream &is, bool binary);
  
    /// Destructor
    ~CuMatrix() { Destroy(); }
  
    inline const Matrix<Real> &Mat() const {
      return *(reinterpret_cast<const Matrix<Real>* >(this));
    }
    inline Matrix<Real> &Mat() {
      return *(reinterpret_cast<Matrix<Real>* >(this));
    }
  
    /// Here, A is interpreted as a matrix of probabilities, and "elements" as a list
    /// of posteriors (possibly zero-one), and "*this" as a matrix of derivatives
    /// w.r.t. the log-probs.
    /// This function does: for each element { row, column, weight } indexed i in
    /// the vector "elements", let x(i) = A(row(i), column(i)); then it does
    /// (*this)(row(i), column(i)) += weight(i) / x(i), and
    /// *tot_objf = \sum_i weight(i) * log(x(i)), and
    /// *tot_weight = \sum_i weight(i)
    /// Preconditions: A must be strictly positive, and no (row, column) pair
    /// may be repeated within "elements"
    void CompObjfAndDeriv(const std::vector<MatrixElement<Real> > &elements,
                          const CuMatrix<Real> &A,
                          Real *tot_objf,
                          Real *tot_weight);
  
   private:
    void Destroy();
  };
  
  
  /// This class is used for a piece of a CuMatrix.
  template<typename Real>
  class CuSubMatrix: public CuMatrixBase<Real> {
   public:
    inline CuSubMatrix(const CuMatrixBase<Real> &mat,
                       const MatrixIndexT row_offset,
                       const MatrixIndexT num_rows,
                       const MatrixIndexT col_offset,
                       const MatrixIndexT num_cols);
  
    // This constructor should be used with caution; it can be used for
    // constructing 'fake' submatrices if you want to play with
    // the stride. 'data' should point to GPU data if you're using the
    // GPU.
    inline CuSubMatrix(const Real *data,
                       const MatrixIndexT num_rows,
                       const MatrixIndexT num_cols,
                       const MatrixIndexT stride);
  
    /// This type of constructor is needed for Range() to work [in CuMatrix base
    /// class]. Cannot make it explicit or that breaks.
    inline CuSubMatrix<Real> (const CuSubMatrix &other):
    CuMatrixBase<Real> (other.data_, other.num_rows_, other.num_cols_,
                        other.stride_) {}
   private:
    /// Disallow assignment.
    CuSubMatrix<Real> &operator = (const CuSubMatrix<Real> &other);
  };
  
  
  template<typename Real>
  bool ApproxEqual(const CuMatrixBase<Real> &A,
                   const CuMatrixBase<Real> &B, Real tol = 0.01) {
    return A.ApproxEqual(B, tol);
  }
  
  template<typename Real>
  inline void AssertEqual(const CuMatrixBase<Real> &A,
                          const CuMatrixBase<Real> &B, float tol = 0.01) {
    KALDI_ASSERT(A.ApproxEqual(B, tol));
  }
  
  template<typename Real>
  bool SameDim(const CuMatrixBase<Real> &M, const CuMatrixBase<Real> &N) {
    return (M.NumRows() == N.NumRows() && M.NumCols() == N.NumCols());
  }
  
  template<typename Real>
  bool SameDimAndStride(const CuMatrixBase<Real> &M, const CuMatrixBase<Real> &N) {
    return (M.NumRows() == N.NumRows() && M.NumCols() == N.NumCols()
            && M.Stride() == N.Stride());
  }
  
  /// I/O
  template<typename Real>
  std::ostream &operator << (std::ostream &out, const CuMatrixBase<Real> &mat);
  
  
  template<typename Real>
  template<typename OtherReal>
  Matrix<Real>::Matrix(const CuMatrixBase<OtherReal> &M,
                       MatrixTransposeType trans) {
    if (trans == kNoTrans) Init(M.NumRows(), M.NumCols(), kDefaultStride);
    else Init(M.NumCols(), M.NumRows(), kDefaultStride);
    M.CopyToMat(this, trans);
  }
  
  template<typename Real>
  template<typename OtherReal>
  void MatrixBase<Real>::CopyFromMat(const CuMatrixBase<OtherReal> &cu,
                                     MatrixTransposeType trans) {
    cu.CopyToMat(this, trans);
  }
  
  
  }  // namespace
  
  
  #include "cudamatrix/cu-matrix-inl.h"
  
  #endif