<commit_before>#include "cnn/nodes.h"

#include <limits>
#include <cmath>
#include <sstream>

using namespace std;

namespace cnn {

inline bool LooksLikeVector(const Dim& d) {
  if (d.ndims() == 1) return true;
  if (d.ndims() > 1) {
    for (int i = 1; i < d.ndims(); ++i)
      if (d[i] != 1) return false;
  }
  return true;
}

string Min::as_string(const vector<string>& arg_names) const {
  ostringstream s;
  s << "min{" << arg_names[0] << ", " << arg_names[1] << "}";
  return s.str();
}

Dim Min::dim_forward(const vector<Dim>& xs) const {
  if (xs.size() != 2 || xs[0] != xs[1]) {
    cerr << "Bad arguments in Min: " << xs << endl;
    throw std::invalid_argument("invalid arguments to Min");
  }
  return xs[0].bd >= xs[1].bd ? xs[0] : xs[1];
}

string Max::as_string(const vector<string>& arg_names) const {
  ostringstream s;
  s << "max{" << arg_names[0] << ", " << arg_names[1] << "}";
  return s.str();
}

Dim Max::dim_forward(const vector<Dim>& xs) const {
  if (xs.size() != 2 || xs[0] != xs[1]) {
    cerr << "Bad arguments in Max: " << xs << endl;
    throw std::invalid_argument("invalid arguments to Max");
  }
  return xs[0].bd >= xs[1].bd ? xs[0] : xs[1];
}

string TraceOfProduct::as_string(const vector<string>& arg_names) const {
  ostringstream s;
  s << "Tr(" << arg_names[0] << " * " << arg_names[1] << "^T)";
  return s.str();
}

Dim TraceOfProduct::dim_forward(const vector<Dim>& xs) const {
  if (xs.size() != 2 || xs[0] != xs[1]) {
    cerr << "Bad arguments in TraceOfProduct: " << xs << endl;
    throw std::invalid_argument("invalid arguments to TraceOfProduct");
  }
  return Dim({1}, max(xs[0].bd, xs[1].bd));
}

string ConstScalarMultiply::as_string(const vector<string>& arg_names) const {
  ostringstream s;
  s << arg_names[0] << " * " << alpha;
  return s.str();
}

Dim ConstScalarMultiply::dim_forward(const vector<Dim>& xs) const {
  if (xs.size() != 1) {
    cerr << "ConstScalarMultiply expects one argument: " << xs << endl;
    throw std::invalid_argument("ConstScalarMultiply expects one argument");
  }
  return xs[0];
}

string DotProduct::as_string(const vector<string>& arg_names) const {
  ostringstream s;
  s << arg_names[0] << "^T . " << arg_names[1];
  return s.str();
}

Dim DotProduct::dim_forward(const vector<Dim>& xs) const {
  if (xs.size() != 2 ||
      !LooksLikeVector(xs[0]) ||
      !LooksLikeVector(xs[1]) ||
      xs[0].rows() != xs[1].rows()) {
    cerr << "Bad arguments to DotProduct: " << xs << endl;
    throw std::invalid_argument("Bad arguments to DotProduct");
  }
  return Dim({1}, max(xs[0].bd, xs[1].bd));
}

string Transpose::as_string(const vector<string>& arg_names) const {
  ostringstream s;
  s << arg_names[0] << "^T";
  return s.str();
}

Dim Transpose::dim_forward(const vector<Dim>& xs) const {
  if (xs.size() != 1) {
    cerr << "Bad arguments to Transpose: " << xs << endl;
    throw std::invalid_argument("Bad arguments to Transpose");
  }
  return xs[0].transpose();
}

string Reshape::as_string(const vector<string>& arg_names) const {
  ostringstream s;
  s << "reshape(" << arg_names[0] << " --> " << to << ')';
  return s.str();
}

Dim Reshape::dim_forward(const vector<Dim>& xs) const {
  assert(xs.size() == 1);
  assert(xs[0].size() == to.size());
  return to;
}

string SumColumns::as_string(const vector<string>& arg_names) const {
  ostringstream s;
  s << "sum_cols(matrix=" << arg_names[0];
  if (arg_names.size() == 2) s << ", col_weighting=" << arg_names[1];
  s << ')';
  return s.str();
}

Dim SumColumns::dim_forward(const vector<Dim>& xs) const {
  assert(xs.size() == 1 || xs.size() == 2);
  int bd = (xs.size() == 1 ? xs[0].bd : max(xs[0].bd, xs[1].bd));
  return Dim({xs[0].rows()}, bd);

}

string KMHNGram::as_string(const vector<string>& arg_names) const {
  ostringstream s;
  s << "kmh-ngram(" << arg_names[0] << ')';
  return s.str();
}

Dim KMHNGram::dim_forward(const vector<Dim>& xs) const {
  assert(xs[0].ndims() == 2);
  const unsigned new_cols = xs[0].cols() - n + 1;
  if (new_cols < 1) {
    ostringstream s; s << "Bad input dimensions in KMHNGram: " << xs;
    throw std::invalid_argument(s.str());
  }
  return Dim({xs[0][0], new_cols});
}

string InnerProduct3D_1D::as_string(const vector<string>& arg_names) const {
  ostringstream s;
  s << "dot(" << arg_names[0] << "," << arg_names[1] << ')';
  if (arg_names.size() == 3) s << " + " << arg_names[2];
  return s.str();
}

Dim InnerProduct3D_1D::dim_forward(const vector<Dim>& xs) const {
  if (xs.size() != 2 && xs.size() != 3)
    throw std::invalid_argument("Expected two or three arguments in InnerProduct3D_1D");
  if (xs[0].ndims() != 3 ||
      xs[1].ndims() != 1 ||
      xs[0].size(2) != xs[1].size(0)) {
    ostringstream s; s << "Bad input dimensions in InnerProduct3D_1D: " << xs;
    throw std::invalid_argument(s.str());
  }
  Dim d({xs[0].size(0), xs[0].size(1)}, max(xs[0].bd, xs[1].bd));
  if(xs.size() == 3) d.bd = max(d.bd, xs[2].bd);
  if (xs.size() == 3 && xs[2] != d) {
    ostringstream s; s << "Bad input dimensions in InnerProduct3D_1D: " << xs;
    throw std::invalid_argument(s.str());
  }
  return d;
}

string GaussianNoise::as_string(const vector<string>& arg_names) const {
  ostringstream s;
  s << arg_names[0] << " + N(0," << stddev << ')';
  return s.str();
}

Dim GaussianNoise::dim_forward(const vector<Dim>& xs) const {
  assert(xs.size() == 1);
  return xs[0];
}

string Dropout::as_string(const vector<string>& arg_names) const {
  ostringstream s;
  s << "dropout(" << arg_names[0] << ",p=" << p << ')';
  return s.str();
}

Dim Dropout::dim_forward(const vector<Dim>& xs) const {
  assert(xs.size() == 1);
  return xs[0];
}

string BlockDropout::as_string(const vector<string>& arg_names) const {
  ostringstream s;
  s << "block_dropout(" << arg_names[0] << ",dropout_probability=" << dropout_probability << ')';
  return s.str();
}

Dim BlockDropout::dim_forward(const vector<Dim>& xs) const {
  assert(xs.size() == 1);
  return xs[0];
}

string ConstantPlusX::as_string(const vector<string>& arg_names) const {
  ostringstream s;
  s << c << " + " << arg_names[0];
  return s.str();
}

Dim ConstantPlusX::dim_forward(const vector<Dim>& xs) const {
  assert(xs.size() == 1);
  return xs[0];
}

string ConstantMinusX::as_string(const vector<string>& arg_names) const {
  ostringstream s;
  s << c << " - " << arg_names[0];
  return s.str();
}

Dim ConstantMinusX::dim_forward(const vector<Dim>& xs) const {
  assert(xs.size() == 1);
  return xs[0];
}

string LogSumExp::as_string(const vector<string>& arg_names) const {
  ostringstream s;
  s << "log(exp " << arg_names[0];
  for (unsigned i = 1; i < arg_names.size(); ++i)
    s << " + exp " << arg_names[i];
  s << ")";
  return s.str();
}

Dim LogSumExp::dim_forward(const vector<Dim>& xs) const {
  Dim d = xs[0].truncate();
  for (unsigned i = 1; i < xs.size(); ++i) {
    if (d.single_batch() != xs[i].truncate().single_batch()) {
      ostringstream s; s << "Mismatched input dimensions in LogSumExp: " << xs;
      throw std::invalid_argument(s.str());
    }
    d.bd = max(xs[i].bd, d.bd);
  }
  return d;
}
string Sum::as_string(const vector<string>& arg_names) const {
  ostringstream s;
  s << arg_names[0];
  for (unsigned i = 1; i < arg_names.size(); ++i)
    s << " + " << arg_names[i];
  return s.str();
}

Dim Sum::dim_forward(const vector<Dim>& xs) const {
  Dim d = xs[0].truncate();
  for (unsigned i = 1; i < xs.size(); ++i) {
    if (d.single_batch() != xs[i].truncate().single_batch()) {
      ostringstream s; s << "Mismatched input dimensions in Sum: " << xs;
      throw std::invalid_argument(s.str());
    }
    d.bd = max(xs[i].bd, d.bd);
  }
  return d;
}

string SumBatches::as_string(const vector<string>& arg_names) const {
  ostringstream s;
  s << "sum_batches( " << arg_names[0] << " )";
  return s.str();
}

Dim SumBatches::dim_forward(const vector<Dim>& xs) const {
  assert(xs.size() == 1);
  return xs[0].single_batch();
}

string Average::as_string(const vector<string>& arg_names) const {
  ostringstream s;
  s << "average(" << arg_names[0];
  for (unsigned i = 1; i < arg_names.size(); ++i)
    s << ", " << arg_names[i];
  s << ")";
  return s.str();
}

Dim Average::dim_forward(const vector<Dim>& xs) const {
  Dim d(xs[0]);
  for (unsigned i = 1; i < xs.size(); ++i) {
    if (xs[0].single_batch() != xs[1].single_batch()) {
      ostringstream s; s << "Mismatched input dimensions in Average: " << xs;
      throw std::invalid_argument(s.str());
    }
    d.bd = max(xs[i].bd, d.bd);
  }
  return d;
}

string Sqrt::as_string(const vector<string>& arg_names) const {
  ostringstream s;
  s << "sqrt(" << arg_names[0] << ')';
  return s.str();
}

Dim Sqrt::dim_forward(const vector<Dim>& xs) const {
  assert(xs.size() == 1);
  return xs[0];
}

/*string Erf::as_string(const vector<string>& arg_names) const {
  ostringstream s;
  s << "erf(" << arg_names[0] << ')';
  return s.str();
}

Dim Erf::dim_forward(const vector<Dim>& xs) const {
  assert(xs.size() == 1);
  return xs[0];
}
*/

string Tanh::as_string(const vector<string>& arg_names) const {
  ostringstream s;
  s << "tanh(" << arg_names[0] << ')';
  return s.str();
}

Dim Tanh::dim_forward(const vector<Dim>& xs) const {
  assert(xs.size() == 1);
  return xs[0];
}

string Square::as_string(const vector<string>& arg_names) const {
  ostringstream s;
  s << "square(" << arg_names[0] << ')';
  return s.str();
}

Dim Square::dim_forward(const vector<Dim>& xs) const {
  assert(xs.size() == 1);
  return xs[0];
}

string Cube::as_string(const vector<string>& arg_names) const {
  ostringstream s;
  s << "cube(" << arg_names[0] << ')';
  return s.str();
}

Dim Cube::dim_forward(const vector<Dim>& xs) const {
  assert(xs.size() == 1);
  return xs[0];
}

string Exp::as_string(const vector<string>& arg_names) const {
  ostringstream os;
  os << "exp(" << arg_names[0] << ')';
  return os.str();
}

Dim Exp::dim_forward(const vector<Dim>& xs) const {
  assert(xs.size() == 1);
  return xs[0];
}

/*
string LogGamma::as_string(const vector<string>& arg_names) const {
  ostringstream os;
  os << "lgamma(" << arg_names[0] << ')';
  return os.str();
}

Dim LogGamma::dim_forward(const vector<Dim>& xs) const {
  assert(xs.size() == 1);
  return xs[0];
}
*/

string Log::as_string(const vector<string>& arg_names) const {
  ostringstream os;
  os << "log(" << arg_names[0] << ')';
  return os.str();
}

Dim Log::dim_forward(const vector<Dim>& xs) const {
  assert(xs.size() == 1);
  return xs[0];
}

string Concatenate::as_string(const vector<string>& arg_names) const {
  ostringstream os;
  os << "concat(" << arg_names[0];
  for (unsigned i = 1; i < arg_names.size(); ++i) {
    os << ',' << arg_names[i];
  }
  os << ')';
  return os.str();
}

Dim Concatenate::dim_forward(const vector<Dim>& xs) const {
  unsigned new_rows = 0;
  Dim dr = xs[0];
  if (LooksLikeVector(dr)) dr.resize(1);
  for (auto c : xs) {
    if (LooksLikeVector(c)) c.resize(1);
    new_rows += c[0];
    dr.set(0, c[0]);
    if (dr != c) {
      ostringstream s; s << "Bad input dimensions in Concatenate: " << xs;
      throw std::invalid_argument(s.str());
    }
    dr.bd = max(dr.bd, c.bd);
  }
  dr.set(0, new_rows);
  return dr;
}

string ConcatenateColumns::as_string(const vector<string>& arg_names) const {
  ostringstream os;
  os << "concat_cols(" << arg_names[0];
  for (unsigned i = 1; i < arg_names.size(); ++i) {
    os << ',' << arg_names[i];
  }
  os << ')';
  return os.str();
}

Dim ConcatenateColumns::dim_forward(const vector<Dim>& xs) const {
  assert(xs.size() > 0);
  unsigned rows = xs[0][0];
  unsigned new_cols = 0;
  unsigned bd = 1;
  for (auto& d : xs) {
    if (d[0] != rows) {
      ostringstream s; s << "Bad input dimensions in ConcatenateColumns: " << xs;
      throw std::invalid_argument(s.str());
    }
    new_cols += d[1];
    bd = max(bd, d.bd);
  }
  return Dim({rows, new_cols}, bd);
}

string PairwiseRankLoss::as_string(const vector<string>& arg_names) const {
  ostringstream os;
  os << "max(0, " << margin << " - " << arg_names[0] << " + " << arg_names[1] << ')';
  return os.str();
}

Dim PairwiseRankLoss::dim_forward(const vector<Dim>& xs) const {
  if (xs.size() != 2 ||
      xs[0] != xs[1] ||
      xs[0].rows() != 1 ||
      (xs[0].ndims() != 1 && xs[0].ndims() != 2)) {
    ostringstream s; s << "Bad input dimensions in PairwiseRankLoss: " << xs;
    throw std::invalid_argument(s.str());
  }
  return xs[0].bd >= xs[1].bd ? xs[0] : xs[1];
}

string Hinge::as_string(const vector<string>& arg_names) const {
  ostringstream os;
  os << "hinge(" << arg_names[0] << ", pe=" << pelement << ", m=" << margin << ')';
  return os.str();
}

Dim Hinge::dim_forward(const vector<Dim>& xs) const {
  if (xs.size() != 1 || !LooksLikeVector(xs[0])) {
    ostringstream s; s << "Bad input dimensions in Hinge: " << xs;
    throw std::invalid_argument(s.str());
  }
  return Dim({1}, xs[0].bd);
}

string Identity::as_string(const vector<string>& arg_names) const {
  return arg_names[0];
}

Dim Identity::dim_forward(const vector<Dim>& xs) const {
  assert(xs.size() == 1);
  return xs[0];
}

string Softmax::as_string(const vector<string>& arg_names) const {
  ostringstream s;
  s << "softmax(" << arg_names[0] << ')';
  return s.str();
}

Dim Softmax::dim_forward(const vector<Dim>& xs) const {
  assert(xs.size() == 1);
  if (!LooksLikeVector(xs[0])) {
    ostringstream s; s << "Bad input dimensions in Softmax: " << xs;
    throw std::invalid_argument(s.str());
  }
  return xs[0];
}

string SoftSign::as_string(const vector<string>& arg_names) const {
  ostringstream s;
  s << "softsign(" << arg_names[0] << ')';
  return s.str();
}

Dim SoftSign::dim_forward(const vector<Dim>& xs) const {
  assert(xs.size() == 1);
  if (!LooksLikeVector(xs[0])) {
    ostringstream s; s << "Bad input dimensions in SoftSign: " << xs;
    throw std::invalid_argument(s.str());
  }
  return xs[0];
}

/*
string PickNegLogSoftmax::as_string(const vector<string>& arg_names) const {
  ostringstream s;
  if(pval) {
    s << "log_softmax(" << arg_names[0] << ")_{" << *pval << '}';
  } else {
    s << "log_softmax(" << arg_names[0] << ")_{";
    string sep = "";
    for(auto v : *pvals) { s << sep << v; sep = ","; }
    s << '}';
  }
  return s.str();
}

Dim PickNegLogSoftmax::dim_forward(const vector<Dim>& xs) const {
  assert(xs.size() == 1);
  if (!LooksLikeVector(xs[0])) {
    ostringstream s; s << "Bad input dimensions in PickNegLogSoftmax: " << xs;
    throw std::invalid_argument(s.str());
  }
  return Dim({1}, xs[0].bd);
}
*/

string LogSoftmax::as_string(const vector<string>& arg_names) const {
  ostringstream s;
  s << "log_softmax(" << arg_names[0] << ')';
  return s.str();
}

Dim LogSoftmax::dim_forward(const vector<Dim>& xs) const {
  assert(xs.size() == 1);
  if (!LooksLikeVector(xs[0])) {
    ostringstream s; s << "Bad input dimensions in LogSoftmax: " << xs;
    throw std::invalid_argument(s.str());
  }
  return xs[0];
}

string RestrictedLogSoftmax::as_string(const vector<string>& arg_names) const {
  ostringstream s;
  s << "r_log_softmax(" << arg_names[0] << ')';
  return s.str();
}

Dim RestrictedLogSoftmax::dim_forward(const vector<Dim>& xs) const {<commit_msg>allow data parallel for log-softmax<commit_after>#include "cnn/nodes.h"

#include <limits>
#include <cmath>
#include <sstream>

using namespace std;

namespace cnn {

    inline bool LooksLikeVector(const Dim& d) {
        if (d.ndims() == 1) return true;
        if (d.ndims() > 1) {
            for (int i = 1; i < d.ndims(); ++i)
            if (d[i] != 1) return false;
        }
        return true;
    }

    inline bool LooksLikeMatrix(const Dim& d) {
        if (d.ndims() == 2) return true;
        if (d.ndims() > 2) {
            for (int i = 2; i < d.ndims(); ++i)
            if (d[i] != 1) return false;
        }
        return true;
}

string Min::as_string(const vector<string>& arg_names) const {
  ostringstream s;
  s << "min{" << arg_names[0] << ", " << arg_names[1] << "}";
  return s.str();
}

Dim Min::dim_forward(const vector<Dim>& xs) const {
  if (xs.size() != 2 || xs[0] != xs[1]) {
    cerr << "Bad arguments in Min: " << xs << endl;
    throw std::invalid_argument("invalid arguments to Min");
  }
  return xs[0].bd >= xs[1].bd ? xs[0] : xs[1];
}

string Max::as_string(const vector<string>& arg_names) const {
  ostringstream s;
  s << "max{" << arg_names[0] << ", " << arg_names[1] << "}";
  return s.str();
}

Dim Max::dim_forward(const vector<Dim>& xs) const {
  if (xs.size() != 2 || xs[0] != xs[1]) {
    cerr << "Bad arguments in Max: " << xs << endl;
    throw std::invalid_argument("invalid arguments to Max");
  }
  return xs[0].bd >= xs[1].bd ? xs[0] : xs[1];
}

string TraceOfProduct::as_string(const vector<string>& arg_names) const {
  ostringstream s;
  s << "Tr(" << arg_names[0] << " * " << arg_names[1] << "^T)";
  return s.str();
}

Dim TraceOfProduct::dim_forward(const vector<Dim>& xs) const {
  if (xs.size() != 2 || xs[0] != xs[1]) {
    cerr << "Bad arguments in TraceOfProduct: " << xs << endl;
    throw std::invalid_argument("invalid arguments to TraceOfProduct");
  }
  return Dim({1}, max(xs[0].bd, xs[1].bd));
}

string ConstScalarMultiply::as_string(const vector<string>& arg_names) const {
  ostringstream s;
  s << arg_names[0] << " * " << alpha;
  return s.str();
}

Dim ConstScalarMultiply::dim_forward(const vector<Dim>& xs) const {
  if (xs.size() != 1) {
    cerr << "ConstScalarMultiply expects one argument: " << xs << endl;
    throw std::invalid_argument("ConstScalarMultiply expects one argument");
  }
  return xs[0];
}

string DotProduct::as_string(const vector<string>& arg_names) const {
  ostringstream s;
  s << arg_names[0] << "^T . " << arg_names[1];
  return s.str();
}

Dim DotProduct::dim_forward(const vector<Dim>& xs) const {
  if (xs.size() != 2 ||
      !LooksLikeVector(xs[0]) ||
      !LooksLikeVector(xs[1]) ||
      xs[0].rows() != xs[1].rows()) {
    cerr << "Bad arguments to DotProduct: " << xs << endl;
    throw std::invalid_argument("Bad arguments to DotProduct");
  }
  return Dim({1}, max(xs[0].bd, xs[1].bd));
}

string Transpose::as_string(const vector<string>& arg_names) const {
  ostringstream s;
  s << arg_names[0] << "^T";
  return s.str();
}

Dim Transpose::dim_forward(const vector<Dim>& xs) const {
  if (xs.size() != 1) {
    cerr << "Bad arguments to Transpose: " << xs << endl;
    throw std::invalid_argument("Bad arguments to Transpose");
  }
  return xs[0].transpose();
}

string Reshape::as_string(const vector<string>& arg_names) const {
  ostringstream s;
  s << "reshape(" << arg_names[0] << " --> " << to << ')';
  return s.str();
}

Dim Reshape::dim_forward(const vector<Dim>& xs) const {
  assert(xs.size() == 1);
  assert(xs[0].size() == to.size());
  return to;
}

string SumColumns::as_string(const vector<string>& arg_names) const {
  ostringstream s;
  s << "sum_cols(matrix=" << arg_names[0];
  if (arg_names.size() == 2) s << ", col_weighting=" << arg_names[1];
  s << ')';
  return s.str();
}

Dim SumColumns::dim_forward(const vector<Dim>& xs) const {
  assert(xs.size() == 1 || xs.size() == 2);
  int bd = (xs.size() == 1 ? xs[0].bd : max(xs[0].bd, xs[1].bd));
  return Dim({xs[0].rows()}, bd);

}

string KMHNGram::as_string(const vector<string>& arg_names) const {
  ostringstream s;
  s << "kmh-ngram(" << arg_names[0] << ')';
  return s.str();
}

Dim KMHNGram::dim_forward(const vector<Dim>& xs) const {
  assert(xs[0].ndims() == 2);
  const unsigned new_cols = xs[0].cols() - n + 1;
  if (new_cols < 1) {
    ostringstream s; s << "Bad input dimensions in KMHNGram: " << xs;
    throw std::invalid_argument(s.str());
  }
  return Dim({xs[0][0], new_cols});
}

string InnerProduct3D_1D::as_string(const vector<string>& arg_names) const {
  ostringstream s;
  s << "dot(" << arg_names[0] << "," << arg_names[1] << ')';
  if (arg_names.size() == 3) s << " + " << arg_names[2];
  return s.str();
}

Dim InnerProduct3D_1D::dim_forward(const vector<Dim>& xs) const {
  if (xs.size() != 2 && xs.size() != 3)
    throw std::invalid_argument("Expected two or three arguments in InnerProduct3D_1D");
  if (xs[0].ndims() != 3 ||
      xs[1].ndims() != 1 ||
      xs[0].size(2) != xs[1].size(0)) {
    ostringstream s; s << "Bad input dimensions in InnerProduct3D_1D: " << xs;
    throw std::invalid_argument(s.str());
  }
  Dim d({xs[0].size(0), xs[0].size(1)}, max(xs[0].bd, xs[1].bd));
  if(xs.size() == 3) d.bd = max(d.bd, xs[2].bd);
  if (xs.size() == 3 && xs[2] != d) {
    ostringstream s; s << "Bad input dimensions in InnerProduct3D_1D: " << xs;
    throw std::invalid_argument(s.str());
  }
  return d;
}

string GaussianNoise::as_string(const vector<string>& arg_names) const {
  ostringstream s;
  s << arg_names[0] << " + N(0," << stddev << ')';
  return s.str();
}

Dim GaussianNoise::dim_forward(const vector<Dim>& xs) const {
  assert(xs.size() == 1);
  return xs[0];
}

string Dropout::as_string(const vector<string>& arg_names) const {
  ostringstream s;
  s << "dropout(" << arg_names[0] << ",p=" << p << ')';
  return s.str();
}

Dim Dropout::dim_forward(const vector<Dim>& xs) const {
  assert(xs.size() == 1);
  return xs[0];
}

string BlockDropout::as_string(const vector<string>& arg_names) const {
  ostringstream s;
  s << "block_dropout(" << arg_names[0] << ",dropout_probability=" << dropout_probability << ')';
  return s.str();
}

Dim BlockDropout::dim_forward(const vector<Dim>& xs) const {
  assert(xs.size() == 1);
  return xs[0];
}

string ConstantPlusX::as_string(const vector<string>& arg_names) const {
  ostringstream s;
  s << c << " + " << arg_names[0];
  return s.str();
}

Dim ConstantPlusX::dim_forward(const vector<Dim>& xs) const {
  assert(xs.size() == 1);
  return xs[0];
}

string ConstantMinusX::as_string(const vector<string>& arg_names) const {
  ostringstream s;
  s << c << " - " << arg_names[0];
  return s.str();
}

Dim ConstantMinusX::dim_forward(const vector<Dim>& xs) const {
  assert(xs.size() == 1);
  return xs[0];
}

string LogSumExp::as_string(const vector<string>& arg_names) const {
  ostringstream s;
  s << "log(exp " << arg_names[0];
  for (unsigned i = 1; i < arg_names.size(); ++i)
    s << " + exp " << arg_names[i];
  s << ")";
  return s.str();
}

Dim LogSumExp::dim_forward(const vector<Dim>& xs) const {
  Dim d = xs[0].truncate();
  for (unsigned i = 1; i < xs.size(); ++i) {
    if (d.single_batch() != xs[i].truncate().single_batch()) {
      ostringstream s; s << "Mismatched input dimensions in LogSumExp: " << xs;
      throw std::invalid_argument(s.str());
    }
    d.bd = max(xs[i].bd, d.bd);
  }
  return d;
}
string Sum::as_string(const vector<string>& arg_names) const {
  ostringstream s;
  s << arg_names[0];
  for (unsigned i = 1; i < arg_names.size(); ++i)
    s << " + " << arg_names[i];
  return s.str();
}

Dim Sum::dim_forward(const vector<Dim>& xs) const {
  Dim d = xs[0].truncate();
  for (unsigned i = 1; i < xs.size(); ++i) {
    if (d.single_batch() != xs[i].truncate().single_batch()) {
      ostringstream s; s << "Mismatched input dimensions in Sum: " << xs;
      throw std::invalid_argument(s.str());
    }
    d.bd = max(xs[i].bd, d.bd);
  }
  return d;
}

string SumBatches::as_string(const vector<string>& arg_names) const {
  ostringstream s;
  s << "sum_batches( " << arg_names[0] << " )";
  return s.str();
}

Dim SumBatches::dim_forward(const vector<Dim>& xs) const {
  assert(xs.size() == 1);
  return xs[0].single_batch();
}

string Average::as_string(const vector<string>& arg_names) const {
  ostringstream s;
  s << "average(" << arg_names[0];
  for (unsigned i = 1; i < arg_names.size(); ++i)
    s << ", " << arg_names[i];
  s << ")";
  return s.str();
}

Dim Average::dim_forward(const vector<Dim>& xs) const {
  Dim d(xs[0]);
  for (unsigned i = 1; i < xs.size(); ++i) {
    if (xs[0].single_batch() != xs[1].single_batch()) {
      ostringstream s; s << "Mismatched input dimensions in Average: " << xs;
      throw std::invalid_argument(s.str());
    }
    d.bd = max(xs[i].bd, d.bd);
  }
  return d;
}

string Sqrt::as_string(const vector<string>& arg_names) const {
  ostringstream s;
  s << "sqrt(" << arg_names[0] << ')';
  return s.str();
}

Dim Sqrt::dim_forward(const vector<Dim>& xs) const {
  assert(xs.size() == 1);
  return xs[0];
}

/*string Erf::as_string(const vector<string>& arg_names) const {
  ostringstream s;
  s << "erf(" << arg_names[0] << ')';
  return s.str();
}

Dim Erf::dim_forward(const vector<Dim>& xs) const {
  assert(xs.size() == 1);
  return xs[0];
}
*/

string Tanh::as_string(const vector<string>& arg_names) const {
  ostringstream s;
  s << "tanh(" << arg_names[0] << ')';
  return s.str();
}

Dim Tanh::dim_forward(const vector<Dim>& xs) const {
  assert(xs.size() == 1);
  return xs[0];
}

string Square::as_string(const vector<string>& arg_names) const {
  ostringstream s;
  s << "square(" << arg_names[0] << ')';
  return s.str();
}

Dim Square::dim_forward(const vector<Dim>& xs) const {
  assert(xs.size() == 1);
  return xs[0];
}

string Cube::as_string(const vector<string>& arg_names) const {
  ostringstream s;
  s << "cube(" << arg_names[0] << ')';
  return s.str();
}

Dim Cube::dim_forward(const vector<Dim>& xs) const {
  assert(xs.size() == 1);
  return xs[0];
}

string Exp::as_string(const vector<string>& arg_names) const {
  ostringstream os;
  os << "exp(" << arg_names[0] << ')';
  return os.str();
}

Dim Exp::dim_forward(const vector<Dim>& xs) const {
  assert(xs.size() == 1);
  return xs[0];
}

/*
string LogGamma::as_string(const vector<string>& arg_names) const {
  ostringstream os;
  os << "lgamma(" << arg_names[0] << ')';
  return os.str();
}

Dim LogGamma::dim_forward(const vector<Dim>& xs) const {
  assert(xs.size() == 1);
  return xs[0];
}
*/

string Log::as_string(const vector<string>& arg_names) const {
  ostringstream os;
  os << "log(" << arg_names[0] << ')';
  return os.str();
}

Dim Log::dim_forward(const vector<Dim>& xs) const {
  assert(xs.size() == 1);
  return xs[0];
}

string Concatenate::as_string(const vector<string>& arg_names) const {
  ostringstream os;
  os << "concat(" << arg_names[0];
  for (unsigned i = 1; i < arg_names.size(); ++i) {
    os << ',' << arg_names[i];
  }
  os << ')';
  return os.str();
}

Dim Concatenate::dim_forward(const vector<Dim>& xs) const {
  unsigned new_rows = 0;
  Dim dr = xs[0];
  if (LooksLikeVector(dr)) dr.resize(1);
  for (auto c : xs) {
    if (LooksLikeVector(c)) c.resize(1);
    new_rows += c[0];
    dr.set(0, c[0]);
    if (dr != c) {
      ostringstream s; s << "Bad input dimensions in Concatenate: " << xs;
      throw std::invalid_argument(s.str());
    }
    dr.bd = max(dr.bd, c.bd);
  }
  dr.set(0, new_rows);
  return dr;
}

string ConcatenateColumns::as_string(const vector<string>& arg_names) const {
  ostringstream os;
  os << "concat_cols(" << arg_names[0];
  for (unsigned i = 1; i < arg_names.size(); ++i) {
    os << ',' << arg_names[i];
  }
  os << ')';
  return os.str();
}

Dim ConcatenateColumns::dim_forward(const vector<Dim>& xs) const {
  assert(xs.size() > 0);
  unsigned rows = xs[0][0];
  unsigned new_cols = 0;
  unsigned bd = 1;
  for (auto& d : xs) {
    if (d[0] != rows) {
      ostringstream s; s << "Bad input dimensions in ConcatenateColumns: " << xs;
      throw std::invalid_argument(s.str());
    }
    new_cols += d[1];
    bd = max(bd, d.bd);
  }
  return Dim({rows, new_cols}, bd);
}

string PairwiseRankLoss::as_string(const vector<string>& arg_names) const {
  ostringstream os;
  os << "max(0, " << margin << " - " << arg_names[0] << " + " << arg_names[1] << ')';
  return os.str();
}

Dim PairwiseRankLoss::dim_forward(const vector<Dim>& xs) const {
  if (xs.size() != 2 ||
      xs[0] != xs[1] ||
      xs[0].rows() != 1 ||
      (xs[0].ndims() != 1 && xs[0].ndims() != 2)) {
    ostringstream s; s << "Bad input dimensions in PairwiseRankLoss: " << xs;
    throw std::invalid_argument(s.str());
  }
  return xs[0].bd >= xs[1].bd ? xs[0] : xs[1];
}

string Hinge::as_string(const vector<string>& arg_names) const {
  ostringstream os;
  os << "hinge(" << arg_names[0] << ", pe=" << pelement << ", m=" << margin << ')';
  return os.str();
}

Dim Hinge::dim_forward(const vector<Dim>& xs) const {
  if (xs.size() != 1 || !LooksLikeVector(xs[0])) {
    ostringstream s; s << "Bad input dimensions in Hinge: " << xs;
    throw std::invalid_argument(s.str());
  }
  return Dim({1}, xs[0].bd);
}

string Identity::as_string(const vector<string>& arg_names) const {
  return arg_names[0];
}

Dim Identity::dim_forward(const vector<Dim>& xs) const {
  assert(xs.size() == 1);
  return xs[0];
}

string Softmax::as_string(const vector<string>& arg_names) const {
  ostringstream s;
  s << "softmax(" << arg_names[0] << ')';
  return s.str();
}

Dim Softmax::dim_forward(const vector<Dim>& xs) const {
  assert(xs.size() == 1);
  if (!LooksLikeVector(xs[0])) {
    ostringstream s; s << "Bad input dimensions in Softmax: " << xs;
    throw std::invalid_argument(s.str());
  }
  return xs[0];
}

string SoftSign::as_string(const vector<string>& arg_names) const {
  ostringstream s;
  s << "softsign(" << arg_names[0] << ')';
  return s.str();
}

Dim SoftSign::dim_forward(const vector<Dim>& xs) const {
  assert(xs.size() == 1);
  if (!LooksLikeVector(xs[0])) {
    ostringstream s; s << "Bad input dimensions in SoftSign: " << xs;
    throw std::invalid_argument(s.str());
  }
  return xs[0];
}

/*
string PickNegLogSoftmax::as_string(const vector<string>& arg_names) const {
  ostringstream s;
  if(pval) {
    s << "log_softmax(" << arg_names[0] << ")_{" << *pval << '}';
  } else {
    s << "log_softmax(" << arg_names[0] << ")_{";
    string sep = "";
    for(auto v : *pvals) { s << sep << v; sep = ","; }
    s << '}';
  }
  return s.str();
}

Dim PickNegLogSoftmax::dim_forward(const vector<Dim>& xs) const {
  assert(xs.size() == 1);
  if (!LooksLikeVector(xs[0])) {
    ostringstream s; s << "Bad input dimensions in PickNegLogSoftmax: " << xs;
    throw std::invalid_argument(s.str());
  }
  return Dim({1}, xs[0].bd);
}
*/

string LogSoftmax::as_string(const vector<string>& arg_names) const {
  ostringstream s;
  s << "log_softmax(" << arg_names[0] << ')';
  return s.str();
}

Dim LogSoftmax::dim_forward(const vector<Dim>& xs) const {
    assert(xs.size() == 1 || xs.size() == 2);
    if (!LooksLikeVector(xs[0]) && !LooksLikeMatrix(xs[0])) {
        ostringstream s; s << "Bad input dimensions in PickElement: " << xs;
        throw std::invalid_argument(s.str());
    }
    return xs[0];
}

string RestrictedLogSoftmax::as_string(const vector<string>& arg_names) const {
  ostringstream s;
  s << "r_log_softmax(" << arg_names[0] << ')';
  return s.str();
}

Dim RestrictedLogSoftmax::dim_forward(const vector<Dim>& xs) const {
{'fraction_non_alphanumeric': 0.2009748172217709, 'fraction_numerical': 0.014492282696994313, 'mean_word_length': 3.8688498655276065, 'pattern_counts': {'":': 0, '<': 962, '<?xml version=': 0, '>': 170, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4696352', 'n_tokens_mistral': 12065, 'n_tokens_neox': 11575, 'n_words': 2682}