// nnet3/nnet-combined-component.cc

// Copyright 2015-2018  Johns Hopkins University (author: Daniel Povey)
//                2015  Daniel Galvez
//                2018  Hang Lyu

// 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.

#include <iterator>
#include <sstream>
#include <algorithm>
#include <iomanip>
#include "nnet3/nnet-combined-component.h"
#include "nnet3/nnet-parse.h"
#include "cudamatrix/cu-math.h"

namespace kaldi {
namespace nnet3 {

// Constructors for the convolution component
ConvolutionComponent::ConvolutionComponent():
    UpdatableComponent(),
    input_x_dim_(0), input_y_dim_(0), input_z_dim_(0),
    filt_x_dim_(0), filt_y_dim_(0),
    filt_x_step_(0), filt_y_step_(0),
    input_vectorization_(kZyx) { }

ConvolutionComponent::ConvolutionComponent(
    const ConvolutionComponent &component):
    UpdatableComponent(component),
    input_x_dim_(component.input_x_dim_),
    input_y_dim_(component.input_y_dim_),
    input_z_dim_(component.input_z_dim_),
    filt_x_dim_(component.filt_x_dim_),
    filt_y_dim_(component.filt_y_dim_),
    filt_x_step_(component.filt_x_step_),
    filt_y_step_(component.filt_y_step_),
    input_vectorization_(component.input_vectorization_),
    filter_params_(component.filter_params_),
    bias_params_(component.bias_params_) { }

ConvolutionComponent::ConvolutionComponent(
    const CuMatrixBase<BaseFloat> &filter_params,
    const CuVectorBase<BaseFloat> &bias_params,
    int32 input_x_dim, int32 input_y_dim, int32 input_z_dim,
    int32 filt_x_dim, int32 filt_y_dim,
    int32 filt_x_step, int32 filt_y_step,
    TensorVectorizationType input_vectorization,
    BaseFloat learning_rate):
    input_x_dim_(input_x_dim),
    input_y_dim_(input_y_dim),
    input_z_dim_(input_z_dim),
    filt_x_dim_(filt_x_dim),
    filt_y_dim_(filt_y_dim),
    filt_x_step_(filt_x_step),
    filt_y_step_(filt_y_step),
    input_vectorization_(input_vectorization),
    filter_params_(filter_params),
    bias_params_(bias_params){
  KALDI_ASSERT(filter_params.NumRows() == bias_params.Dim() &&
               bias_params.Dim() != 0);
  KALDI_ASSERT(filter_params.NumCols() == filt_x_dim * filt_y_dim * input_z_dim);
  SetUnderlyingLearningRate(learning_rate);
  is_gradient_ = false;
}

// aquire input dim
int32 ConvolutionComponent::InputDim() const {
  return input_x_dim_ * input_y_dim_ * input_z_dim_;
}

// aquire output dim
int32 ConvolutionComponent::OutputDim() const {
  int32 num_x_steps = (1 + (input_x_dim_ - filt_x_dim_) / filt_x_step_);
  int32 num_y_steps = (1 + (input_y_dim_ - filt_y_dim_) / filt_y_step_);
  int32 num_filters = filter_params_.NumRows();
  return num_x_steps * num_y_steps * num_filters;
}

// initialize the component using hyperparameters
void ConvolutionComponent::Init(
    int32 input_x_dim, int32 input_y_dim, int32 input_z_dim,
    int32 filt_x_dim, int32 filt_y_dim,
    int32 filt_x_step, int32 filt_y_step, int32 num_filters,
    TensorVectorizationType input_vectorization,
    BaseFloat param_stddev, BaseFloat bias_stddev) {
  input_x_dim_ = input_x_dim;
  input_y_dim_ = input_y_dim;
  input_z_dim_ = input_z_dim;
  filt_x_dim_ = filt_x_dim;
  filt_y_dim_ = filt_y_dim;
  filt_x_step_ = filt_x_step;
  filt_y_step_ = filt_y_step;
  input_vectorization_ = input_vectorization;
  KALDI_ASSERT((input_x_dim_ - filt_x_dim_) % filt_x_step_ == 0);
  KALDI_ASSERT((input_y_dim_ - filt_y_dim_) % filt_y_step_ == 0);
  int32 filter_dim = filt_x_dim_ * filt_y_dim_ * input_z_dim_;
  filter_params_.Resize(num_filters, filter_dim);
  bias_params_.Resize(num_filters);
  KALDI_ASSERT(param_stddev >= 0.0 && bias_stddev >= 0.0);
  filter_params_.SetRandn();
  filter_params_.Scale(param_stddev);
  bias_params_.SetRandn();
  bias_params_.Scale(bias_stddev);
}

// initialize the component using predefined matrix file
void ConvolutionComponent::Init(
    int32 input_x_dim, int32 input_y_dim, int32 input_z_dim,
    int32 filt_x_dim, int32 filt_y_dim,
    int32 filt_x_step, int32 filt_y_step,
    TensorVectorizationType input_vectorization,
    std::string matrix_filename) {
  input_x_dim_ = input_x_dim;
  input_y_dim_ = input_y_dim;
  input_z_dim_ = input_z_dim;
  filt_x_dim_ = filt_x_dim;
  filt_y_dim_ = filt_y_dim;
  filt_x_step_ = filt_x_step;
  filt_y_step_ = filt_y_step;
  input_vectorization_ = input_vectorization;
  CuMatrix<BaseFloat> mat;
  ReadKaldiObject(matrix_filename, &mat);
  int32 filter_dim = (filt_x_dim_ * filt_y_dim_ * input_z_dim_);
  int32 num_filters = mat.NumRows();
  KALDI_ASSERT(mat.NumCols() == (filter_dim + 1));
  filter_params_.Resize(num_filters, filter_dim);
  bias_params_.Resize(num_filters);
  filter_params_.CopyFromMat(mat.Range(0, num_filters, 0, filter_dim));
  bias_params_.CopyColFromMat(mat, filter_dim);
}

// display information about component
std::string ConvolutionComponent::Info() const {
  std::ostringstream stream;
  stream << UpdatableComponent::Info()
         << ", input-x-dim=" << input_x_dim_
         << ", input-y-dim=" << input_y_dim_
         << ", input-z-dim=" << input_z_dim_
         << ", filt-x-dim=" << filt_x_dim_
         << ", filt-y-dim=" << filt_y_dim_
         << ", filt-x-step=" << filt_x_step_
         << ", filt-y-step=" << filt_y_step_
         << ", input-vectorization=" << input_vectorization_
         << ", num-filters=" << filter_params_.NumRows();
  PrintParameterStats(stream, "filter-params", filter_params_);
  PrintParameterStats(stream, "bias-params", bias_params_, true);
  return stream.str();
}

// initialize the component using configuration file
void ConvolutionComponent::InitFromConfig(ConfigLine *cfl) {
  bool ok = true;
  std::string matrix_filename;
  int32 input_x_dim = -1, input_y_dim = -1, input_z_dim = -1,
        filt_x_dim = -1, filt_y_dim = -1,
        filt_x_step = -1, filt_y_step = -1,
        num_filters = -1;
  std::string input_vectorization_order = "zyx";
  InitLearningRatesFromConfig(cfl);
  ok = ok && cfl->GetValue("input-x-dim", &input_x_dim);
  ok = ok && cfl->GetValue("input-y-dim", &input_y_dim);
  ok = ok && cfl->GetValue("input-z-dim", &input_z_dim);
  ok = ok && cfl->GetValue("filt-x-dim", &filt_x_dim);
  ok = ok && cfl->GetValue("filt-y-dim", &filt_y_dim);
  ok = ok && cfl->GetValue("filt-x-step", &filt_x_step);
  ok = ok && cfl->GetValue("filt-y-step", &filt_y_step);

  if (!ok)
    KALDI_ERR << "Bad initializer " << cfl->WholeLine();
  // optional argument
  TensorVectorizationType input_vectorization;
  cfl->GetValue("input-vectorization-order", &input_vectorization_order);
  if (input_vectorization_order.compare("zyx") == 0) {
    input_vectorization = kZyx;
  } else if (input_vectorization_order.compare("yzx") == 0) {
    input_vectorization = kYzx;
  } else {
    KALDI_ERR << "Unknown or unsupported input vectorization order "
              << input_vectorization_order
              << " accepted candidates are 'yzx' and 'zyx'";
  }

  if (cfl->GetValue("matrix", &matrix_filename)) {
    // initialize from prefined parameter matrix
    Init(input_x_dim, input_y_dim, input_z_dim,
         filt_x_dim, filt_y_dim,
         filt_x_step, filt_y_step,
         input_vectorization,
         matrix_filename);
  } else {
    ok = ok && cfl->GetValue("num-filters", &num_filters);
    if (!ok)
      KALDI_ERR << "Bad initializer " << cfl->WholeLine();
    // initialize from configuration
    int32 filter_input_dim = filt_x_dim * filt_y_dim * input_z_dim;
    BaseFloat param_stddev = 1.0 / std::sqrt(filter_input_dim), bias_stddev = 1.0;
    cfl->GetValue("param-stddev", &param_stddev);
    cfl->GetValue("bias-stddev", &bias_stddev);
    Init(input_x_dim, input_y_dim, input_z_dim,
         filt_x_dim, filt_y_dim, filt_x_step, filt_y_step, num_filters,
         input_vectorization, param_stddev, bias_stddev);
  }
  if (cfl->HasUnusedValues())
    KALDI_ERR << "Could not process these elements in initializer: "
              << cfl->UnusedValues();
  if (!ok)
    KALDI_ERR << "Bad initializer " << cfl->WholeLine();
}

// Inline methods to convert from tensor index i.e., (x,y,z) index
// to index in yzx or zyx vectorized tensors
inline int32 YzxVectorIndex(int32 x, int32 y, int32 z,
                            int32 input_x_dim,
                            int32 input_y_dim,
                            int32 input_z_dim) {
  KALDI_PARANOID_ASSERT(x < input_x_dim && y < input_y_dim && z < input_z_dim);
  return (input_y_dim * input_z_dim) * x + (input_y_dim) * z + y;
}

inline int32 ZyxVectorIndex(int32 x, int32 y, int32 z,
                            int32 input_x_dim,
                            int32 input_y_dim,
                            int32 input_z_dim) {
  KALDI_PARANOID_ASSERT(x < input_x_dim && y < input_y_dim && z < input_z_dim);
  return (input_y_dim * input_z_dim) * x + (input_z_dim) * y + z;
}

// Method to convert from a matrix representing a minibatch of vectorized
// 3D tensors to patches for convolution, each patch corresponds to
// one dot product in the convolution
void ConvolutionComponent::InputToInputPatches(
    const CuMatrixBase<BaseFloat>& in,
    CuMatrix<BaseFloat> *patches) const{
  int32 num_x_steps = (1 + (input_x_dim_ - filt_x_dim_) / filt_x_step_);
  int32 num_y_steps = (1 + (input_y_dim_ - filt_y_dim_) / filt_y_step_);
  const int32 filt_x_step = filt_x_step_,
              filt_y_step = filt_y_step_,
              filt_x_dim = filt_x_dim_,
              filt_y_dim = filt_y_dim_,
              input_x_dim = input_x_dim_,
              input_y_dim = input_y_dim_,
              input_z_dim = input_z_dim_,
              filter_dim = filter_params_.NumCols();

  std::vector<int32> column_map(patches->NumCols());
  int32 column_map_size = column_map.size();
  for (int32 x_step = 0; x_step < num_x_steps; x_step++) {
    for (int32 y_step = 0; y_step < num_y_steps; y_step++)  {
      int32 patch_number = x_step * num_y_steps + y_step;
      int32 patch_start_index = patch_number * filter_dim;
      for (int32 x = 0, index = patch_start_index; x < filt_x_dim; x++)  {
        for (int32 y = 0; y < filt_y_dim; y++)  {
          for (int32 z = 0; z < input_z_dim; z++, index++)  {
            KALDI_ASSERT(index < column_map_size);
            if (input_vectorization_ == kZyx)  {
              column_map[index] = ZyxVectorIndex(x_step * filt_x_step + x,
                                                 y_step * filt_y_step + y, z,
                                                 input_x_dim, input_y_dim,
                                                 input_z_dim);
            } else if (input_vectorization_ == kYzx)  {
              column_map[index] = YzxVectorIndex(x_step * filt_x_step + x,
                                                  y_step * filt_y_step + y, z,
                                                  input_x_dim, input_y_dim,
                                                  input_z_dim);
            }
          }
        }
      }
    }
  }
  CuArray<int32> cu_cols(column_map);
  patches->CopyCols(in, cu_cols);
}


// propagation function
// see function declaration in nnet-simple-component.h for details
void* ConvolutionComponent::Propagate(const ComponentPrecomputedIndexes *indexes,
                                         const CuMatrixBase<BaseFloat> &in,
                                         CuMatrixBase<BaseFloat> *out) const {
  const int32 num_x_steps = (1 + (input_x_dim_ - filt_x_dim_) / filt_x_step_),
              num_y_steps = (1 + (input_y_dim_ - filt_y_dim_) / filt_y_step_),
              num_filters = filter_params_.NumRows(),
              num_frames = in.NumRows(),
              filter_dim = filter_params_.NumCols();
  KALDI_ASSERT((*out).NumRows() == num_frames &&
               (*out).NumCols() == (num_filters * num_x_steps * num_y_steps));

  CuMatrix<BaseFloat> patches(num_frames,
                              num_x_steps * num_y_steps * filter_dim,
                              kUndefined);
  InputToInputPatches(in, &patches);
  CuSubMatrix<BaseFloat>* filter_params_elem = new CuSubMatrix<BaseFloat>(
      filter_params_, 0, filter_params_.NumRows(), 0, filter_params_.NumCols());
  std::vector<CuSubMatrix<BaseFloat>* > tgt_batch, patch_batch,
      filter_params_batch;

  for (int32 x_step = 0; x_step < num_x_steps; x_step++)  {
    for (int32 y_step = 0; y_step < num_y_steps; y_step++)  {
      int32 patch_number = x_step * num_y_steps + y_step;
      tgt_batch.push_back(new CuSubMatrix<BaseFloat>(
              out->ColRange(patch_number * num_filters, num_filters)));
      patch_batch.push_back(new CuSubMatrix<BaseFloat>(
              patches.ColRange(patch_number * filter_dim, filter_dim)));
      filter_params_batch.push_back(filter_params_elem);
      tgt_batch[patch_number]->AddVecToRows(1.0, bias_params_, 1.0); // add bias
    }
  }
  // apply all filters
  AddMatMatBatched<BaseFloat>(1.0, tgt_batch, patch_batch,
                              kNoTrans, filter_params_batch,
                              kTrans, 1.0);
  // release memory
  delete filter_params_elem;
  for (int32 p = 0; p < tgt_batch.size(); p++) {
    delete tgt_batch[p];
    delete patch_batch[p];
  }
  return NULL;
}

// scale the parameters
void ConvolutionComponent::Scale(BaseFloat scale) {
  if (scale == 0.0) {
    filter_params_.SetZero();
    bias_params_.SetZero();
  } else {
    filter_params_.Scale(scale);
    bias_params_.Scale(scale);
  }
}

// add another convolution component
void ConvolutionComponent::Add(BaseFloat alpha, const Component &other_in) {
  const ConvolutionComponent *other =
      dynamic_cast<const ConvolutionComponent*>(&other_in);
  KALDI_ASSERT(other != NULL);
  filter_params_.AddMat(alpha, other->filter_params_);
  bias_params_.AddVec(alpha, other->bias_params_);
}

/*
 This function transforms a vector of lists into a list of vectors,
 padded with -1.
 @param[in] The input vector of lists. Let in.size() be D, and let
            the longest list length (i.e. the max of in[i].size()) be L.
 @param[out] The output list of vectors. The length of the list will
            be L, each vector-dimension will be D (i.e. out[i].size() == D),
            and if in[i] == j, then for some k we will have that
            out[k][j] = i. The output vectors are padded with -1
            where necessary if not all the input lists have the same side.
*/
void RearrangeIndexes(const std::vector<std::vector<int32> > &in,
                                                std::vector<std::vector<int32> > *out) {
  int32 D = in.size();
  int32 L = 0;
  for (int32 i = 0; i < D; i++)
    if (in[i].size() > L)
      L = in[i].size();
  out->resize(L);
  for (int32 i = 0; i < L; i++)
    (*out)[i].resize(D, -1);
  for (int32 i = 0; i < D; i++) {
    for (int32 j = 0; j < in[i].size(); j++) {
      (*out)[j][i] = in[i][j];
    }
  }
}

// Method to compute the input derivative matrix from the input derivatives
// for patches, where each patch corresponds to one dot product
// in the convolution
void ConvolutionComponent::InderivPatchesToInderiv(
    const CuMatrix<BaseFloat>& in_deriv_patches,
    CuMatrixBase<BaseFloat> *in_deriv) const {

  const int32 num_x_steps = (1 + (input_x_dim_ - filt_x_dim_) / filt_x_step_),
              num_y_steps = (1 + (input_y_dim_ - filt_y_dim_) / filt_y_step_),
              filt_x_step = filt_x_step_,
              filt_y_step = filt_y_step_,
              filt_x_dim = filt_x_dim_,
              filt_y_dim = filt_y_dim_,
              input_x_dim = input_x_dim_,
              input_y_dim = input_y_dim_,
              input_z_dim = input_z_dim_,
              filter_dim = filter_params_.NumCols();

  // Compute the reverse column_map from the matrix with input
  // derivative patches to input derivative matrix
  std::vector<std::vector<int32> > reverse_column_map(in_deriv->NumCols());
  int32 rev_col_map_size = reverse_column_map.size();
  for (int32 x_step = 0; x_step < num_x_steps; x_step++) {
    for (int32 y_step = 0; y_step < num_y_steps; y_step++)  {
      int32 patch_number = x_step * num_y_steps + y_step;
      int32 patch_start_index = patch_number * filter_dim;
      for (int32 x = 0, index = patch_start_index; x < filt_x_dim; x++)  {
        for (int32 y = 0; y < filt_y_dim; y++)  {
          for (int32 z = 0; z < input_z_dim; z++, index++)  {
            int32 vector_index;
            if (input_vectorization_ == kZyx)  {
              vector_index = ZyxVectorIndex(x_step * filt_x_step + x,
                                            y_step * filt_y_step + y, z,
                                            input_x_dim, input_y_dim,
                                            input_z_dim);
            } else {
              KALDI_ASSERT(input_vectorization_ == kYzx);
              vector_index = YzxVectorIndex(x_step * filt_x_step + x,
                                            y_step * filt_y_step + y, z,
                                            input_x_dim, input_y_dim,
                                            input_z_dim);
            }
            KALDI_ASSERT(vector_index < rev_col_map_size);
            reverse_column_map[vector_index].push_back(index);
          }
        }
      }
    }
  }
  std::vector<std::vector<int32> > rearranged_column_map;
  RearrangeIndexes(reverse_column_map, &rearranged_column_map);
  for (int32 p = 0; p < rearranged_column_map.size(); p++) {
    CuArray<int32> cu_cols(rearranged_column_map[p]);
    in_deriv->AddCols(in_deriv_patches, cu_cols);
  }
}

// back propagation function
// see function declaration in nnet-simple-component.h for details
void ConvolutionComponent::Backprop(const std::string &debug_info,
                                    const ComponentPrecomputedIndexes *indexes,
                                    const CuMatrixBase<BaseFloat> &in_value,
                                    const CuMatrixBase<BaseFloat> &, // out_value,
                                    const CuMatrixBase<BaseFloat> &out_deriv,
                                    void *memo,
                                    Component *to_update_in,
                                    CuMatrixBase<BaseFloat> *in_deriv) const {
  ConvolutionComponent *to_update =
      dynamic_cast<ConvolutionComponent*>(to_update_in);
  const int32 num_x_steps = (1 + (input_x_dim_ - filt_x_dim_) / filt_x_step_),
              num_y_steps = (1 + (input_y_dim_ - filt_y_dim_) / filt_y_step_),
              num_filters = filter_params_.NumRows(),
              num_frames = out_deriv.NumRows(),
              filter_dim = filter_params_.NumCols();

  KALDI_ASSERT(out_deriv.NumRows() == num_frames &&
               out_deriv.NumCols() ==
               (num_filters * num_x_steps * num_y_steps));

  // Compute inderiv patches
  CuMatrix<BaseFloat> in_deriv_patches(num_frames,
                                       num_x_steps * num_y_steps * filter_dim,
                                       kSetZero);

  std::vector<CuSubMatrix<BaseFloat>* > patch_deriv_batch, out_deriv_batch,
      filter_params_batch;
  CuSubMatrix<BaseFloat>* filter_params_elem = new CuSubMatrix<BaseFloat>(
      filter_params_, 0, filter_params_.NumRows(), 0, filter_params_.NumCols());

  for (int32 x_step = 0; x_step < num_x_steps; x_step++)  {
    for (int32 y_step = 0; y_step < num_y_steps; y_step++)  {
      int32 patch_number = x_step * num_y_steps + y_step;

      patch_deriv_batch.push_back(new CuSubMatrix<BaseFloat>(
              in_deriv_patches.ColRange(
              patch_number * filter_dim, filter_dim)));
      out_deriv_batch.push_back(new CuSubMatrix<BaseFloat>(out_deriv.ColRange(
              patch_number * num_filters, num_filters)));
      filter_params_batch.push_back(filter_params_elem);
    }
  }
  AddMatMatBatched<BaseFloat>(1.0, patch_deriv_batch,
                              out_deriv_batch, kNoTrans,
                              filter_params_batch, kNoTrans, 0.0);

  if (in_deriv) {
    // combine the derivatives from the individual input deriv patches
    // to compute input deriv matrix
    InderivPatchesToInderiv(in_deriv_patches, in_deriv);
  }

  if (to_update != NULL)  {
    to_update->Update(debug_info, in_value, out_deriv, out_deriv_batch);
  }

  // release memory
  delete filter_params_elem;
  for (int32 p = 0; p < patch_deriv_batch.size(); p++) {
    delete patch_deriv_batch[p];
    delete out_deriv_batch[p];
  }
}


// update parameters
// see function declaration in nnet-simple-component.h for details
void ConvolutionComponent::Update(const std::string &debug_info,
                                  const CuMatrixBase<BaseFloat> &in_value,
                                  const CuMatrixBase<BaseFloat> &out_deriv,
                                  const std::vector<CuSubMatrix<BaseFloat> *>& out_deriv_batch) {
  // useful dims
  const int32 num_x_steps = (1 + (input_x_dim_ - filt_x_dim_) / filt_x_step_),
              num_y_steps = (1 + (input_y_dim_ - filt_y_dim_) / filt_y_step_),
              num_filters = filter_params_.NumRows(),
              num_frames = out_deriv.NumRows(),
              filter_dim = filter_params_.NumCols();
  KALDI_ASSERT(out_deriv.NumRows() == num_frames &&
               out_deriv.NumCols() ==
               (num_filters * num_x_steps * num_y_steps));


  CuMatrix<BaseFloat> filters_grad;
  CuVector<BaseFloat> bias_grad;

  CuMatrix<BaseFloat> input_patches(num_frames,
                                    filter_dim * num_x_steps * num_y_steps,
                                    kUndefined);
  InputToInputPatches(in_value, &input_patches);

  filters_grad.Resize(num_filters, filter_dim, kSetZero); // reset
  bias_grad.Resize(num_filters, kSetZero); // reset

  // create a single large matrix holding the smaller matrices
  // from the vector container filters_grad_batch along the rows
  CuMatrix<BaseFloat> filters_grad_blocks_batch(
      num_x_steps * num_y_steps * filters_grad.NumRows(),
      filters_grad.NumCols());

  std::vector<CuSubMatrix<BaseFloat>* > filters_grad_batch, input_patch_batch;

  for (int32 x_step = 0; x_step < num_x_steps; x_step++)  {
    for (int32 y_step = 0; y_step < num_y_steps; y_step++)  {
      int32 patch_number = x_step * num_y_steps + y_step;
      filters_grad_batch.push_back(new CuSubMatrix<BaseFloat>(
          filters_grad_blocks_batch.RowRange(
              patch_number * filters_grad.NumRows(), filters_grad.NumRows())));

      input_patch_batch.push_back(new CuSubMatrix<BaseFloat>(
              input_patches.ColRange(patch_number * filter_dim, filter_dim)));
    }
  }

  AddMatMatBatched<BaseFloat>(1.0, filters_grad_batch, out_deriv_batch, kTrans,
                              input_patch_batch, kNoTrans, 1.0);

  // add the row blocks together to filters_grad
  filters_grad.AddMatBlocks(1.0, filters_grad_blocks_batch);

  // create a matrix holding the col blocks sum of out_deriv
  CuMatrix<BaseFloat> out_deriv_col_blocks_sum(out_deriv.NumRows(),
                                               num_filters);

  // add the col blocks together to out_deriv_col_blocks_sum
  out_deriv_col_blocks_sum.AddMatBlocks(1.0, out_deriv);

  bias_grad.AddRowSumMat(1.0, out_deriv_col_blocks_sum, 1.0);

  // release memory
  for (int32 p = 0; p < input_patch_batch.size(); p++) {
    delete filters_grad_batch[p];
    delete input_patch_batch[p];
  }

  //
  // update
  //
  filter_params_.AddMat(learning_rate_, filters_grad);
  bias_params_.AddVec(learning_rate_, bias_grad);
}

void ConvolutionComponent::Read(std::istream &is, bool binary) {
  ReadUpdatableCommon(is, binary);  // Read opening tag and learning rate.
  ExpectToken(is, binary, "<InputXDim>");
  ReadBasicType(is, binary, &input_x_dim_);
  ExpectToken(is, binary, "<InputYDim>");
  ReadBasicType(is, binary, &input_y_dim_);
  ExpectToken(is, binary, "<InputZDim>");
  ReadBasicType(is, binary, &input_z_dim_);
  ExpectToken(is, binary, "<FiltXDim>");
  ReadBasicType(is, binary, &filt_x_dim_);
  ExpectToken(is, binary, "<FiltYDim>");
  ReadBasicType(is, binary, &filt_y_dim_);
  ExpectToken(is, binary, "<FiltXStep>");
  ReadBasicType(is, binary, &filt_x_step_);
  ExpectToken(is, binary, "<FiltYStep>");
  ReadBasicType(is, binary, &filt_y_step_);
  ExpectToken(is, binary, "<InputVectorization>");
  int32 input_vectorization;
  ReadBasicType(is, binary, &input_vectorization);
  input_vectorization_ = static_cast<TensorVectorizationType>(input_vectorization);
  ExpectToken(is, binary, "<FilterParams>");
  filter_params_.Read(is, binary);
  ExpectToken(is, binary, "<BiasParams>");
  bias_params_.Read(is, binary);
  std::string tok;
  ReadToken(is, binary, &tok);
  if (tok == "<IsGradient>") {
    ReadBasicType(is, binary, &is_gradient_);
    ExpectToken(is, binary, "</ConvolutionComponent>");
  } else {
    is_gradient_ = false;
    KALDI_ASSERT(tok == "</ConvolutionComponent>");
  }
}

void ConvolutionComponent::Write(std::ostream &os, bool binary) const {
  WriteUpdatableCommon(os, binary);  // write opening tag and learning rate.
  WriteToken(os, binary, "<InputXDim>");
  WriteBasicType(os, binary, input_x_dim_);
  WriteToken(os, binary, "<InputYDim>");
  WriteBasicType(os, binary, input_y_dim_);
  WriteToken(os, binary, "<InputZDim>");
  WriteBasicType(os, binary, input_z_dim_);
  WriteToken(os, binary, "<FiltXDim>");
  WriteBasicType(os, binary, filt_x_dim_);
  WriteToken(os, binary, "<FiltYDim>");
  WriteBasicType(os, binary, filt_y_dim_);
  WriteToken(os, binary, "<FiltXStep>");
  WriteBasicType(os, binary, filt_x_step_);
  WriteToken(os, binary, "<FiltYStep>");
  WriteBasicType(os, binary, filt_y_step_);
  WriteToken(os, binary, "<InputVectorization>");
  WriteBasicType(os, binary, static_cast<int32>(input_vectorization_));
  WriteToken(os, binary, "<FilterParams>");
  filter_params_.Write(os, binary);
  WriteToken(os, binary, "<BiasParams>");
  bias_params_.Write(os, binary);
  WriteToken(os, binary, "<IsGradient>");
  WriteBasicType(os, binary, is_gradient_);
  WriteToken(os, binary, "</ConvolutionComponent>");
}

BaseFloat ConvolutionComponent::DotProduct(const UpdatableComponent &other_in) const {
  const ConvolutionComponent *other =
      dynamic_cast<const ConvolutionComponent*>(&other_in);
  return TraceMatMat(filter_params_, other->filter_params_, kTrans)
         + VecVec(bias_params_, other->bias_params_);
}

Component* ConvolutionComponent::Copy() const {
  ConvolutionComponent *ans = new ConvolutionComponent(*this);
  return ans;
}

void ConvolutionComponent::PerturbParams(BaseFloat stddev) {
  CuMatrix<BaseFloat> temp_filter_params(filter_params_);
  temp_filter_params.SetRandn();
  filter_params_.AddMat(stddev, temp_filter_params);

  CuVector<BaseFloat> temp_bias_params(bias_params_);
  temp_bias_params.SetRandn();
  bias_params_.AddVec(stddev, temp_bias_params);
}

void ConvolutionComponent::SetParams(const VectorBase<BaseFloat> &bias,
                                     const MatrixBase<BaseFloat> &filter) {
  bias_params_ = bias;
  filter_params_ = filter;
  KALDI_ASSERT(bias_params_.Dim() == filter_params_.NumRows());
}

int32 ConvolutionComponent::NumParameters() const {
  return (filter_params_.NumCols() + 1) * filter_params_.NumRows();
}

void ConvolutionComponent::Vectorize(VectorBase<BaseFloat> *params) const {
  KALDI_ASSERT(params->Dim() == this->NumParameters());
  int32 num_filter_params = filter_params_.NumCols() * filter_params_.NumRows();
  params->Range(0, num_filter_params).CopyRowsFromMat(filter_params_);
  params->Range(num_filter_params, bias_params_.Dim()).CopyFromVec(bias_params_);
}
void ConvolutionComponent::UnVectorize(const VectorBase<BaseFloat> &params) {
  KALDI_ASSERT(params.Dim() == this->NumParameters());
  int32 num_filter_params = filter_params_.NumCols() * filter_params_.NumRows();
  filter_params_.CopyRowsFromVec(params.Range(0, num_filter_params));
  bias_params_.CopyFromVec(params.Range(num_filter_params, bias_params_.Dim()));
}

// aquire input dim
int32 MaxpoolingComponent::InputDim() const {
  return input_x_dim_ * input_y_dim_ * input_z_dim_;
}

MaxpoolingComponent::MaxpoolingComponent(
    const MaxpoolingComponent &component):
    input_x_dim_(component.input_x_dim_),
    input_y_dim_(component.input_y_dim_),
    input_z_dim_(component.input_z_dim_),
    pool_x_size_(component.pool_x_size_),
    pool_y_size_(component.pool_y_size_),
    pool_z_size_(component.pool_z_size_),
    pool_x_step_(component.pool_x_step_),
    pool_y_step_(component.pool_y_step_),
    pool_z_step_(component.pool_z_step_) { }

// aquire output dim
int32 MaxpoolingComponent::OutputDim() const {
  int32 num_pools_x = 1 + (input_x_dim_ - pool_x_size_) / pool_x_step_;
  int32 num_pools_y = 1 + (input_y_dim_ - pool_y_size_) / pool_y_step_;
  int32 num_pools_z = 1 + (input_z_dim_ - pool_z_size_) / pool_z_step_;
  return num_pools_x * num_pools_y * num_pools_z;
}

// check the component parameters
void MaxpoolingComponent::Check() const {
  // sanity check of the max pooling parameters
  KALDI_ASSERT(input_x_dim_ > 0);
  KALDI_ASSERT(input_y_dim_ > 0);
  KALDI_ASSERT(input_z_dim_ > 0);
  KALDI_ASSERT(pool_x_size_ > 0);
  KALDI_ASSERT(pool_y_size_ > 0);
  KALDI_ASSERT(pool_z_size_ > 0);
  KALDI_ASSERT(pool_x_step_ > 0);
  KALDI_ASSERT(pool_y_step_ > 0);
  KALDI_ASSERT(pool_z_step_ > 0);
  KALDI_ASSERT(input_x_dim_ >= pool_x_size_);
  KALDI_ASSERT(input_y_dim_ >= pool_y_size_);
  KALDI_ASSERT(input_z_dim_ >= pool_z_size_);
  KALDI_ASSERT(pool_x_size_ >= pool_x_step_);
  KALDI_ASSERT(pool_y_size_ >= pool_y_step_);
  KALDI_ASSERT(pool_z_size_ >= pool_z_step_);
  KALDI_ASSERT((input_x_dim_ - pool_x_size_) % pool_x_step_  == 0);
  KALDI_ASSERT((input_y_dim_ - pool_y_size_) % pool_y_step_  == 0);
  KALDI_ASSERT((input_z_dim_ - pool_z_size_) % pool_z_step_  == 0);
}

// initialize the component using configuration file
void MaxpoolingComponent::InitFromConfig(ConfigLine *cfl) {
  bool ok = true;

  ok = ok && cfl->GetValue("input-x-dim", &input_x_dim_);
  ok = ok && cfl->GetValue("input-y-dim", &input_y_dim_);
  ok = ok && cfl->GetValue("input-z-dim", &input_z_dim_);
  ok = ok && cfl->GetValue("pool-x-size", &pool_x_size_);
  ok = ok && cfl->GetValue("pool-y-size", &pool_y_size_);
  ok = ok && cfl->GetValue("pool-z-size", &pool_z_size_);
  ok = ok && cfl->GetValue("pool-x-step", &pool_x_step_);
  ok = ok && cfl->GetValue("pool-y-step", &pool_y_step_);
  ok = ok && cfl->GetValue("pool-z-step", &pool_z_step_);

  if (cfl->HasUnusedValues())
    KALDI_ERR << "Could not process these elements in initializer: "
              << cfl->UnusedValues();
  if (!ok)
    KALDI_ERR << "Bad initializer " << cfl->WholeLine();

  Check();
}

// Method to convert from a matrix representing a minibatch of vectorized
// 3D tensors to patches for 3d max pooling, each patch corresponds to
// the nodes having the same local coordinatenodes from each pool
void MaxpoolingComponent::InputToInputPatches(
    const CuMatrixBase<BaseFloat>& in,
    CuMatrix<BaseFloat> *patches) const{
  int32 num_pools_x = 1 + (input_x_dim_ - pool_x_size_) / pool_x_step_;
  int32 num_pools_y = 1 + (input_y_dim_ - pool_y_size_) / pool_y_step_;
  int32 num_pools_z = 1 + (input_z_dim_ - pool_z_size_) / pool_z_step_;

  std::vector<int32> column_map(patches->NumCols());
  int32 column_map_size = column_map.size();
  for (int32 x = 0, index =0; x < pool_x_size_; x++) {
    for (int32 y = 0; y < pool_y_size_; y++) {
      for (int32 z = 0; z < pool_z_size_; z++) {
        // given the local node coordinate, group them from each pool
        // to form a patch
        for (int32 x_pool = 0; x_pool < num_pools_x; x_pool++) {
          for (int32 y_pool = 0; y_pool < num_pools_y; y_pool++) {
            for (int32 z_pool = 0; z_pool < num_pools_z; z_pool++, index++) {
              KALDI_ASSERT(index < column_map_size);
              column_map[index] = (x_pool * pool_x_step_ + x) * input_y_dim_ * input_z_dim_ +
                                  (y_pool * pool_y_step_ + y) * input_z_dim_ +
                                  (z_pool * pool_z_step_ + z);

            }
          }
        }
      }
    }
  }
  CuArray<int32> cu_cols(column_map);
  patches->CopyCols(in, cu_cols);
}

/*
  This is the 3d max pooling propagate function.
  It is assumed that each row of the input matrix
  is a vectorized 3D-tensor of type zxy.
  Similar to the propagate function of ConvolutionComponent,
  the input matrix is first arranged into patches so that
  pools (with / without overlapping) could be
  processed in a parallelizable manner.
  The output matrix is also a vectorized 3D-tensor of type zxy.
*/

void* MaxpoolingComponent::Propagate(const ComponentPrecomputedIndexes *indexes,
                                    const CuMatrixBase<BaseFloat> &in,
                                    CuMatrixBase<BaseFloat> *out) const {
  int32 num_frames = in.NumRows();
  int32 num_pools = OutputDim();
  int32 pool_size = pool_x_size_ * pool_y_size_ * pool_z_size_;
  CuMatrix<BaseFloat> patches(num_frames, num_pools * pool_size, kUndefined);
  InputToInputPatches(in, &patches);

  out->Set(-1e20); // reset a large negative value
  for (int32 q = 0; q < pool_size; q++)
    out->Max(patches.ColRange(q * num_pools, num_pools));
  return NULL;
}

// Method to compute the input derivative matrix from the input derivatives
// for patches, where each patch corresponds to
// the nodes having the same local coordinatenodes from each pool
void MaxpoolingComponent::InderivPatchesToInderiv(
    const CuMatrix<BaseFloat>& in_deriv_patches,
    CuMatrixBase<BaseFloat> *in_deriv) const {

  int32 num_pools_x = 1 + (input_x_dim_ - pool_x_size_) / pool_x_step_;
  int32 num_pools_y = 1 + (input_y_dim_ - pool_y_size_) / pool_y_step_;
  int32 num_pools_z = 1 + (input_z_dim_ - pool_z_size_) / pool_z_step_;

  std::vector<std::vector<int32> > reverse_column_map(in_deriv->NumCols());
  int32 rev_col_map_size = reverse_column_map.size();
  for (int32 x = 0, index = 0; x < pool_x_size_; x++) {
    for (int32 y = 0; y < pool_y_size_; y++) {
      for (int32 z = 0; z < pool_z_size_; z++) {

        for (int32 x_pool = 0; x_pool < num_pools_x; x_pool++) {
          for (int32 y_pool = 0; y_pool < num_pools_y; y_pool++) {
            for (int32 z_pool = 0; z_pool < num_pools_z; z_pool++, index++) {
              int32 vector_index = (x_pool * pool_x_step_ + x) * input_y_dim_ * input_z_dim_ +
                                  (y_pool * pool_y_step_ + y) * input_z_dim_ +
                                  (z_pool * pool_z_step_ + z);

              KALDI_ASSERT(vector_index < rev_col_map_size);
              reverse_column_map[vector_index].push_back(index);
            }
          }
        }
      }
    }
  }
  std::vector<std::vector<int32> > rearranged_column_map;
  RearrangeIndexes(reverse_column_map, &rearranged_column_map);
  for (int32 p = 0; p < rearranged_column_map.size(); p++) {
    CuArray<int32> cu_cols(rearranged_column_map[p]);
    in_deriv->AddCols(in_deriv_patches, cu_cols);
  }
}

/*
  3d max pooling backpropagate function
  This function backpropagate the error from
  out_deriv to in_deriv.
  In order to select the node in each pool to
  backpropagate the error, it has to compare
  the output pool value stored in the out_value
  matrix with each of its input pool member node
  stroed in the in_value matrix.
*/
void MaxpoolingComponent::Backprop(const std::string &debug_info,
                                   const ComponentPrecomputedIndexes *indexes,
                                   const CuMatrixBase<BaseFloat> &in_value,
                                   const CuMatrixBase<BaseFloat> &out_value,
                                   const CuMatrixBase<BaseFloat> &out_deriv,
                                   void *memo,
                                   Component *, // to_update,
                                   CuMatrixBase<BaseFloat> *in_deriv) const {
  if (!in_deriv)
    return;

  int32 num_frames = in_value.NumRows();
  int32 num_pools = OutputDim();
  int32 pool_size = pool_x_size_ * pool_y_size_ * pool_z_size_;
  CuMatrix<BaseFloat> patches(num_frames, num_pools * pool_size, kUndefined);
  InputToInputPatches(in_value, &patches);

  for (int32 q = 0; q < pool_size; q++) {
    // zero-out mask
    CuMatrix<BaseFloat> mask;
    out_value.EqualElementMask(patches.ColRange(q * num_pools, num_pools), &mask);
    mask.MulElements(out_deriv);
    patches.ColRange(q * num_pools, num_pools).CopyFromMat(mask);
  }

  // combine the derivatives from the individual input deriv patches
  // to compute input deriv matrix
  InderivPatchesToInderiv(patches, in_deriv);
}

void MaxpoolingComponent::Read(std::istream &is, bool binary) {
  ExpectOneOrTwoTokens(is, binary, "<MaxpoolingComponent>", "<InputXDim>");
  ReadBasicType(is, binary, &input_x_dim_);
  ExpectToken(is, binary, "<InputYDim>");
  ReadBasicType(is, binary, &input_y_dim_);
  ExpectToken(is, binary, "<InputZDim>");
  ReadBasicType(is, binary, &input_z_dim_);
  ExpectToken(is, binary, "<PoolXSize>");
  ReadBasicType(is, binary, &pool_x_size_);
  ExpectToken(is, binary, "<PoolYSize>");
  ReadBasicType(is, binary, &pool_y_size_);
  ExpectToken(is, binary, "<PoolZSize>");
  ReadBasicType(is, binary, &pool_z_size_);
  ExpectToken(is, binary, "<PoolXStep>");
  ReadBasicType(is, binary, &pool_x_step_);
  ExpectToken(is, binary, "<PoolYStep>");
  ReadBasicType(is, binary, &pool_y_step_);
  ExpectToken(is, binary, "<PoolZStep>");
  ReadBasicType(is, binary, &pool_z_step_);
  ExpectToken(is, binary, "</MaxpoolingComponent>");
  Check();
}

void MaxpoolingComponent::Write(std::ostream &os, bool binary) const {
  WriteToken(os, binary, "<MaxpoolingComponent>");
  WriteToken(os, binary, "<InputXDim>");
  WriteBasicType(os, binary, input_x_dim_);
  WriteToken(os, binary, "<InputYDim>");
  WriteBasicType(os, binary, input_y_dim_);
  WriteToken(os, binary, "<InputZDim>");
  WriteBasicType(os, binary, input_z_dim_);
  WriteToken(os, binary, "<PoolXSize>");
  WriteBasicType(os, binary, pool_x_size_);
  WriteToken(os, binary, "<PoolYSize>");
  WriteBasicType(os, binary, pool_y_size_);
  WriteToken(os, binary, "<PoolZSize>");
  WriteBasicType(os, binary, pool_z_size_);
  WriteToken(os, binary, "<PoolXStep>");
  WriteBasicType(os, binary, pool_x_step_);
  WriteToken(os, binary, "<PoolYStep>");
  WriteBasicType(os, binary, pool_y_step_);
  WriteToken(os, binary, "<PoolZStep>");
  WriteBasicType(os, binary, pool_z_step_);
  WriteToken(os, binary, "</MaxpoolingComponent>");
}

// display information about component
std::string MaxpoolingComponent::Info() const {
  std::ostringstream stream;
  stream << Type()
         << ", input-x-dim=" << input_x_dim_
         << ", input-y-dim=" << input_y_dim_
         << ", input-z-dim=" << input_z_dim_
         << ", pool-x-size=" << pool_x_size_
         << ", pool-y-size=" << pool_y_size_
         << ", pool-z-size=" << pool_z_size_
         << ", pool-x-step=" << pool_x_step_
         << ", pool-y-step=" << pool_y_step_
         << ", pool-z-step=" << pool_z_step_;
  return stream.str();
}


int32 LstmNonlinearityComponent::InputDim() const {
  int32 cell_dim = value_sum_.NumCols();
  return cell_dim * 5 + (use_dropout_ ? 3 : 0);
}

int32 LstmNonlinearityComponent::OutputDim() const {
  int32 cell_dim = value_sum_.NumCols();
  return cell_dim * 2;
}


void LstmNonlinearityComponent::Read(std::istream &is, bool binary) {
  ReadUpdatableCommon(is, binary);  // Read opening tag and learning rate.
  ExpectToken(is, binary, "<Params>");
  params_.Read(is, binary);
  ExpectToken(is, binary, "<ValueAvg>");
  value_sum_.Read(is, binary);
  ExpectToken(is, binary, "<DerivAvg>");
  deriv_sum_.Read(is, binary);
  ExpectToken(is, binary, "<SelfRepairConfig>");
  self_repair_config_.Read(is, binary);
  ExpectToken(is, binary, "<SelfRepairProb>");
  self_repair_total_.Read(is, binary);

  std::string tok;
  ReadToken(is, binary, &tok);
  if (tok == "<UseDropout>") {
    ReadBasicType(is, binary, &use_dropout_);
    ReadToken(is, binary, &tok);
  } else {
    use_dropout_ = false;
  }
  KALDI_ASSERT(tok == "<Count>");
  ReadBasicType(is, binary, &count_);

  // For the on-disk format, we normalze value_sum_, deriv_sum_ and
  // self_repair_total_ by dividing by the count, but in memory they are scaled
  // by the count.  [for self_repair_total_, the scaling factor is count_ *
  // cell_dim].
  value_sum_.Scale(count_);
  deriv_sum_.Scale(count_);
  int32 cell_dim = params_.NumCols();
  self_repair_total_.Scale(count_ * cell_dim);

  InitNaturalGradient();

  ExpectToken(is, binary, "</LstmNonlinearityComponent>");

}

void LstmNonlinearityComponent::Write(std::ostream &os, bool binary) const {
  WriteUpdatableCommon(os, binary);  // Read opening tag and learning rate.

  WriteToken(os, binary, "<Params>");
  params_.Write(os, binary);
  WriteToken(os, binary, "<ValueAvg>");
  {
    Matrix<BaseFloat> value_avg(value_sum_);
    if (count_ != 0.0)
      value_avg.Scale(1.0 / count_);
    value_avg.Write(os, binary);
  }
  WriteToken(os, binary, "<DerivAvg>");
  {
    Matrix<BaseFloat> deriv_avg(deriv_sum_);
    if (count_ != 0.0)
      deriv_avg.Scale(1.0 / count_);
    deriv_avg.Write(os, binary);
  }
  WriteToken(os, binary, "<SelfRepairConfig>");
  self_repair_config_.Write(os, binary);
  WriteToken(os, binary, "<SelfRepairProb>");
  {
    int32 cell_dim = params_.NumCols();
    Vector<BaseFloat> self_repair_prob(self_repair_total_);
    if (count_ != 0.0)
      self_repair_prob.Scale(1.0 / (count_ * cell_dim));
    self_repair_prob.Write(os, binary);
  }
  if (use_dropout_) {
    // only write this if true; we have back-compat code in reading anyway.
    // this makes the models without dropout easier to read with older code.
    WriteToken(os, binary, "<UseDropout>");
    WriteBasicType(os, binary, use_dropout_);
  }
  WriteToken(os, binary, "<Count>");
  WriteBasicType(os, binary, count_);
  WriteToken(os, binary, "</LstmNonlinearityComponent>");
}



std::string LstmNonlinearityComponent::Info() const {
  std::ostringstream stream;
  int32 cell_dim = params_.NumCols();
  stream << UpdatableComponent::Info() << ", cell-dim=" << cell_dim
         << ", use-dropout=" << (use_dropout_ ? "true" : "false");
  PrintParameterStats(stream, "w_ic", params_.Row(0));
  PrintParameterStats(stream, "w_fc", params_.Row(1));
  PrintParameterStats(stream, "w_oc", params_.Row(2));

  // Note: some of the following code mirrors the code in
  // UpdatableComponent::Info(), in nnet-component-itf.cc.
  if (count_ > 0) {
    stream << ", count=" << std::setprecision(3) << count_
           << std::setprecision(6);
  }
  static const char *nonlin_names[] = { "i_t_sigmoid", "f_t_sigmoid", "c_t_tanh",
                                        "o_t_sigmoid", "m_t_tanh" };
  for (int32 i = 0; i < 5; i++) {
    stream << ", " << nonlin_names[i] << "={";
    stream << " self-repair-lower-threshold=" << self_repair_config_(i)
           << ", self-repair-scale=" << self_repair_config_(i + 5);

    if (count_ != 0) {
      BaseFloat self_repaired_proportion =
          self_repair_total_(i) / (count_ * cell_dim);
      stream << ", self-repaired-proportion=" << self_repaired_proportion;
      Vector<double> value_sum(value_sum_.Row(i)),
          deriv_sum(deriv_sum_.Row(i));
      Vector<BaseFloat> value_avg(value_sum), deriv_avg(deriv_sum);
      value_avg.Scale(1.0 / count_);
      deriv_avg.Scale(1.0 / count_);
      stream << ", value-avg=" << SummarizeVector(value_avg)
             << ", deriv-avg=" << SummarizeVector(deriv_avg);
    }
    stream << " }";
  }
  return stream.str();
}


Component* LstmNonlinearityComponent::Copy() const {
  return new LstmNonlinearityComponent(*this);
}

void LstmNonlinearityComponent::ZeroStats() {
  value_sum_.SetZero();
  deriv_sum_.SetZero();
  self_repair_total_.SetZero();
  count_ = 0.0;
}

void LstmNonlinearityComponent::Scale(BaseFloat scale) {
  if (scale == 0.0) {
    params_.SetZero();
    value_sum_.SetZero();
    deriv_sum_.SetZero();
    self_repair_total_.SetZero();
    count_ = 0.0;
  } else {
    params_.Scale(scale);
    value_sum_.Scale(scale);
    deriv_sum_.Scale(scale);
    self_repair_total_.Scale(scale);
    count_ *= scale;
  }
}

void LstmNonlinearityComponent::Add(BaseFloat alpha,
                                    const Component &other_in) {
  const LstmNonlinearityComponent *other =
      dynamic_cast<const LstmNonlinearityComponent*>(&other_in);
  KALDI_ASSERT(other != NULL);
  params_.AddMat(alpha, other->params_);
  value_sum_.AddMat(alpha, other->value_sum_);
  deriv_sum_.AddMat(alpha, other->deriv_sum_);
  self_repair_total_.AddVec(alpha, other->self_repair_total_);
  count_ += alpha * other->count_;
}

void LstmNonlinearityComponent::PerturbParams(BaseFloat stddev) {
  CuMatrix<BaseFloat> temp_params(params_.NumRows(), params_.NumCols());
  temp_params.SetRandn();
  params_.AddMat(stddev, temp_params);
}

BaseFloat LstmNonlinearityComponent::DotProduct(
    const UpdatableComponent &other_in) const {
  const LstmNonlinearityComponent *other =
      dynamic_cast<const LstmNonlinearityComponent*>(&other_in);
  KALDI_ASSERT(other != NULL);
  return TraceMatMat(params_, other->params_, kTrans);
}

int32 LstmNonlinearityComponent::NumParameters() const {
  return params_.NumRows() * params_.NumCols();
}

void LstmNonlinearityComponent::Vectorize(VectorBase<BaseFloat> *params) const {
  KALDI_ASSERT(params->Dim() == NumParameters());
  params->CopyRowsFromMat(params_);
}


void LstmNonlinearityComponent::UnVectorize(
    const VectorBase<BaseFloat> &params)  {
  KALDI_ASSERT(params.Dim() == NumParameters());
  params_.CopyRowsFromVec(params);
}


void* LstmNonlinearityComponent::Propagate(
    const ComponentPrecomputedIndexes *, // indexes
    const CuMatrixBase<BaseFloat> &in,
    CuMatrixBase<BaseFloat> *out) const {
  cu::ComputeLstmNonlinearity(in, params_, out);
  return NULL;
}


void LstmNonlinearityComponent::Backprop(
    const std::string &debug_info,
    const ComponentPrecomputedIndexes *indexes,
    const CuMatrixBase<BaseFloat> &in_value,
    const CuMatrixBase<BaseFloat> &, // out_value,
    const CuMatrixBase<BaseFloat> &out_deriv,
    void *memo,
    Component *to_update_in,
    CuMatrixBase<BaseFloat> *in_deriv) const {

  if (to_update_in == NULL) {
    cu::BackpropLstmNonlinearity(in_value, params_, out_deriv,
                                 deriv_sum_, self_repair_config_,
                                 count_, in_deriv,
                                 (CuMatrixBase<BaseFloat>*) NULL,
                                 (CuMatrixBase<double>*) NULL,
                                 (CuMatrixBase<double>*) NULL,
                                 (CuMatrixBase<BaseFloat>*) NULL);
  } else {
    LstmNonlinearityComponent *to_update =
        dynamic_cast<LstmNonlinearityComponent*>(to_update_in);
    KALDI_ASSERT(to_update != NULL);

    int32 cell_dim = params_.NumCols();
    CuMatrix<BaseFloat> params_deriv(3, cell_dim, kUndefined);
    CuMatrix<BaseFloat> self_repair_total(5, cell_dim, kUndefined);

    cu::BackpropLstmNonlinearity(in_value, params_, out_deriv,
                                 deriv_sum_, self_repair_config_,
                                 count_, in_deriv, &params_deriv,
                                 &(to_update->value_sum_),
                                 &(to_update->deriv_sum_),
                                 &self_repair_total);

    CuVector<BaseFloat> self_repair_total_sum(5);
    self_repair_total_sum.AddColSumMat(1.0, self_repair_total, 0.0);
    to_update->self_repair_total_.AddVec(1.0, self_repair_total_sum);
    to_update->count_ += static_cast<double>(in_value.NumRows());

    BaseFloat scale = 1.0;
    if (!to_update->is_gradient_) {
      to_update->preconditioner_.PreconditionDirections(
          &params_deriv, &scale);
    }
    to_update->params_.AddMat(to_update->learning_rate_ * scale,
                              params_deriv);
  }
}

LstmNonlinearityComponent::LstmNonlinearityComponent(
    const LstmNonlinearityComponent &other):
    UpdatableComponent(other),
    params_(other.params_),
    use_dropout_(other.use_dropout_),
    value_sum_(other.value_sum_),
    deriv_sum_(other.deriv_sum_),
    self_repair_config_(other.self_repair_config_),
    self_repair_total_(other.self_repair_total_),
    count_(other.count_),
    preconditioner_(other.preconditioner_) { }

void LstmNonlinearityComponent::Init(
    int32 cell_dim, bool use_dropout,
    BaseFloat param_stddev,
    BaseFloat tanh_self_repair_threshold,
    BaseFloat sigmoid_self_repair_threshold,
    BaseFloat self_repair_scale) {
  KALDI_ASSERT(cell_dim > 0 && param_stddev >= 0.0 &&
               tanh_self_repair_threshold >= 0.0 &&
               tanh_self_repair_threshold <= 1.0 &&
               sigmoid_self_repair_threshold >= 0.0 &&
               sigmoid_self_repair_threshold <= 0.25 &&
               self_repair_scale >= 0.0 && self_repair_scale <= 0.1);
  use_dropout_ = use_dropout;
  params_.Resize(3, cell_dim);
  params_.SetRandn();
  params_.Scale(param_stddev);
  value_sum_.Resize(5, cell_dim);
  deriv_sum_.Resize(5, cell_dim);
  self_repair_config_.Resize(10);
  self_repair_config_.Range(0, 5).Set(sigmoid_self_repair_threshold);
  self_repair_config_(2) = tanh_self_repair_threshold;
  self_repair_config_(4) = tanh_self_repair_threshold;
  self_repair_config_.Range(5, 5).Set(self_repair_scale);
  self_repair_total_.Resize(5);
  count_ = 0.0;
  InitNaturalGradient();

}

void LstmNonlinearityComponent::InitNaturalGradient() {
  // As regards the configuration for the natural-gradient preconditioner, we
  // don't make it configurable from the command line-- it's unlikely that any
  // differences from changing this would be substantial enough to effectively
  // tune the configuration.  Because the preconditioning code doesn't 'see' the
  // derivatives from individual frames, but only averages over the minibatch,
  // there is a fairly small amount of data available to estimate the Fisher
  // information matrix, so we set the rank, update period and
  // num-samples-history to smaller values than normal.
  preconditioner_.SetRank(20);
  preconditioner_.SetUpdatePeriod(2);
  preconditioner_.SetNumSamplesHistory(1000.0);
}

/// virtual
void LstmNonlinearityComponent::FreezeNaturalGradient(bool freeze) {
  preconditioner_.Freeze(freeze);
}

void LstmNonlinearityComponent::InitFromConfig(ConfigLine *cfl) {
  InitLearningRatesFromConfig(cfl);
  bool ok = true;
  bool use_dropout = false;
  int32 cell_dim;
  // these self-repair thresholds are the normal defaults for tanh and sigmoid
  // respectively.  If, later on, we decide that we want to support different
  // self-repair config values for the individual sigmoid and tanh
  // nonlinearities, we can modify this code then.
  BaseFloat tanh_self_repair_threshold = 0.2,
      sigmoid_self_repair_threshold = 0.05,
      self_repair_scale = 1.0e-05;
  // param_stddev is the stddev of the parameters.  it may be better to
  // use a smaller value but this was the default in the python scripts
  // for a while.
  BaseFloat param_stddev = 1.0;
  ok = ok && cfl->GetValue("cell-dim", &cell_dim);
  cfl->GetValue("param-stddev", &param_stddev);
  cfl->GetValue("tanh-self-repair-threshold",
                &tanh_self_repair_threshold);
  cfl->GetValue("sigmoid-self-repair-threshold",
                &sigmoid_self_repair_threshold);
  cfl->GetValue("self-repair-scale", &self_repair_scale);
  cfl->GetValue("use-dropout", &use_dropout);

  // We may later on want to make it possible to initialize the different
  // parameters w_ic, w_fc and w_oc with different biases.  We'll implement
  // that when and if it's needed.

  if (cfl->HasUnusedValues())
    KALDI_ERR << "Could not process these elements in initializer: "
              << cfl->UnusedValues();
  if (ok) {
    Init(cell_dim, use_dropout, param_stddev, tanh_self_repair_threshold,
         sigmoid_self_repair_threshold, self_repair_scale);
  } else {
    KALDI_ERR << "Invalid initializer for layer of type "
              << Type() << ": \"" << cfl->WholeLine() << "\"";
  }
}

void LstmNonlinearityComponent::ConsolidateMemory() {
  OnlineNaturalGradient preconditioner_temp(preconditioner_);
  preconditioner_.Swap(&preconditioner_);
}


int32 GruNonlinearityComponent::InputDim() const {
  if (recurrent_dim_ == cell_dim_) {
    // non-projected GRU.
    return 4 * cell_dim_;
  } else {
    return 3 * cell_dim_ + 2 * recurrent_dim_;
  }
}

int32 GruNonlinearityComponent::OutputDim() const {
  return 2 * cell_dim_;
}


std::string GruNonlinearityComponent::Info() const {
  std::ostringstream stream;
  stream << UpdatableComponent::Info()
         << ", cell-dim=" << cell_dim_
         << ", recurrent-dim=" << recurrent_dim_;
  PrintParameterStats(stream, "w_h", w_h_);
  stream << ", self-repair-threshold=" << self_repair_threshold_
         << ", self-repair-scale=" << self_repair_scale_;
  if (count_ > 0) {  // c.f. NonlinearComponent::Info().
    stream << ", count=" << std::setprecision(3) << count_
           << std::setprecision(6);
    stream << ", self-repaired-proportion="
           << (self_repair_total_ / (count_ * cell_dim_));
    Vector<double> value_avg_dbl(value_sum_);
    Vector<BaseFloat> value_avg(value_avg_dbl);
    value_avg.Scale(1.0 / count_);
    stream << ", value-avg=" << SummarizeVector(value_avg);
    Vector<double> deriv_avg_dbl(deriv_sum_);
    Vector<BaseFloat> deriv_avg(deriv_avg_dbl);
    deriv_avg.Scale(1.0 / count_);
    stream << ", deriv-avg=" << SummarizeVector(deriv_avg);
  }
  // natural-gradient parameters.
  stream << ", alpha=" << preconditioner_in_.GetAlpha()
         << ", rank-in=" << preconditioner_in_.GetRank()
         << ", rank-out=" << preconditioner_out_.GetRank()
         << ", update-period="
         << preconditioner_in_.GetUpdatePeriod();
  return stream.str();
}

void GruNonlinearityComponent::InitFromConfig(ConfigLine *cfl) {
  cell_dim_ = -1;
  recurrent_dim_ = -1;
  self_repair_threshold_ = 0.2;
  self_repair_scale_ = 1.0e-05;

  InitLearningRatesFromConfig(cfl);
  if (!cfl->GetValue("cell-dim", &cell_dim_) || cell_dim_ <= 0)
    KALDI_ERR << "cell-dim > 0 is required for GruNonlinearityComponent.";

  BaseFloat param_stddev = 1.0 / std::sqrt(cell_dim_),
      alpha = 4.0;
  int32 rank_in = 20, rank_out = 80,
      update_period = 4;

  cfl->GetValue("recurrent-dim", &recurrent_dim_);
  cfl->GetValue("self-repair-threshold", &self_repair_threshold_);
  cfl->GetValue("self-repair-scale", &self_repair_scale_);
  cfl->GetValue("param-stddev", &param_stddev);
  cfl->GetValue("alpha", &alpha);
  cfl->GetValue("rank-in", &rank_in);
  cfl->GetValue("rank-out", &rank_out);
  cfl->GetValue("update-period", &update_period);

  if (recurrent_dim_ < 0)
    recurrent_dim_ = cell_dim_;
  if (recurrent_dim_ == 0 || recurrent_dim_ > cell_dim_)
    KALDI_ERR << "Invalid values for cell-dim and recurrent-dim";

  w_h_.Resize(cell_dim_, recurrent_dim_);
  w_h_.SetRandn();
  w_h_.Scale(param_stddev);

  preconditioner_in_.SetAlpha(alpha);
  preconditioner_in_.SetRank(rank_in);
  preconditioner_in_.SetUpdatePeriod(update_period);
  preconditioner_out_.SetAlpha(alpha);
  preconditioner_out_.SetRank(rank_out);
  preconditioner_out_.SetUpdatePeriod(update_period);

  count_ = 0.0;
  self_repair_total_ = 0.0;
  value_sum_.Resize(cell_dim_);
  deriv_sum_.Resize(cell_dim_);

  Check();
}

void* GruNonlinearityComponent::Propagate(
    const ComponentPrecomputedIndexes *indexes,
    const CuMatrixBase<BaseFloat> &in,
    CuMatrixBase<BaseFloat> *out) const {
  KALDI_ASSERT(in.NumRows() == out->NumRows() &&
               in.NumCols() == InputDim() &&
               out->NumCols() == OutputDim());
  // If recurrent_dim_ != cell_dim_, this is projected GRU and we
  // are computing:
  //  (z_t, r_t, hpart_t, c_{t-1}, s_{t-1}) -> (h_t, c_t).
  // Otherwise (no projection), it's
  //  (z_t, r_t, hpart_t, y_{t-1},) -> (h_t, y_t).
  // but to understand this code, it's better to rename y to c:
  //  (z_t, r_t, hpart_t, c_{t-1}) -> (h_t, c_t).
  int32 num_rows = in.NumRows(),
      c = cell_dim_,
      r =  recurrent_dim_;
  CuSubMatrix<BaseFloat> z_t(in, 0, num_rows, 0, c),
      r_t(in, 0, num_rows, c, r),
      hpart_t(in, 0, num_rows, c + r, c),
      c_t1(in, 0, num_rows, c + r + c, c);
  // note: the variable named 'c_t1' actually represents
  // y_{t-1} for non-projected GRUs.

  // By setting s_t1 to the last recurrent_dim_ rows of 'in', we get something
  // that represents s_{t-1} for recurrent setups and y_{t-1} (which we're
  // renaming to c_{t-1}) for non-projected GRUs.  The key thing is that
  // in the non-projected case, the variables c_t1 and s_t1 point to the
  // same memory.
  CuSubMatrix<BaseFloat> s_t1(in, 0, num_rows, in.NumCols() - r, r);

  // note: for non-projected GRUs, c_t below is actually y_t.
  CuSubMatrix<BaseFloat> h_t(*out, 0, num_rows, 0, c),
      c_t(*out, 0, num_rows, c, c);

  // sdotr is the only temporary storage we need in the forward pass.
  CuMatrix<BaseFloat> sdotr(num_rows, r);
  sdotr.AddMatMatElements(1.0, r_t, s_t1, 0.0);
  // now sdotr = r_t \dot s_{t-1}.
  h_t.CopyFromMat(hpart_t);
  // now h_t = hpart_t (note: hpart_t actually means U^h x_t).
  h_t.AddMatMat(1.0, sdotr, kNoTrans, w_h_, kTrans, 1.0);
  // now h_t = hpart_t + W^h (s_{t-1} \dot r_t).
  h_t.Tanh(h_t);
  // now, h_t = tanh(hpart_t + W^h (s_{t-1} \dot r_t)).

  c_t.CopyFromMat(h_t);
  // now c_t = h_t
  c_t.AddMatMatElements(-1.0, z_t, h_t, 1.0);
  // now c_t = (1 - z_t) \dot h_t.
  c_t.AddMatMatElements(1.0, z_t, c_t1, 1.0);
  // now c_t = (1 - z_t) \dot h_t  +  z_t \dot c_{t-1}.
  return NULL;
}

void GruNonlinearityComponent::Backprop(
    const std::string &debug_info,
    const ComponentPrecomputedIndexes *, // indexes
    const CuMatrixBase<BaseFloat> &in_value,
    const CuMatrixBase<BaseFloat> &out_value,
    const CuMatrixBase<BaseFloat> &out_deriv,
    void *memo,
    Component *to_update_in,
    CuMatrixBase<BaseFloat> *in_deriv) const {
  KALDI_ASSERT(SameDim(out_value, out_deriv) &&
               in_value.NumRows() == out_value.NumRows() &&
               in_value.NumCols() == InputDim() &&
               out_value.NumCols() == OutputDim() &&
               (in_deriv == NULL || SameDim(in_value, *in_deriv)) &&
               memo == NULL);
  GruNonlinearityComponent *to_update =
      dynamic_cast<GruNonlinearityComponent*>(to_update_in);
  KALDI_ASSERT(in_deriv != NULL || to_update != NULL);
  int32 num_rows = in_value.NumRows(),
      c = cell_dim_,
      r = recurrent_dim_;

  // To understand what's going on here, compare this code with the
  // corresponding 'forward' code in Propagate().


  CuSubMatrix<BaseFloat> z_t(in_value, 0, num_rows, 0, c),
      r_t(in_value, 0, num_rows, c, r),
      hpart_t(in_value, 0, num_rows, c + r, c),
      c_t1(in_value, 0, num_rows, c + r + c, c),
      s_t1(in_value, 0, num_rows, in_value.NumCols() - r, r);


  // The purpose of this 'in_deriv_ptr' is so that we can create submatrices
  // like z_t_deriv without the code crashing.  If in_deriv is NULL these point
  // to 'in_value', and we'll be careful never to actually write to these
  // sub-matrices, which aside from being conceptually wrong would violate the
  // const semantics of this function.
  const CuMatrixBase<BaseFloat> *in_deriv_ptr =
      (in_deriv == NULL ? &in_value : in_deriv);
  CuSubMatrix<BaseFloat> z_t_deriv(*in_deriv_ptr, 0, num_rows, 0, c),
      r_t_deriv(*in_deriv_ptr, 0, num_rows, c, r),
      hpart_t_deriv(*in_deriv_ptr, 0, num_rows, c + r, c),
      c_t1_deriv(*in_deriv_ptr, 0, num_rows, c + r + c, c),
      s_t1_deriv(*in_deriv_ptr, 0, num_rows, in_value.NumCols() - r, r);

  // Note: the output h_t is never actually used in the GRU computation (we only
  // output it because we want the value to be cached to save computation in the
  // backprop), so we expect that the 'h_t_deriv', if we extracted it in the
  // obvious way, would be all zeros.
  // We create a different, local h_t_deriv
  // variable that backpropagates the derivative from c_t_deriv.
  CuSubMatrix<BaseFloat> h_t(out_value, 0, num_rows, 0, c),
      c_t(out_value, 0, num_rows, c, c),
      c_t_deriv(out_deriv, 0, num_rows, c, c);
  CuMatrix<BaseFloat> h_t_deriv(num_rows, c, kUndefined);

  {  // we initialize h_t_deriv with the derivative from 'out_deriv'.
    // In real life in a GRU, this would always be zero; but in testing
    // code it may be nonzero and we include this term so that
    // the tests don't fail.  Note: if you were to remove these
    // lines, you'd have to change 'h_t_deriv.AddMat(1.0, c_t_deriv);' below
    // to a CopyFromMat() call.
    CuSubMatrix<BaseFloat> h_t_deriv_in(out_deriv, 0, num_rows, 0, c);
    h_t_deriv.CopyFromMat(h_t_deriv_in);
  }


  // sdotr is the same variable as used in the forward pass, it will contain
  // r_t \dot s_{t-1}.
  CuMatrix<BaseFloat> sdotr(num_rows, r);
  sdotr.AddMatMatElements(1.0, r_t, s_t1, 0.0);


  { // This block does the
    // backprop corresponding to the
    // forward-pass expression: c_t = (1 - z_t) \dot h_t + z_t \dot c_{t-1}.

    // First do: h_t_deriv = c_t_deriv \dot (1 - z_t).
    h_t_deriv.AddMat(1.0, c_t_deriv);
    h_t_deriv.AddMatMatElements(-1.0, c_t_deriv, z_t, 1.0);

    if (in_deriv) {
      // these should be self-explanatory if you study
      // the expression "c_t = (1 - z_t) \dot h_t + z_t \dot c_{t-1}".
      z_t_deriv.AddMatMatElements(-1.0, c_t_deriv, h_t, 1.0);
      z_t_deriv.AddMatMatElements(1.0, c_t_deriv, c_t1, 1.0);
      c_t1_deriv.AddMatMatElements(1.0, c_t_deriv, z_t, 1.0);
    }
  }

  h_t_deriv.DiffTanh(h_t, h_t_deriv);
  if (to_update)
    to_update->TanhStatsAndSelfRepair(h_t, &h_t_deriv);


  if (to_update)
    to_update->UpdateParameters(sdotr, h_t_deriv);

  // At this point, 'h_t_deriv' contains the derivative w.r.t.
  // the argument of the tanh function, i.e. w.r.t. the expression:
  //    hpart_t + W^h (s_{t-1} \dot r_t).
  // The next block propagates this to the derivatives for
  // hpart_t, s_{t-1} and r_t.
  if (in_deriv) {
    hpart_t_deriv.AddMat(1.0, h_t_deriv);

    // We re-use the memory that we used for s_{t-1} \dot r_t,
    // for its derivative.
    CuMatrix<BaseFloat> &sdotr_deriv(sdotr);
    sdotr_deriv.AddMatMat(1.0, h_t_deriv, kNoTrans, w_h_, kNoTrans, 0.0);

    // we add to all the input-derivatives instead of setting them,
    // because we chose to export the flag kBackpropAdds.
    r_t_deriv.AddMatMatElements(1.0, sdotr_deriv, s_t1, 1.0);
    s_t1_deriv.AddMatMatElements(1.0, sdotr_deriv, r_t, 1.0);
  }
}


void GruNonlinearityComponent::TanhStatsAndSelfRepair(
    const CuMatrixBase<BaseFloat> &h_t,
    CuMatrixBase<BaseFloat> *h_t_deriv) {
  KALDI_ASSERT(SameDim(h_t, *h_t_deriv));

  // we use this probability (hardcoded for now) to limit the stats accumulation
  // and self-repair code to running on about half of the minibatches.
  BaseFloat repair_and_stats_probability = 0.5;
  if (RandUniform() > repair_and_stats_probability)
    return;

  // OK, accumulate stats.
  // For the next few lines, compare with TanhComponent::StoreStats(), which is where
  // we got this code.
  // tanh_deriv is the function derivative of the tanh function,
  // tanh'(x) = tanh(x) * (1.0 - tanh(x)).  h_t corresponds to tanh(x).
  CuMatrix<BaseFloat> tanh_deriv(h_t);
  tanh_deriv.ApplyPow(2.0);
  tanh_deriv.Scale(-1.0);
  tanh_deriv.Add(1.0);

  count_ += h_t.NumRows();
  CuVector<BaseFloat> temp(cell_dim_);
  temp.AddRowSumMat(1.0, h_t, 0.0);
  value_sum_.AddVec(1.0, temp);
  temp.AddRowSumMat(1.0, tanh_deriv, 0.0);
  deriv_sum_.AddVec(1.0, temp);

  if (count_ <= 0.0) {
    // this would be rather pathological if it happened.
    return;
  }

  // The rest of this function contains code modified from
  // TanhComponent::RepairGradients().

  // thresholds_vec is actually a 1-row matrix.  (the ApplyHeaviside
  // function isn't defined for vectors).
  CuMatrix<BaseFloat> thresholds(1, cell_dim_);
  CuSubVector<BaseFloat> thresholds_vec(thresholds, 0);
  thresholds_vec.AddVec(-1.0, deriv_sum_);
  thresholds_vec.Add(self_repair_threshold_ * count_);
  thresholds.ApplyHeaviside();
  self_repair_total_ += thresholds_vec.Sum();

  // there is a comment explaining what we are doing with
  // 'thresholds_vec', at this point in TanhComponent::RepairGradients().
  // We won't repeat it here.

  h_t_deriv->AddMatDiagVec(-self_repair_scale_ / repair_and_stats_probability,
                           h_t, kNoTrans, thresholds_vec);
}

void GruNonlinearityComponent::UpdateParameters(
    const CuMatrixBase<BaseFloat> &sdotr,
    const CuMatrixBase<BaseFloat> &h_t_deriv) {
  if (is_gradient_) {
    // 'simple' update, no natural gradient.  Compare
    // with AffineComponent::UpdateSimple().
    w_h_.AddMatMat(learning_rate_, h_t_deriv, kTrans,
                   sdotr, kNoTrans, 1.0);
  } else {
    // the natural-gradient update.
    CuMatrix<BaseFloat> in_value_temp(sdotr),
        out_deriv_temp(h_t_deriv);

    // These "scale" values get will get multiplied into the learning rate.
    BaseFloat in_scale, out_scale;

    preconditioner_in_.PreconditionDirections(&in_value_temp, &in_scale);
    preconditioner_out_.PreconditionDirections(&out_deriv_temp, &out_scale);

    BaseFloat local_lrate = learning_rate_ * in_scale * out_scale;
    w_h_.AddMatMat(local_lrate, out_deriv_temp, kTrans,
                   in_value_temp, kNoTrans, 1.0);
  }
}



void GruNonlinearityComponent::Read(std::istream &is, bool binary) {
  ReadUpdatableCommon(is, binary);
  ExpectToken(is, binary, "<CellDim>");
  ReadBasicType(is, binary, &cell_dim_);
  ExpectToken(is, binary, "<RecurrentDim>");
  ReadBasicType(is, binary, &recurrent_dim_);
  ExpectToken(is, binary, "<w_h>");
  w_h_.Read(is, binary);
  ExpectToken(is, binary, "<ValueAvg>");
  value_sum_.Read(is, binary);
  ExpectToken(is, binary, "<DerivAvg>");
  deriv_sum_.Read(is, binary);
  ExpectToken(is, binary, "<SelfRepairTotal>");
  ReadBasicType(is, binary, &self_repair_total_);
  ExpectToken(is, binary, "<Count>");
  ReadBasicType(is, binary, &count_);
  value_sum_.Scale(count_);  // we read in the averages, not the sums.
  deriv_sum_.Scale(count_);
  ExpectToken(is, binary, "<SelfRepairThreshold>");
  ReadBasicType(is, binary, &self_repair_threshold_);
  ExpectToken(is, binary, "<SelfRepairScale>");
  ReadBasicType(is, binary, &self_repair_scale_);
  BaseFloat alpha;
  int32 rank_in, rank_out, update_period;
  ExpectToken(is, binary, "<Alpha>");
  ReadBasicType(is, binary, &alpha);
  ExpectToken(is, binary, "<RankInOut>");
  ReadBasicType(is, binary, &rank_in);
  ReadBasicType(is, binary, &rank_out);
  ExpectToken(is, binary, "<UpdatePeriod>");
  ReadBasicType(is, binary, &update_period);
  preconditioner_in_.SetRank(rank_in);
  preconditioner_out_.SetRank(rank_out);
  preconditioner_in_.SetAlpha(alpha);
  preconditioner_out_.SetAlpha(alpha);
  preconditioner_in_.SetUpdatePeriod(update_period);
  preconditioner_out_.SetUpdatePeriod(update_period);
  ExpectToken(is, binary, "</GruNonlinearityComponent>");
}

void GruNonlinearityComponent::Write(std::ostream &os, bool binary) const {
  WriteUpdatableCommon(os, binary);
  WriteToken(os, binary, "<CellDim>");
  WriteBasicType(os, binary, cell_dim_);
  WriteToken(os, binary, "<RecurrentDim>");
  WriteBasicType(os, binary, recurrent_dim_);
  WriteToken(os, binary, "<w_h>");
  w_h_.Write(os, binary);
  {
    // Write the value and derivative stats in a count-normalized way, for
    // greater readability in text form.
    WriteToken(os, binary, "<ValueAvg>");
    Vector<BaseFloat> temp(value_sum_);
    if (count_ != 0.0) temp.Scale(1.0 / count_);
    temp.Write(os, binary);
    WriteToken(os, binary, "<DerivAvg>");
    temp.CopyFromVec(deriv_sum_);
    if (count_ != 0.0) temp.Scale(1.0 / count_);
    temp.Write(os, binary);
  }
  WriteToken(os, binary, "<SelfRepairTotal>");
  WriteBasicType(os, binary, self_repair_total_);
  WriteToken(os, binary, "<Count>");
  WriteBasicType(os, binary, count_);
  WriteToken(os, binary, "<SelfRepairThreshold>");
  WriteBasicType(os, binary, self_repair_threshold_);
  WriteToken(os, binary, "<SelfRepairScale>");
  WriteBasicType(os, binary, self_repair_scale_);

  BaseFloat alpha = preconditioner_in_.GetAlpha();
  int32 rank_in = preconditioner_in_.GetRank(),
      rank_out = preconditioner_out_.GetRank(),
      update_period = preconditioner_in_.GetUpdatePeriod();
  WriteToken(os, binary, "<Alpha>");
  WriteBasicType(os, binary, alpha);
  WriteToken(os, binary, "<RankInOut>");
  WriteBasicType(os, binary, rank_in);
  WriteBasicType(os, binary, rank_out);
  WriteToken(os, binary, "<UpdatePeriod>");
  WriteBasicType(os, binary, update_period);
  WriteToken(os, binary, "</GruNonlinearityComponent>");
}

void GruNonlinearityComponent::Scale(BaseFloat scale) {
  if (scale == 0.0) {
    w_h_.SetZero();
    value_sum_.SetZero();
    deriv_sum_.SetZero();
    self_repair_total_ = 0.0;
    count_ = 0.0;
  } else {
    w_h_.Scale(scale);
    value_sum_.Scale(scale);
    deriv_sum_.Scale(scale);
    self_repair_total_ *= scale;
    count_ *= scale;
  }
}

void GruNonlinearityComponent::Add(BaseFloat alpha,
                                   const Component &other_in) {
  const GruNonlinearityComponent *other =
      dynamic_cast<const GruNonlinearityComponent*>(&other_in);
  KALDI_ASSERT(other != NULL);
  w_h_.AddMat(alpha, other->w_h_);
  value_sum_.AddVec(alpha, other->value_sum_);
  deriv_sum_.AddVec(alpha, other->deriv_sum_);
  self_repair_total_ += alpha * other->self_repair_total_;
  count_ += alpha * other->count_;
}

void GruNonlinearityComponent::ZeroStats() {
  value_sum_.SetZero();
  deriv_sum_.SetZero();
  self_repair_total_ = 0.0;
  count_ = 0.0;
}

void GruNonlinearityComponent::Check() const {
  KALDI_ASSERT(cell_dim_ > 0 && recurrent_dim_ > 0 &&
               recurrent_dim_ <= cell_dim_ &&
               self_repair_threshold_ >= 0.0 &&
               self_repair_scale_ >= 0.0 );
  KALDI_ASSERT(w_h_.NumRows() == cell_dim_ &&
               w_h_.NumCols() == recurrent_dim_);
  KALDI_ASSERT(value_sum_.Dim() == cell_dim_ &&
               deriv_sum_.Dim() == cell_dim_);
}

void GruNonlinearityComponent::PerturbParams(BaseFloat stddev) {
  CuMatrix<BaseFloat> temp_params(w_h_.NumRows(), w_h_.NumCols());
  temp_params.SetRandn();
  w_h_.AddMat(stddev, temp_params);
}

BaseFloat GruNonlinearityComponent::DotProduct(
    const UpdatableComponent &other_in) const {
  const GruNonlinearityComponent *other =
      dynamic_cast<const GruNonlinearityComponent*>(&other_in);
  KALDI_ASSERT(other != NULL);
  return TraceMatMat(w_h_, other->w_h_, kTrans);
}

int32 GruNonlinearityComponent::NumParameters() const {
  return w_h_.NumRows() * w_h_.NumCols();
}

void GruNonlinearityComponent::Vectorize(VectorBase<BaseFloat> *params) const {
  KALDI_ASSERT(params->Dim() == NumParameters());
  params->CopyRowsFromMat(w_h_);
}


void GruNonlinearityComponent::UnVectorize(
    const VectorBase<BaseFloat> &params)  {
  KALDI_ASSERT(params.Dim() == NumParameters());
  w_h_.CopyRowsFromVec(params);
}

void GruNonlinearityComponent::FreezeNaturalGradient(bool freeze) {
  preconditioner_in_.Freeze(freeze);
  preconditioner_out_.Freeze(freeze);
}

GruNonlinearityComponent::GruNonlinearityComponent(
    const GruNonlinearityComponent &other):
    UpdatableComponent(other),
    cell_dim_(other.cell_dim_),
    recurrent_dim_(other.recurrent_dim_),
    w_h_(other.w_h_),
    value_sum_(other.value_sum_),
    deriv_sum_(other.deriv_sum_),
    self_repair_total_(other.self_repair_total_),
    count_(other.count_),
    self_repair_threshold_(other.self_repair_threshold_),
    self_repair_scale_(other.self_repair_scale_),
    preconditioner_in_(other.preconditioner_in_),
    preconditioner_out_(other.preconditioner_out_) {
  Check();
}


int32 OutputGruNonlinearityComponent::InputDim() const {
  return 3 * cell_dim_;
}

int32 OutputGruNonlinearityComponent::OutputDim() const {
  return 2 * cell_dim_;
}


std::string OutputGruNonlinearityComponent::Info() const {
  std::ostringstream stream;
  stream << UpdatableComponent::Info()
         << ", cell-dim=" << cell_dim_;
  PrintParameterStats(stream, "w_h", w_h_);
  stream << ", self-repair-threshold=" << self_repair_threshold_
         << ", self-repair-scale=" << self_repair_scale_;
  if (count_ > 0) {  // c.f. NonlinearComponent::Info().
    stream << ", count=" << std::setprecision(3) << count_
           << std::setprecision(6);
    stream << ", self-repaired-proportion="
           << (self_repair_total_ / (count_ * cell_dim_));
    Vector<double> value_avg_dbl(value_sum_);
    Vector<BaseFloat> value_avg(value_avg_dbl);
    value_avg.Scale(1.0 / count_);
    stream << ", value-avg=" << SummarizeVector(value_avg);
    Vector<double> deriv_avg_dbl(deriv_sum_);
    Vector<BaseFloat> deriv_avg(deriv_avg_dbl);
    deriv_avg.Scale(1.0 / count_);
    stream << ", deriv-avg=" << SummarizeVector(deriv_avg);
  }
  // natural-gradient parameters.
  stream << ", alpha=" << preconditioner_.GetAlpha()
         << ", rank=" << preconditioner_.GetRank()
         << ", update-period="
         << preconditioner_.GetUpdatePeriod();
  return stream.str();
}

void OutputGruNonlinearityComponent::InitFromConfig(ConfigLine *cfl) {
  cell_dim_ = -1;
  self_repair_threshold_ = 0.2;
  self_repair_scale_ = 1.0e-05;

  InitLearningRatesFromConfig(cfl);
  if (!cfl->GetValue("cell-dim", &cell_dim_) || cell_dim_ <= 0)
    KALDI_ERR << "cell-dim > 0 is required for GruNonlinearityComponent.";

  BaseFloat param_mean = 0.0, param_stddev = 1.0, 
      alpha = 4.0;
  int32 rank=8,
      update_period = 10;

  cfl->GetValue("self-repair-threshold", &self_repair_threshold_);
  cfl->GetValue("self-repair-scale", &self_repair_scale_);
  cfl->GetValue("param-mean", &param_mean);
  cfl->GetValue("param-stddev", &param_stddev);
  cfl->GetValue("alpha", &alpha);
  cfl->GetValue("rank", &rank);
  cfl->GetValue("update-period", &update_period);


  w_h_.Resize(cell_dim_);
  w_h_.SetRandn();
  w_h_.Scale(param_stddev);
  w_h_.Add(param_mean);

  preconditioner_.SetAlpha(alpha);
  preconditioner_.SetRank(rank);
  preconditioner_.SetUpdatePeriod(update_period);

  count_ = 0.0;
  self_repair_total_ = 0.0;
  value_sum_.Resize(cell_dim_);
  deriv_sum_.Resize(cell_dim_);

  Check();
}

void* OutputGruNonlinearityComponent::Propagate(
    const ComponentPrecomputedIndexes *indexes,
    const CuMatrixBase<BaseFloat> &in,
    CuMatrixBase<BaseFloat> *out) const {
  KALDI_ASSERT(in.NumRows() == out->NumRows() &&
               in.NumCols() == InputDim() &&
               out->NumCols() == OutputDim());
  // This component implements the function
  // (z_t, hpart_t, c_{t-1}) -> (h_t, c_t)
  // of dimensions
  // (cell_dim, cell_dim, cell_dim) -> (cell_dim, cell_dim),
  // where:
  // h_t = \tanh( hpart_t + W^h \dot c_{t-1} )
  // c_t = (1 - z_t) \dot h_t + z_t \dot c_{t-1}.
  int32 num_rows = in.NumRows(),
      c = cell_dim_;
  CuSubMatrix<BaseFloat> z_t(in, 0, num_rows, 0, c),
      hpart_t(in, 0, num_rows, c, c),
      c_t1(in, 0, num_rows, c + c, c);

  CuSubMatrix<BaseFloat> h_t(*out, 0, num_rows, 0, c),
      c_t(*out, 0, num_rows, c, c);

  h_t.CopyFromMat(c_t1);
  // now h_t = c_{t-1}
  h_t.MulColsVec(w_h_);
  // now h_t = W^h \dot c_{t-1}
  h_t.AddMat(1.0, hpart_t, kNoTrans);
  // now h_t = hpart_t + W^h \dot c_{t-1}.(note: hpart_t actually means U^h x_t).
  h_t.Tanh(h_t);
  // now, h_t = tanh(hpart_t + W^h \dot c_{t-1}).

  c_t.CopyFromMat(h_t);
  // now c_t = h_t
  c_t.AddMatMatElements(-1.0, z_t, h_t, 1.0);
  // now c_t = (1 - z_t) \dot h_t.
  c_t.AddMatMatElements(1.0, z_t, c_t1, 1.0);
  // now c_t = (1 - z_t) \dot h_t  +  z_t \dot c_{t-1}.
  return NULL;
}

void OutputGruNonlinearityComponent::Backprop(
    const std::string &debug_info,
    const ComponentPrecomputedIndexes *, // indexes
    const CuMatrixBase<BaseFloat> &in_value,
    const CuMatrixBase<BaseFloat> &out_value,
    const CuMatrixBase<BaseFloat> &out_deriv,
    void *memo,
    Component *to_update_in,
    CuMatrixBase<BaseFloat> *in_deriv) const {
  KALDI_ASSERT(SameDim(out_value, out_deriv) &&
               in_value.NumRows() == out_value.NumRows() &&
               in_value.NumCols() == InputDim() &&
               out_value.NumCols() == OutputDim() &&
               (in_deriv == NULL || SameDim(in_value, *in_deriv)) &&
               memo == NULL);
  OutputGruNonlinearityComponent *to_update =
      dynamic_cast<OutputGruNonlinearityComponent*>(to_update_in);
  KALDI_ASSERT(in_deriv != NULL || to_update != NULL);
  int32 num_rows = in_value.NumRows(),
      c = cell_dim_;

  // To understand what's going on here, compare this code with the
  // corresponding 'forward' code in Propagate().


  CuSubMatrix<BaseFloat> z_t(in_value, 0, num_rows, 0, c),
      hpart_t(in_value, 0, num_rows, c, c),
      c_t1(in_value, 0, num_rows, c + c, c);

  // The purpose of this 'in_deriv_ptr' is so that we can create submatrices
  // like z_t_deriv without the code crashing.  If in_deriv is NULL these point
  // to 'in_value', and we'll be careful never to actually write to these
  // sub-matrices, which aside from being conceptually wrong would violate the
  // const semantics of this function.
  const CuMatrixBase<BaseFloat> *in_deriv_ptr =
      (in_deriv == NULL ? &in_value : in_deriv);
  CuSubMatrix<BaseFloat> z_t_deriv(*in_deriv_ptr, 0, num_rows, 0, c),
      hpart_t_deriv(*in_deriv_ptr, 0, num_rows, c, c),
      c_t1_deriv(*in_deriv_ptr, 0, num_rows, c + c, c);

  // Note: the output h_t is never actually used in the GRU computation (we only
  // output it because we want the value to be cached to save computation in the
  // backprop), so we expect that the 'h_t_deriv', if we extracted it in the
  // obvious way, would be all zeros.
  // We create a different, local h_t_deriv
  // variable that backpropagates the derivative from c_t_deriv.
  CuSubMatrix<BaseFloat> h_t(out_value, 0, num_rows, 0, c),
      c_t(out_value, 0, num_rows, c, c),
      c_t_deriv(out_deriv, 0, num_rows, c, c);
  CuMatrix<BaseFloat> h_t_deriv(num_rows, c, kUndefined);

  {  // we initialize h_t_deriv with the derivative from 'out_deriv'.
    // In real life in a GRU, this would always be zero; but in testing
    // code it may be nonzero and we include this term so that
    // the tests don't fail.  Note: if you were to remove these
    // lines, you'd have to change 'h_t_deriv.AddMat(1.0, c_t_deriv);' below
    // to a CopyFromMat() call.
    CuSubMatrix<BaseFloat> h_t_deriv_in(out_deriv, 0, num_rows, 0, c);
    h_t_deriv.CopyFromMat(h_t_deriv_in);
  }


  { // This block does the
    // backprop corresponding to the
    // forward-pass expression: c_t = (1 - z_t) \dot h_t + z_t \dot c_{t-1}.

    // First do: h_t_deriv = c_t_deriv \dot (1 - z_t).
    h_t_deriv.AddMat(1.0, c_t_deriv);
    h_t_deriv.AddMatMatElements(-1.0, c_t_deriv, z_t, 1.0);

    if (in_deriv) {
      // these should be self-explanatory if you study
      // the expression "c_t = (1 - z_t) \dot h_t + z_t \dot c_{t-1}".
      z_t_deriv.AddMatMatElements(-1.0, c_t_deriv, h_t, 1.0);
      z_t_deriv.AddMatMatElements(1.0, c_t_deriv, c_t1, 1.0);
      c_t1_deriv.AddMatMatElements(1.0, c_t_deriv, z_t, 1.0);
    }
  }

  h_t_deriv.DiffTanh(h_t, h_t_deriv);
  if (to_update)
    to_update->TanhStatsAndSelfRepair(h_t, &h_t_deriv);
  
  if (to_update)
    to_update->UpdateParameters(c_t1, h_t_deriv);
  // At this point, 'h_t_deriv' contains the derivative w.r.t.
  // the argument of the tanh function, i.e. w.r.t. the expression:
  //    hpart_t + W^h \dot c_{t-1}.
  // The next block propagates this to the derivative for h_part_t and c_t1
  // The derivative of z_t has already been finished.
  if (in_deriv) {
    hpart_t_deriv.AddMat(1.0, h_t_deriv);

    // Currently, c_t1_deriv contains the derivative from
    // c_t = (1 - z_t) \dot h_t + z_t \dot c_{t-1}
    // Now compute the h_t = \tanh(hpart_t + W^h \dot c_{t-1}) part
    h_t_deriv.MulColsVec(w_h_);
    // Combine the two parts
    c_t1_deriv.AddMat(1.0, h_t_deriv);
  }
}


void OutputGruNonlinearityComponent::TanhStatsAndSelfRepair(
    const CuMatrixBase<BaseFloat> &h_t,
    CuMatrixBase<BaseFloat> *h_t_deriv) {
  KALDI_ASSERT(SameDim(h_t, *h_t_deriv));

  // we use this probability (hardcoded for now) to limit the stats accumulation
  // and self-repair code to running on about half of the minibatches.
  BaseFloat repair_and_stats_probability = 0.5;
  if (RandUniform() > repair_and_stats_probability)
    return;

  // OK, accumulate stats.
  // For the next few lines, compare with TanhComponent::StoreStats(), which is where
  // we got this code.
  // tanh_deriv is the function derivative of the tanh function,
  // tanh'(x) = tanh(x) * (1.0 - tanh(x)).  h_t corresponds to tanh(x).
  CuMatrix<BaseFloat> tanh_deriv(h_t);
  tanh_deriv.ApplyPow(2.0);
  tanh_deriv.Scale(-1.0);
  tanh_deriv.Add(1.0);

  count_ += h_t.NumRows();
  CuVector<BaseFloat> temp(cell_dim_);
  temp.AddRowSumMat(1.0, h_t, 0.0);
  value_sum_.AddVec(1.0, temp);
  temp.AddRowSumMat(1.0, tanh_deriv, 0.0);
  deriv_sum_.AddVec(1.0, temp);

  if (count_ <= 0.0) {
    // this would be rather pathological if it happened.
    return;
  }

  // The rest of this function contains code modified from
  // TanhComponent::RepairGradients().

  // thresholds_vec is actually a 1-row matrix.  (the ApplyHeaviside
  // function isn't defined for vectors).
  CuMatrix<BaseFloat> thresholds(1, cell_dim_);
  CuSubVector<BaseFloat> thresholds_vec(thresholds, 0);
  thresholds_vec.AddVec(-1.0, deriv_sum_);
  thresholds_vec.Add(self_repair_threshold_ * count_);
  thresholds.ApplyHeaviside();
  self_repair_total_ += thresholds_vec.Sum();

  // there is a comment explaining what we are doing with
  // 'thresholds_vec', at this point in TanhComponent::RepairGradients().
  // We won't repeat it here.

  h_t_deriv->AddMatDiagVec(-self_repair_scale_ / repair_and_stats_probability,
                           h_t, kNoTrans, thresholds_vec);
}

void OutputGruNonlinearityComponent::UpdateParameters(
    const CuMatrixBase<BaseFloat> &c_t1_value,
    const CuMatrixBase<BaseFloat> &h_t_deriv) {
  if (is_gradient_) {
    // 'simple' update, no natural gradient.  Compare
    // with PerElementScaleComponent::UpdateSimple().
    w_h_.AddDiagMatMat(learning_rate_, h_t_deriv, kTrans,
                       c_t1_value, kNoTrans, 1.0);
  } else {
    // the natural-gradient update.
    CuMatrix<BaseFloat> derivs_per_frame(c_t1_value);
    derivs_per_frame.MulElements(h_t_deriv);

    // This "scale" value gets will get multiplied into the learning rate.
    BaseFloat scale;

    preconditioner_.PreconditionDirections(&derivs_per_frame, &scale);

    CuVector<BaseFloat> delta_w_h(w_h_.Dim());
    delta_w_h.AddRowSumMat(scale * learning_rate_, derivs_per_frame);
    w_h_.AddVec(1.0, delta_w_h);
  }
}



void OutputGruNonlinearityComponent::Read(std::istream &is, bool binary) {
  ReadUpdatableCommon(is, binary);
  ExpectToken(is, binary, "<CellDim>");
  ReadBasicType(is, binary, &cell_dim_);
  ExpectToken(is, binary, "<w_h>");
  w_h_.Read(is, binary);
  ExpectToken(is, binary, "<ValueAvg>");
  value_sum_.Read(is, binary);
  ExpectToken(is, binary, "<DerivAvg>");
  deriv_sum_.Read(is, binary);
  ExpectToken(is, binary, "<SelfRepairTotal>");
  ReadBasicType(is, binary, &self_repair_total_);
  ExpectToken(is, binary, "<Count>");
  ReadBasicType(is, binary, &count_);
  value_sum_.Scale(count_);  // we read in the averages, not the sums.
  deriv_sum_.Scale(count_);
  ExpectToken(is, binary, "<SelfRepairThreshold>");
  ReadBasicType(is, binary, &self_repair_threshold_);
  ExpectToken(is, binary, "<SelfRepairScale>");
  ReadBasicType(is, binary, &self_repair_scale_);
  BaseFloat alpha;
  int32 rank, update_period;
  ExpectToken(is, binary, "<Alpha>");
  ReadBasicType(is, binary, &alpha);
  ExpectToken(is, binary, "<Rank>");
  ReadBasicType(is, binary, &rank);
  ExpectToken(is, binary, "<UpdatePeriod>");
  ReadBasicType(is, binary, &update_period);
  preconditioner_.SetRank(rank);
  preconditioner_.SetAlpha(alpha);
  preconditioner_.SetUpdatePeriod(update_period);
  ExpectToken(is, binary, "</OutputGruNonlinearityComponent>");
}

void OutputGruNonlinearityComponent::Write(std::ostream &os, bool binary) const {
  WriteUpdatableCommon(os, binary);
  WriteToken(os, binary, "<CellDim>");
  WriteBasicType(os, binary, cell_dim_);
  WriteToken(os, binary, "<w_h>");
  w_h_.Write(os, binary);
  {
    // Write the value and derivative stats in a count-normalized way, for
    // greater readability in text form.
    WriteToken(os, binary, "<ValueAvg>");
    Vector<BaseFloat> temp(value_sum_);
    if (count_ != 0.0) temp.Scale(1.0 / count_);
    temp.Write(os, binary);
    WriteToken(os, binary, "<DerivAvg>");
    temp.CopyFromVec(deriv_sum_);
    if (count_ != 0.0) temp.Scale(1.0 / count_);
    temp.Write(os, binary);
  }
  WriteToken(os, binary, "<SelfRepairTotal>");
  WriteBasicType(os, binary, self_repair_total_);
  WriteToken(os, binary, "<Count>");
  WriteBasicType(os, binary, count_);
  WriteToken(os, binary, "<SelfRepairThreshold>");
  WriteBasicType(os, binary, self_repair_threshold_);
  WriteToken(os, binary, "<SelfRepairScale>");
  WriteBasicType(os, binary, self_repair_scale_);

  BaseFloat alpha = preconditioner_.GetAlpha();
  int32 rank = preconditioner_.GetRank(),
      update_period = preconditioner_.GetUpdatePeriod();
  WriteToken(os, binary, "<Alpha>");
  WriteBasicType(os, binary, alpha);
  WriteToken(os, binary, "<Rank>");
  WriteBasicType(os, binary, rank);
  WriteToken(os, binary, "<UpdatePeriod>");
  WriteBasicType(os, binary, update_period);
  WriteToken(os, binary, "</OutputGruNonlinearityComponent>");
}

void OutputGruNonlinearityComponent::Scale(BaseFloat scale) {
  if (scale == 0.0) {
    w_h_.SetZero();
    value_sum_.SetZero();
    deriv_sum_.SetZero();
    self_repair_total_ = 0.0;
    count_ = 0.0;
  } else {
    w_h_.Scale(scale);
    value_sum_.Scale(scale);
    deriv_sum_.Scale(scale);
    self_repair_total_ *= scale;
    count_ *= scale;
  }
}

void OutputGruNonlinearityComponent::Add(BaseFloat alpha,
                                   const Component &other_in) {
  const OutputGruNonlinearityComponent *other =
      dynamic_cast<const OutputGruNonlinearityComponent*>(&other_in);
  KALDI_ASSERT(other != NULL);
  w_h_.AddVec(alpha, other->w_h_);
  value_sum_.AddVec(alpha, other->value_sum_);
  deriv_sum_.AddVec(alpha, other->deriv_sum_);
  self_repair_total_ += alpha * other->self_repair_total_;
  count_ += alpha * other->count_;
}

void OutputGruNonlinearityComponent::ZeroStats() {
  value_sum_.SetZero();
  deriv_sum_.SetZero();
  self_repair_total_ = 0.0;
  count_ = 0.0;
}

void OutputGruNonlinearityComponent::Check() const {
  KALDI_ASSERT(cell_dim_ > 0 &&
               self_repair_threshold_ >= 0.0 &&
               self_repair_scale_ >= 0.0 );
  KALDI_ASSERT(w_h_.Dim() == cell_dim_);
  KALDI_ASSERT(value_sum_.Dim() == cell_dim_ &&
               deriv_sum_.Dim() == cell_dim_);
}

void OutputGruNonlinearityComponent::PerturbParams(BaseFloat stddev) {
  CuVector<BaseFloat> temp_params(w_h_.Dim());
  temp_params.SetRandn();
  w_h_.AddVec(stddev, temp_params);
}

BaseFloat OutputGruNonlinearityComponent::DotProduct(
    const UpdatableComponent &other_in) const {
  const OutputGruNonlinearityComponent *other =
      dynamic_cast<const OutputGruNonlinearityComponent*>(&other_in);
  KALDI_ASSERT(other != NULL);
  return VecVec(w_h_, other->w_h_);
}

int32 OutputGruNonlinearityComponent::NumParameters() const {
  return w_h_.Dim();
}

void OutputGruNonlinearityComponent::Vectorize(VectorBase<BaseFloat> *params) const {
  KALDI_ASSERT(params->Dim() == NumParameters());
  params->CopyFromVec(w_h_);
}


void OutputGruNonlinearityComponent::UnVectorize(
    const VectorBase<BaseFloat> &params)  {
  KALDI_ASSERT(params.Dim() == NumParameters());
  w_h_.CopyFromVec(params);
}

void OutputGruNonlinearityComponent::FreezeNaturalGradient(bool freeze) {
  preconditioner_.Freeze(freeze);
}

OutputGruNonlinearityComponent::OutputGruNonlinearityComponent(
    const OutputGruNonlinearityComponent &other):
    UpdatableComponent(other),
    cell_dim_(other.cell_dim_),
    w_h_(other.w_h_),
    value_sum_(other.value_sum_),
    deriv_sum_(other.deriv_sum_),
    self_repair_total_(other.self_repair_total_),
    count_(other.count_),
    self_repair_threshold_(other.self_repair_threshold_),
    self_repair_scale_(other.self_repair_scale_),
    preconditioner_(other.preconditioner_) {
  Check();
}

} // namespace nnet3
} // namespace kaldi