// nnet3/nnet-tdnn-component.h

// Copyright      2017  Johns Hopkins University (author: Daniel Povey)

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

// Note: the code defined here was declared in nnet-convolutional-component.h.

#include <iterator>
#include <sstream>
#include <iomanip>
#include "nnet3/nnet-convolutional-component.h"
#include "nnet3/nnet-computation-graph.h"
#include "nnet3/nnet-parse.h"

namespace kaldi {
namespace nnet3 {


TdnnComponent::TdnnComponent():
    orthonormal_constraint_(0.0),
    use_natural_gradient_(true) { }


TdnnComponent::TdnnComponent(
    const TdnnComponent &other):
    UpdatableComponent(other),  // initialize base-class
    time_offsets_(other.time_offsets_),
    linear_params_(other.linear_params_),
    bias_params_(other.bias_params_),
    orthonormal_constraint_(other.orthonormal_constraint_),
    use_natural_gradient_(other.use_natural_gradient_),
    preconditioner_in_(other.preconditioner_in_),
    preconditioner_out_(other.preconditioner_out_) {
  Check();
}


void TdnnComponent::Check() const {
  KALDI_ASSERT(linear_params_.NumRows() > 0 &&
               !time_offsets_.empty() &&
               std::set<int32>(time_offsets_.begin(),
                               time_offsets_.end()).size() ==
               time_offsets_.size() &&
               linear_params_.NumCols() % time_offsets_.size() == 0 &&
               (bias_params_.Dim() == 0 ||
                bias_params_.Dim() == linear_params_.NumRows()));
}

std::string TdnnComponent::Info() const {
  std::ostringstream stream;
  stream << UpdatableComponent::Info();
  if (orthonormal_constraint_ != 0.0)
    stream << ", orthonormal-constraint=" << orthonormal_constraint_;
  stream << ", time-offsets=";
  for (size_t i = 0; i < time_offsets_.size(); i++) {
    if (i != 0) stream << ',';
    stream << time_offsets_[i];
  }
  PrintParameterStats(stream, "linear-params", linear_params_,
                      false, // include_mean
                      true, // include_row_norms
                      true, // include_column_norms
                      GetVerboseLevel() >= 2); // include_singular_values
  if (bias_params_.Dim() == 0) {
    stream << ", has-bias=false";
  } else {
    PrintParameterStats(stream, "bias", bias_params_, true);
  }
  if (!use_natural_gradient_) {
    stream << ", use-natural-gradient=false";
  } else {
    stream << ", rank-in=" << preconditioner_in_.GetRank()
           << ", rank-out=" << preconditioner_out_.GetRank()
           << ", num-samples-history=" << preconditioner_in_.GetNumSamplesHistory()
           << ", update-period=" << preconditioner_in_.GetUpdatePeriod()
           << ", alpha-in=" << preconditioner_in_.GetAlpha()
           << ", alpha-out=" << preconditioner_out_.GetAlpha();
  }
  return stream.str();
}


void TdnnComponent::InitFromConfig(ConfigLine *cfl) {
  // 1. Config values inherited from UpdatableComponent.
  InitLearningRatesFromConfig(cfl);

  // 2. Structural config values
  std::string time_offsets;

  int32 input_dim = -1, output_dim = -1;

  bool ok = cfl->GetValue("time-offsets", &time_offsets) &&
      cfl->GetValue("input-dim", &input_dim) &&
      cfl->GetValue("output-dim", &output_dim);
  if (!ok || input_dim <= 0 || output_dim <= 0 ||
      !SplitStringToIntegers(time_offsets, ",", false, &time_offsets_) ||
      time_offsets_.empty()) {
    KALDI_ERR << "Bad initializer: there is a problem with "
        "time-offsets, input-dim or output-dim (not defined?): "
        << cfl->WholeLine();
  }

  if (std::set<int32>(time_offsets_.begin(),
                      time_offsets_.end()).size() != time_offsets_.size()) {
    KALDI_ERR << "Bad initializer: repeated time-offsets: "
              << cfl->WholeLine();
  }

  // 3. Parameter-initialization configs, "has-bias", and
  // orthonormal-constraint.
  orthonormal_constraint_ = 0.0;
  BaseFloat param_stddev = -1, bias_mean = 0.0, bias_stddev = 1.0;
  bool use_bias = true;
  cfl->GetValue("param-stddev", &param_stddev);
  cfl->GetValue("bias-stddev", &bias_stddev);
  cfl->GetValue("bias-mean", &bias_mean);
  cfl->GetValue("use-bias", &use_bias);
  cfl->GetValue("orthonormal-constraint", &orthonormal_constraint_);
  if (param_stddev < 0.0) {
    param_stddev = 1.0 / sqrt(input_dim * time_offsets_.size());
  }
  // initialize the parameters.
  linear_params_.Resize(output_dim,
                        input_dim * time_offsets_.size());
  linear_params_.SetRandn();
  linear_params_.Scale(param_stddev);

  if (use_bias) {
    bias_params_.Resize(output_dim);
    bias_params_.SetRandn();
    bias_params_.Scale(bias_stddev);
    bias_params_.Add(bias_mean);
  } else {
    bias_params_.Resize(0);
  }

  // 4. Natural-gradient related configs.
  use_natural_gradient_ = true;
  int32 rank_out = -1, rank_in = -1;
  BaseFloat alpha_out = 4.0, alpha_in = 4.0,
      num_samples_history = 2000.0;
  cfl->GetValue("use-natural-gradient", &use_natural_gradient_);
  cfl->GetValue("rank-in", &rank_in);
  cfl->GetValue("rank-out", &rank_out);
  cfl->GetValue("alpha-in", &alpha_in);
  cfl->GetValue("alpha-out", &alpha_out);
  cfl->GetValue("num-samples-history", &num_samples_history);

  int32 spliced_input_dim =
      input_dim * static_cast<int32>(time_offsets_.size());
  if (rank_in < 0)
    rank_in = std::min<int32>(20, (spliced_input_dim + 1) / 2);
  preconditioner_in_.SetRank(rank_in);
  if (rank_out < 0)
    rank_out = std::min<int32>(80, (output_dim + 1) / 2);
  preconditioner_out_.SetRank(rank_out);
  preconditioner_in_.SetNumSamplesHistory(num_samples_history);
  preconditioner_out_.SetNumSamplesHistory(num_samples_history);

  preconditioner_in_.SetAlpha(alpha_in);
  preconditioner_out_.SetAlpha(alpha_out);

  preconditioner_in_.SetUpdatePeriod(4);
  preconditioner_out_.SetUpdatePeriod(4);
}

void* TdnnComponent::Propagate(
    const ComponentPrecomputedIndexes *indexes_in,
    const CuMatrixBase<BaseFloat> &in,
    CuMatrixBase<BaseFloat> *out) const {
  const PrecomputedIndexes *indexes =
      dynamic_cast<const PrecomputedIndexes*>(indexes_in);
  KALDI_ASSERT(indexes != NULL);

  if (bias_params_.Dim() != 0)
    out->CopyRowsFromVec(bias_params_);
  // if bias_params_.Dim() == 0 we don't need to zero 'out' at
  // this point because in that case we set the flag kPropagateAdds,
  // so the calling code knows that the Propagate function *adds to*
  // the 'out' matrix, so it should (typicaly) be zeroed before calling
  // Propagate().

  KALDI_ASSERT(indexes->row_offsets.size() == time_offsets_.size());

  int32 num_offsets = time_offsets_.size(),
      input_dim = InputDim();
  for (int32 i = 0; i < num_offsets; i++) {
    CuSubMatrix<BaseFloat> in_part = GetInputPart(in, out->NumRows(),
                                                  indexes->row_stride,
                                                  indexes->row_offsets[i]);
    CuSubMatrix<BaseFloat> linear_params_part(linear_params_,
                                              0, linear_params_.NumRows(),
                                              i * input_dim, input_dim);
    out->AddMatMat(1.0, in_part, kNoTrans, linear_params_part, kTrans, 1.0);
  }
  return NULL;
}

void TdnnComponent::Backprop(
    const std::string &debug_info,
    const ComponentPrecomputedIndexes *indexes_in,
    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 {
  const PrecomputedIndexes *indexes =
      dynamic_cast<const PrecomputedIndexes*>(indexes_in);
  KALDI_ASSERT(indexes != NULL &&
               indexes->row_offsets.size() == time_offsets_.size());
  int32 num_offsets = time_offsets_.size(),
      input_dim = InputDim();

  if (in_deriv != NULL) {
    // Propagate the derivatives back to the input data.
    for (int32 i = 0; i < num_offsets; i++) {
      CuSubMatrix<BaseFloat> in_deriv_part =
          GetInputPart(*in_deriv, out_deriv.NumRows(),
                       indexes->row_stride, indexes->row_offsets[i]);
      CuSubMatrix<BaseFloat> linear_params_part(linear_params_,
                                                0, linear_params_.NumRows(),
                                                i * input_dim, input_dim);
      // note: this component has the property kBackpropAdds, which is why the
      // final 1.0 is there in the following call (otherwise we'd have to zero
      // *in_deriv first).
      in_deriv_part.AddMatMat(1.0, out_deriv, kNoTrans,
                              linear_params_part, kNoTrans, 1.0);
    }
  }

  if (to_update_in != NULL) {
    TdnnComponent *to_update =
        dynamic_cast<TdnnComponent*>(to_update_in);
    KALDI_ASSERT(to_update != NULL);

    if (to_update->learning_rate_ == 0.0)
      return;

    if (to_update->is_gradient_ || !to_update->use_natural_gradient_)
      to_update->UpdateSimple(*indexes, in_value, out_deriv);
    else
      to_update->UpdateNaturalGradient(*indexes, in_value, out_deriv);
  }
}

void TdnnComponent::UpdateSimple(
    const PrecomputedIndexes &indexes,
    const CuMatrixBase<BaseFloat> &in_value,
    const CuMatrixBase<BaseFloat> &out_deriv) {

  if (bias_params_.Dim() != 0)
    bias_params_.AddRowSumMat(learning_rate_, out_deriv);

  int32 input_dim = in_value.NumCols(),
      num_offsets = time_offsets_.size();
  for (int32 i = 0; i < num_offsets; i++) {
    CuSubMatrix<BaseFloat> in_value_part =
        GetInputPart(in_value, out_deriv.NumRows(),
                     indexes.row_stride,
                     indexes.row_offsets[i]);
    CuSubMatrix<BaseFloat> linear_params_part(linear_params_,
                                              0, linear_params_.NumRows(),
                                              i * input_dim, input_dim);
    linear_params_part.AddMatMat(learning_rate_, out_deriv, kTrans,
                                 in_value_part, kNoTrans, 1.0);
  }
}

void TdnnComponent::UpdateNaturalGradient(
    const PrecomputedIndexes &indexes,
    const CuMatrixBase<BaseFloat> &in_value,
    const CuMatrixBase<BaseFloat> &out_deriv) {

  int32 num_offsets = time_offsets_.size(),
      num_rows = out_deriv.NumRows(),
      input_dim = in_value.NumCols(),
      spliced_input_dim = num_offsets * input_dim,
      augmented_input_dim =
        spliced_input_dim + (bias_params_.Dim() != 0 ? 1 : 0);

  // in_value_temp is the fully spliced input with a column of ones appended to
  // it.
  CuMatrix<BaseFloat> in_value_temp(num_rows,
                                    augmented_input_dim);
  if (bias_params_.Dim() != 0) {
    // set the last column of in_value_temp to 1.0
    in_value_temp.Range(0, num_rows, spliced_input_dim, 1).Set(1.0);
  }

  for (int32 i = 0; i < num_offsets; i++) {
    CuSubMatrix<BaseFloat> in_value_temp_part(in_value_temp,
                                              0, num_rows,
                                              i * input_dim, input_dim),
        in_value_part = GetInputPart(in_value,
                                     num_rows,
                                     indexes.row_stride,
                                     indexes.row_offsets[i]);
    in_value_temp_part.CopyFromMat(in_value_part);
  }

  CuMatrix<BaseFloat> out_deriv_temp(out_deriv);

  // These "scale" values get will get multiplied into the learning rate (faster
  // than having the matrices scaled inside the preconditioning code).
  BaseFloat in_scale, out_scale;

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

  // "scale" is a scaling factor coming from the PreconditionDirections calls
  // (it's faster to have them output a scaling factor than to have them scale
  // their outputs).
  BaseFloat scale = in_scale * out_scale,
      local_lrate = scale * learning_rate_;

  if (bias_params_.Dim() != 0) {
    // this "precon_ones" is what happens to the vector of 1's representing
    // offsets, after multiplication by the preconditioner.
    CuVector<BaseFloat> precon_ones(num_rows);
    precon_ones.CopyColFromMat(in_value_temp, spliced_input_dim);
    bias_params_.AddMatVec(local_lrate, out_deriv_temp, kTrans,
                           precon_ones, 1.0);
  }

  CuSubMatrix<BaseFloat> in_value_precon_part(in_value_temp,
                                              0, num_rows,
                                              0, spliced_input_dim);

  linear_params_.AddMatMat(local_lrate, out_deriv_temp, kTrans,
                           in_value_precon_part, kNoTrans, 1.0);
}

void TdnnComponent::ReorderIndexes(
    std::vector<Index> *input_indexes,
    std::vector<Index> *output_indexes) const {
  using namespace time_height_convolution;

  // The following figures out a regular structure for the input and
  // output indexes, in case there were gaps (which is unlikely in typical
  // situations).
  ConvolutionComputationIo io;
  GetComputationIo(*input_indexes, *output_indexes, &io);
  ModifyComputationIo(&io);

  std::vector<Index> modified_input_indexes,
      modified_output_indexes;
  // The following call ensures that 'modified_input_indexes' and
  // 'modified_output_indexes' have the required ordering (where t has the
  // largest stride and each (n,x) pair is repeated for each 't' value), as well
  // as doing padding (setting t values to kNoTime where it had to insert
  // elements to ensure regular structure).
  GetIndexesForComputation(io, *input_indexes, *output_indexes,
                           &modified_input_indexes,
                           &modified_output_indexes);

  // It will be quite rare that this function actually changes
  // 'input_indexes' or 'output_indexes', because in most cases,
  // the indexes will already have the required structure and
  // ordering.
  input_indexes->swap(modified_input_indexes);
  output_indexes->swap(modified_output_indexes);
}

void TdnnComponent::Write(std::ostream &os, bool binary) const {
  WriteUpdatableCommon(os, binary);  // Write opening tag and learning rate.
  WriteToken(os, binary, "<TimeOffsets>");
  WriteIntegerVector(os, binary, time_offsets_);
  WriteToken(os, binary, "<LinearParams>");
  linear_params_.Write(os, binary);
  WriteToken(os, binary, "<BiasParams>");
  bias_params_.Write(os, binary);
  WriteToken(os, binary, "<OrthonormalConstraint>");
  WriteBasicType(os, binary, orthonormal_constraint_);
  WriteToken(os, binary, "<UseNaturalGradient>");
  WriteBasicType(os, binary, use_natural_gradient_);
  int32 rank_in = preconditioner_in_.GetRank(),
      rank_out = preconditioner_out_.GetRank();
  BaseFloat alpha_in = preconditioner_in_.GetAlpha(),
      alpha_out = preconditioner_out_.GetAlpha(),
      num_samples_history = preconditioner_in_.GetNumSamplesHistory();
  WriteToken(os, binary, "<NumSamplesHistory>");
  WriteBasicType(os, binary, num_samples_history);
  WriteToken(os, binary, "<AlphaInOut>");
  WriteBasicType(os, binary, alpha_in);
  WriteBasicType(os, binary, alpha_out);
  WriteToken(os, binary, "<RankInOut>");
  WriteBasicType(os, binary, rank_in);
  WriteBasicType(os, binary, rank_out);
  WriteToken(os, binary, "</TdnnComponent>");
}

void TdnnComponent::Read(std::istream &is, bool binary) {
  std::string token = ReadUpdatableCommon(is, binary);
  ExpectToken(is, binary, "<TimeOffsets>");
  ReadIntegerVector(is, binary, &time_offsets_);
  ExpectToken(is, binary, "<LinearParams>");
  linear_params_.Read(is, binary);
  ExpectToken(is, binary, "<BiasParams>");
  bias_params_.Read(is, binary);
  ExpectToken(is, binary, "<OrthonormalConstraint>");
  ReadBasicType(is, binary, &orthonormal_constraint_);
  ExpectToken(is, binary, "<UseNaturalGradient>");
  ReadBasicType(is, binary, &use_natural_gradient_);
  int32 rank_in,  rank_out;
  BaseFloat alpha_in, alpha_out,
      num_samples_history;
  ExpectToken(is, binary, "<NumSamplesHistory>");
  ReadBasicType(is, binary, &num_samples_history);
  { // This can be simplified after a while.  It's to read a format of the model
    // that was never checked into master, but with which I (Dan) did many of
    // the experiments while tuning the resnet TDNN-F.
    std::string token;
    ReadToken(is, binary, &token);
    if (token == "<AlphaInOut>") {
      ReadBasicType(is, binary, &alpha_in);
      ReadBasicType(is, binary, &alpha_out);
    } else {
      KALDI_ASSERT(token == "<Alpha>");
      ReadBasicType(is, binary, &alpha_in);
      alpha_out = alpha_in;
    }
  }
  preconditioner_in_.SetAlpha(alpha_in);
  preconditioner_out_.SetAlpha(alpha_out);
  ExpectToken(is, binary, "<RankInOut>");
  ReadBasicType(is, binary, &rank_in);
  ReadBasicType(is, binary, &rank_out);
  preconditioner_in_.SetRank(rank_in);
  preconditioner_out_.SetRank(rank_out);
  preconditioner_in_.SetNumSamplesHistory(num_samples_history);
  preconditioner_out_.SetNumSamplesHistory(num_samples_history);
  // the update periods are not configurable.
  preconditioner_in_.SetUpdatePeriod(4);
  preconditioner_out_.SetUpdatePeriod(4);
  ExpectToken(is, binary, "</TdnnComponent>");
  Check();
}

void TdnnComponent::GetInputIndexes(
    const MiscComputationInfo &misc_info,
    const Index &output_index,
    std::vector<Index> *desired_indexes) const {
  KALDI_ASSERT(output_index.t != kNoTime);
  size_t size = time_offsets_.size();
  desired_indexes->resize(size);
  for (size_t i = 0; i < size; i++) {
    (*desired_indexes)[i].n = output_index.n;
    (*desired_indexes)[i].t = output_index.t + time_offsets_[i];
    (*desired_indexes)[i].x = output_index.x;
  }
}


bool TdnnComponent::IsComputable(
    const MiscComputationInfo &misc_info,
    const Index &output_index,
    const IndexSet &input_index_set,
    std::vector<Index> *used_inputs) const {
  KALDI_ASSERT(output_index.t != kNoTime);
  size_t size = time_offsets_.size();
  Index index(output_index);

  if (used_inputs != NULL) {
    used_inputs->clear();
    used_inputs->reserve(size);
  }
  for (size_t i = 0; i < size; i++) {
    index.t = output_index.t + time_offsets_[i];
    if (input_index_set(index)) {
      if (used_inputs != NULL) {
        // This input index is available.
        used_inputs->push_back(index);
      }
    } else {
      return false;
    }
  }
  return true;
}

// static
CuSubMatrix<BaseFloat> TdnnComponent::GetInputPart(
      const CuMatrixBase<BaseFloat> &input_matrix,
      int32 num_output_rows,
      int32 row_stride,
      int32 row_offset) {
  KALDI_ASSERT(row_offset >= 0 && row_stride >= 1 &&
               input_matrix.NumRows() >=
               row_offset + (row_stride * num_output_rows) - (row_stride - 1));
  // constructor takes args: (data, num_rows, num_cols, stride).
  return CuSubMatrix<BaseFloat>(
      input_matrix.Data() + input_matrix.Stride() * row_offset,
      num_output_rows,
      input_matrix.NumCols(),
      input_matrix.Stride() * row_stride);
}

void TdnnComponent::ModifyComputationIo(
    time_height_convolution::ConvolutionComputationIo *io) {
  if (io->t_step_out == 0) {
    // the 't_step' values may be zero if there was only one (input or output)
    // index so the time-stride could not be determined.  This code fixes them
    // up in that case.  (If there was only one value, the stride is a
    // don't-care actually).
    if (io->t_step_in == 0)
      io->t_step_in = 1;
    io->t_step_out = io->t_step_in;
  }
  // At this point the t_step_{in,out} values will be nonzero.
  KALDI_ASSERT(io->t_step_out % io->t_step_in == 0);
  // The following affects the ordering of the input indexes; it allows us to
  // reshape the input matrix in the way that we need to, in cases where there
  // is subsampling.  See the explanation where the variable was declared in
  // class ConvolutionComputationIo.
  io->reorder_t_in = io->t_step_out / io->t_step_in;

  // make sure that num_t_in is a multiple of io->reorder_t_in by rounding up.
  int32 n = io->reorder_t_in;
  io->num_t_in = n * ((io->num_t_in + n - 1) / n);
}

ComponentPrecomputedIndexes* TdnnComponent::PrecomputeIndexes(
      const MiscComputationInfo &misc_info,
      const std::vector<Index> &input_indexes,
      const std::vector<Index> &output_indexes,
      bool need_backprop) const {
  using namespace time_height_convolution;
  // The following figures out a regular structure for the input and
  // output indexes, in case there were gaps (which is unlikely in typical
  // situations).
  ConvolutionComputationIo io;
  GetComputationIo(input_indexes, output_indexes, &io);
  ModifyComputationIo(&io);

  if (RandInt(0, 10) == 0) {
    // Spot check that the provided indexes have the required properties;
    // this is like calling this->ReorderIndexes() and checking that it
    // doesn't change anything.
    std::vector<Index> modified_input_indexes,
        modified_output_indexes;
    GetIndexesForComputation(io, input_indexes, output_indexes,
                             &modified_input_indexes,
                             &modified_output_indexes);
    KALDI_ASSERT(modified_input_indexes == input_indexes &&
                 modified_output_indexes == output_indexes);
  }


  PrecomputedIndexes *ans = new PrecomputedIndexes();
  ans->row_stride = io.reorder_t_in;
  int32 num_offsets = time_offsets_.size();
  ans->row_offsets.resize(num_offsets);
  for (int32 i = 0; i < num_offsets; i++) {
    // For each offset, work out which row of the input has the same t value as
    // the first t value in the output plus that offset.  That becomes the start
    // row of the corresponding sub-part of the input.
    int32 time_offset = time_offsets_[i],
        required_input_t = io.start_t_out + time_offset,
        input_t = (required_input_t - io.start_t_in) / io.t_step_in;

    KALDI_ASSERT(required_input_t == io.start_t_in + io.t_step_in * input_t);
    // input_t is a kind of normalized time offset in the input, relative to the
    // first 't' value in the input and divided by the t-step in the input, so
    // it's the numbering "as if" the input 't' values were numbered from 0,1,2.
    // To turn input_t into an input row we need to take account of 'reorder_t_in'.
    // If this is 1 then the input row is input_t times io.num_images.
    // Otherwise it's a little more complicated and to understand it you should
    // read the comment where 'reorder_t_in' is declared in convolution.h.
    // Briefly: the part that is an integer multiple of 'reorder_t_in' gets
    // multiplied by io.num_images; the remainder does not.

    int32 n = io.reorder_t_in,
        input_t_multiple = n * (input_t / n), input_t_remainder = input_t % n;
    // note: input_t == input_t_multiple + input_t_remainder .
    int32 input_row_offset = input_t_multiple * io.num_images +
        input_t_remainder;
    ans->row_offsets[i] = input_row_offset;
  }
  return ans;
}

void TdnnComponent::Scale(BaseFloat scale) {
  if (scale == 0.0) {
    linear_params_.SetZero();
    bias_params_.SetZero();
  } else {
    linear_params_.Scale(scale);
    bias_params_.Scale(scale);
  }
}

void TdnnComponent::Add(BaseFloat alpha,
                        const Component &other_in) {
  const TdnnComponent *other =
      dynamic_cast<const TdnnComponent*>(&other_in);
  KALDI_ASSERT(other != NULL);
  linear_params_.AddMat(alpha, other->linear_params_);
  if (bias_params_.Dim() != 0)
    bias_params_.AddVec(alpha, other->bias_params_);
}

void TdnnComponent::PerturbParams(BaseFloat stddev) {
  CuMatrix<BaseFloat> temp_mat(linear_params_.NumRows(),
                               linear_params_.NumCols(), kUndefined);
  temp_mat.SetRandn();
  linear_params_.AddMat(stddev, temp_mat);
  if (bias_params_.Dim() != 0) {
    CuVector<BaseFloat> temp_vec(bias_params_.Dim(), kUndefined);
    temp_vec.SetRandn();
    bias_params_.AddVec(stddev, temp_vec);
  }
}

BaseFloat TdnnComponent::DotProduct(
    const UpdatableComponent &other_in) const {
  const TdnnComponent *other =
      dynamic_cast<const TdnnComponent*>(&other_in);
  KALDI_ASSERT(other != NULL);
  BaseFloat ans = TraceMatMat(linear_params_, other->linear_params_, kTrans);
  if (bias_params_.Dim() != 0)
    ans += VecVec(bias_params_, other->bias_params_);
  return ans;
}

int32 TdnnComponent::NumParameters() const {
  // note: bias_param_.Dim() may actually be zero.
  return linear_params_.NumRows() * linear_params_.NumCols() +
      bias_params_.Dim();
}

void TdnnComponent::Vectorize(
    VectorBase<BaseFloat> *params) const {
  KALDI_ASSERT(params->Dim() == NumParameters());
  int32 linear_size = linear_params_.NumRows() * linear_params_.NumCols(),
      bias_size = bias_params_.Dim();
  params->Range(0, linear_size).CopyRowsFromMat(linear_params_);
  if (bias_size != 0)
    params->Range(linear_size, bias_size).CopyFromVec(bias_params_);
}

void TdnnComponent::UnVectorize(
    const VectorBase<BaseFloat> &params) {
  KALDI_ASSERT(params.Dim() == NumParameters());
  int32 linear_size = linear_params_.NumRows() * linear_params_.NumCols(),
      bias_size = bias_params_.Dim();
  linear_params_.CopyRowsFromVec(params.Range(0, linear_size));
  if (bias_size != 0)
    bias_params_.CopyFromVec(params.Range(linear_size, bias_size));
}

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

TdnnComponent::PrecomputedIndexes*
TdnnComponent::PrecomputedIndexes::Copy() const {
  return new PrecomputedIndexes(*this);
}

void TdnnComponent::PrecomputedIndexes::Write(
    std::ostream &os, bool binary) const {
  WriteToken(os, binary, "<TdnnComponentPrecomputedIndexes>");
  WriteToken(os, binary, "<RowStride>");
  WriteBasicType(os, binary, row_stride);
  WriteToken(os, binary, "<RowOffsets>");
  WriteIntegerVector(os, binary, row_offsets);
  WriteToken(os, binary, "</TdnnComponentPrecomputedIndexes>");
}

void TdnnComponent::PrecomputedIndexes::Read(
    std::istream &is, bool binary) {
  ExpectOneOrTwoTokens(is, binary,
                       "<TdnnComponentPrecomputedIndexes>",
                       "<RowStride>");
  ReadBasicType(is, binary, &row_stride);
  ExpectToken(is, binary, "<RowOffsets>");
  ReadIntegerVector(is, binary, &row_offsets);
  ExpectToken(is, binary, "</TdnnComponentPrecomputedIndexes>");
}

void TdnnComponent::ConsolidateMemory() {
  OnlineNaturalGradient temp_in(preconditioner_in_);
  preconditioner_in_.Swap(&temp_in);
  OnlineNaturalGradient temp_out(preconditioner_out_);
  preconditioner_out_.Swap(&temp_out);
}

} // namespace nnet3
} // namespace kaldi