#include "dynet/tensor-eigen.h"
#include "dynet/nodes-lstm.h"
#include "dynet/matrix-multiply.h"

#include "dynet/functors.h"
#include "dynet/simd-functors.h"
#include "dynet/nodes-impl-macros.h"

using namespace std;

namespace dynet {


// ************* LSTM Gates *************

#ifndef __CUDACC__

  string VanillaLSTMGates::as_string(const vector<string>& arg_names) const {
    ostringstream s;
    s << "vanilla_lstm_gates(" << arg_names[0] << ", " << arg_names[1] << ", " << arg_names[2] << ", " << arg_names[3] << ", " << arg_names[4] << ')';
    return s.str();
  }

  Dim VanillaLSTMGates::dim_forward(const vector<Dim>& xs) const {
    if(dropout){
      DYNET_ARG_CHECK(xs.size() >= 7, "Failed input count check in VanillaLSTMGates");
    }else{
      DYNET_ARG_CHECK(xs.size() >= 5, "Failed input count check in VanillaLSTMGates");
    }
    unsigned num_inputs = dropout?xs.size()-6:xs.size()-4;
    unsigned hidden_dim=xs[num_inputs][0];
    unsigned input_dim=xs[num_inputs+1][1];
    unsigned batch_size=xs[0].bd;
    unsigned inputs_dim_sum=0;
    for(unsigned i=0; i<num_inputs; i++){
        DYNET_ARG_CHECK(xs[i].ndims() == 1, "VanillaLSTMGates: x_t[" << i << "] expected to be a vector");
        DYNET_ARG_CHECK(xs[i].bd == batch_size, "VanillaLSTMGates: x_t has inconsistent batch size");
        inputs_dim_sum += xs[i][0];
    }
    DYNET_ARG_CHECK(inputs_dim_sum == input_dim, "VanillaLSTMGates: x_t has inconsistent dimension");
    DYNET_ARG_CHECK(xs[num_inputs].ndims() == 1, "VanillaLSTMGates: h_tm1 expected to be a vector");
    DYNET_ARG_CHECK(xs[num_inputs+1].ndims() == 2, "VanillaLSTMGates: Wx expected to be a matrix");
    DYNET_ARG_CHECK(xs[num_inputs+2].ndims() == 2, "VanillaLSTMGates: Wh expected to be a matrix");
    DYNET_ARG_CHECK(xs[num_inputs+3].ndims() == 1, "VanillaLSTMGates: b expected to be a vector");
    DYNET_ARG_CHECK(xs[num_inputs+1][0] == hidden_dim * 4, "VanillaLSTMGates: Wx dim 0 expected " << hidden_dim * 4 << ", was " << xs[2][0]);
    DYNET_ARG_CHECK(xs[num_inputs+1][1] == input_dim, "VanillaLSTMGates: Wx dim 1 expected " << input_dim << ", was " << xs[2][1]);
    DYNET_ARG_CHECK(xs[num_inputs+2][0] == hidden_dim * 4, "VanillaLSTMGates: Wh dim 0 expected " << hidden_dim * 4 << ", was " << xs[3][0]);
    DYNET_ARG_CHECK(xs[num_inputs+2][1] == hidden_dim, "VanillaLSTMGates: Wh dim 1 expected " << hidden_dim << ", was " << xs[3][1]);
    DYNET_ARG_CHECK(xs[num_inputs+3][0] == hidden_dim * 4, "VanillaLSTMGates: b dim expected " << hidden_dim * 4 << ", was " << xs[4][0]);
    if(dropout){
      DYNET_ARG_CHECK(xs[num_inputs+4].ndims() == 1, "VanillaLSTMGates: dropout_mask_x expected to be a vector");
      DYNET_ARG_CHECK(xs[num_inputs+5].ndims() == 1, "VanillaLSTMGates: dropout_mask_h expected to be a vector");
      DYNET_ARG_CHECK(xs[num_inputs+4].bd == batch_size || xs[num_inputs+4].bd == 1, "VanillaLSTMGates: dropout_mask_x expected to have batch size 1 or " << batch_size << ", was " << xs[5].bd);
      DYNET_ARG_CHECK(xs[num_inputs+5].bd == batch_size || xs[num_inputs+5].bd == 1, "VanillaLSTMGates: dropout_mask_h expected to have batch size 1 or " << batch_size << ", was " << xs[6].bd);
      DYNET_ARG_CHECK(xs[num_inputs+4][0] == input_dim, "VanillaLSTMGates: dropout_mask_x dim 1 expected " << input_dim << ", was " << xs[5][0]);
      DYNET_ARG_CHECK(xs[num_inputs+5][0] == hidden_dim, "VanillaLSTMGates: dropout_mask_h dim 1 expected " << hidden_dim << ", was " << xs[6][0]);
    }
    return Dim({hidden_dim*4}, batch_size);
  }

  int VanillaLSTMGates::autobatch_sig(const ComputationGraph & cg, SigMap &sm) const {
    Sig s(nt::vanilla_lstm_gates);
    unsigned num_inputs = dropout?args.size()-6:args.size()-4;
    // Assume parameter vectors must be same
    if(dim.bd == 1) {
      for(unsigned i=0; i<num_inputs; i++)
        s.add_dim(cg.nodes[args[i]]->dim);
      // TODO: correct? parameter vectors would be args[num_inputs+1] .. args[num_inputs+3]
      // s.add_dim(cg.nodes[args[num_inputs]]->dim); // not necessary, as will be the same
      s.add_node(args[num_inputs+1]);
      s.add_node(args[num_inputs+2]);
      s.add_node(args[num_inputs+3]);
    } else {
      for(auto nid : args) {
        const Dim & d = cg.nodes[nid]->dim;
        if(d.bd == 1)
          s.add_node(nid);
        else
          s.add_dim(d);
      }
    }
    return sm.get_idx(s);
  }

  std::vector<int> VanillaLSTMGates::autobatch_concat(const ComputationGraph & cg) const {
    vector<int> ret(args.size(), 0);
    unsigned num_inputs = dropout?args.size()-6:args.size()-4;
    if(dim.bd == 1) {
      ret[0] = ret[1] = 1;
      if(dropout)
        ret[num_inputs+4] = ret[num_inputs+5] = 1;
    } else {
      for(size_t i = 0; i < ret.size(); ++i)
        ret[i] = (cg.nodes[args[i]]->dim.bd > 1);
    }
    return ret;
  }

#endif

  template<class MyDevice>
  void VanillaLSTMGates::forward_dev_impl(const MyDevice & dev, const vector<const Tensor*>& xs, Tensor& fx) const {
    // computes affine transforms + nonlinearity
    // gates_i = sigmoid (Wx_i * x_t + Wh_i * h_tm1 + b_i)
    // gates_f = sigmoid (Wx_f * x_t + Wh_f * h_tm1 + b_f + 1)
    // gates_o = sigmoid (Wx_o * x_t + Wh_o * h_tm1 + b_o)
    // gates_g =   tanh  (Wx_g * x_t + Wh_g * h_tm1 + b_g)

    unsigned num_inputs = dropout?xs.size()-6:xs.size()-4;
    Tensor h_tm1(xs[num_inputs]->d, xs[num_inputs]->v, xs[num_inputs]->device, xs[num_inputs]->mem_pool);
    const Tensor *Wx = xs[num_inputs+1];
    const Tensor *Wh = xs[num_inputs+2];
    const Tensor *b  = xs[num_inputs+3];

    unsigned hidden_dim = h_tm1.d[0];
    unsigned input_dim = Wx->d[1];
    unsigned batch_size = xs[0]->d.bd;

    Eigen::DSizes<ptrdiff_t, 2> indices_i(0, 0);
    Eigen::DSizes<ptrdiff_t, 2> indices_f(hidden_dim,0);
    Eigen::DSizes<ptrdiff_t, 2> indices_g(hidden_dim*3,0);
    Eigen::DSizes<ptrdiff_t, 2> sizes_1(hidden_dim, static_cast<ptrdiff_t>(fx.d.bd));
    Eigen::DSizes<ptrdiff_t, 2> sizes_3(hidden_dim*3, static_cast<ptrdiff_t>(fx.d.bd));

    AlignedMemoryPool* scratch_allocator = fx.device->pools[(int)DeviceMempool::SCS];

    Tensor x_t(Dim({input_dim}, batch_size), nullptr, fx.device, fx.mem_pool);
    if(num_inputs==1){
      x_t.v = xs[0]->v;
    } else {
      Tensor tmp(Dim({input_dim}, batch_size), nullptr, fx.device, fx.mem_pool);
      tmp.v = static_cast<float*>(scratch_allocator->allocate(tmp.d.size() * sizeof(float)));
      Eigen::DSizes<ptrdiff_t, 2> indices_tmp(0, 0);
      Eigen::DSizes<ptrdiff_t, 2> sizes_tmp(0, static_cast<ptrdiff_t>(fx.d.bd));
      for(unsigned i=0; i<num_inputs; i++){
        sizes_tmp[0] = xs[i]->d[0];
        tbvec(tmp).slice(indices_tmp, sizes_tmp).device(*dev.edevice) = tbvec(*xs[i]);
        indices_tmp[0] += xs[i]->d[0];
      }
      x_t.v = tmp.v;
    }

    //bias
#ifdef __CUDACC__
    Eigen::array<ptrdiff_t, 3> bcast = {1, 1, batch_size};
    tb<2>(fx).device(*dev.edevice) = tb<2>(*b).broadcast(bcast);
#else
    float *curr_ptr = fx.v, *end_ptr = curr_ptr + fx.d.size(), *in_ptr = b->v;
    do {
      memcpy(curr_ptr, in_ptr, sizeof(float)*b->d[0]);
      curr_ptr += b->d[0];
    } while(curr_ptr != end_ptr);
#endif
    tbvec(fx).slice(indices_f, sizes_1).device(*dev.edevice) += tbvec(fx).slice(indices_f, sizes_1).constant(forget_gate_bias);

    if(dropout){
      Tensor mask_x(Dim({input_dim}, batch_size), nullptr, fx.device, fx.mem_pool);
      if(xs[num_inputs+4]->d.bd == batch_size){
        mask_x.v = xs[num_inputs+4]->v;
      } else {
        mask_x.v = static_cast<float*>(scratch_allocator->allocate(mask_x.d.size() * sizeof(float)));
#ifdef __CUDACC__
        Eigen::array<ptrdiff_t, 2> bcast = {1, batch_size};
        tbvec(mask_x).device(*dev.edevice) = tbvec(*xs[num_inputs+4]).broadcast(bcast);
#else
        float *curr_ptr = mask_x.v, *end_ptr = curr_ptr + mask_x.d.size(), *in_ptr = xs[num_inputs+4]->v;
        do {
          memcpy(curr_ptr, in_ptr, sizeof(float)*xs[num_inputs+4]->d[0]);
          curr_ptr += xs[num_inputs+4]->d[0];
        } while(curr_ptr != end_ptr);
#endif
      }
      Tensor x_t_dropped(Dim({input_dim}, batch_size), nullptr, fx.device, fx.mem_pool);
      x_t_dropped.v = static_cast<float*>(scratch_allocator->allocate(x_t_dropped.d.size() * sizeof(float)));
      tvec(x_t_dropped).device(*dev.edevice) = tvec(x_t) * tvec(mask_x);
      x_t.v = x_t_dropped.v;

      Tensor mask_h(Dim({hidden_dim}, batch_size), nullptr, fx.device, fx.mem_pool);
      if(xs[num_inputs+5]->d.bd == batch_size){
        mask_h.v = xs[num_inputs+5]->v;
      } else {
        mask_h.v = static_cast<float*>(scratch_allocator->allocate(mask_h.d.size() * sizeof(float)));
#ifdef __CUDACC__
        Eigen::array<ptrdiff_t, 2> bcast = {1, batch_size};
        tbvec(mask_h).device(*dev.edevice) = tbvec(*xs[num_inputs+5]).broadcast(bcast);
#else
        float *curr_ptr = mask_h.v, *end_ptr = curr_ptr + mask_h.d.size(), *in_ptr = xs[num_inputs+5]->v;
        do {
          memcpy(curr_ptr, in_ptr, sizeof(float)*xs[num_inputs+5]->d[0]);
          curr_ptr += xs[num_inputs+5]->d[0];
        } while(curr_ptr != end_ptr);
#endif
      }
      Tensor h_tm1_dropped(Dim({hidden_dim}, batch_size), nullptr, fx.device, fx.mem_pool);
      h_tm1_dropped.v = static_cast<float*>(scratch_allocator->allocate(h_tm1_dropped.d.size() * sizeof(float)));
      tvec(h_tm1_dropped).device(*dev.edevice) = tvec(h_tm1) * tvec(mask_h);
      h_tm1.v = h_tm1_dropped.v;
    }
    //matrix mult
    if(weightnoise_std > 0.f){
      Tensor Wx_noisy(Dim({hidden_dim*4, input_dim},1), nullptr, fx.device, fx.mem_pool);
      Wx_noisy.v = static_cast<float*>(scratch_allocator->allocate(Wx_noisy.d.size() * sizeof(float)));
      TensorTools::randomize_normal(Wx_noisy, 0, weightnoise_std);
      tvec(Wx_noisy).device(*dev.edevice) += tvec(*Wx);

      Tensor Wh_noisy(Dim({hidden_dim*4, hidden_dim},1), nullptr, fx.device, fx.mem_pool);
      Wh_noisy.v = static_cast<float*>(scratch_allocator->allocate(Wh_noisy.d.size() * sizeof(float)));
      TensorTools::randomize_normal(Wh_noisy, 0, weightnoise_std);
      tvec(Wh_noisy).device(*dev.edevice) += tvec(*Wh);

      Tensor b_noisy(Dim({hidden_dim*4, 1},1), nullptr, fx.device, fx.mem_pool);
      b_noisy.v = static_cast<float*>(scratch_allocator->allocate(b_noisy.d.size() * sizeof(float)));
      TensorTools::randomize_normal(b_noisy, 0, weightnoise_std);
      tvec(b_noisy).device(*dev.edevice) += tvec(*b);

    } else {
      MatrixMultiply(dev, *Wx, x_t, fx, dev.kSCALAR_ONE);
      MatrixMultiply(dev, *Wh, h_tm1, fx, dev.kSCALAR_ONE);
    }

    // non-linearities
    Tensor fx_ifo(Dim({hidden_dim*3, 1},batch_size), nullptr, fx.device, fx.mem_pool);
    fx_ifo.v = static_cast<float*>(scratch_allocator->allocate(fx_ifo.d.size() * sizeof(float)));
    tbvec(fx_ifo).device(*dev.edevice) = tbvec(fx).slice(indices_i, sizes_3);
    tbvec(fx_ifo).device(*dev.edevice) = tbvec(fx_ifo).unaryExpr(scalar_logistic_sigmoid_op<float>());
    tbvec(fx).slice(indices_i, sizes_3).device(*dev.edevice) = tbvec(fx_ifo);

    Tensor fx_g(Dim({hidden_dim*1, 1},batch_size), nullptr, fx.device, fx.mem_pool);
    fx_g.v = static_cast<float*>(scratch_allocator->allocate(fx_g.d.size() * sizeof(float)));
    tbvec(fx_g).device(*dev.edevice) = tbvec(fx).slice(indices_g, sizes_1);
    tbvec(fx_g).device(*dev.edevice) = tbvec(fx_g).tanh();
    tbvec(fx).slice(indices_g, sizes_1).device(*dev.edevice) = tbvec(fx_g);

    scratch_allocator->free();
  }

  template<class MyDevice>
  void VanillaLSTMGates::backward_dev_impl(const MyDevice & dev,
                               const vector<const Tensor*>& xs,
                               const Tensor& fx,
                               const Tensor& dEdf,
                               unsigned i,
                               Tensor& dEdxi) const {
    unsigned num_inputs = dropout?xs.size()-6:xs.size()-4;
    Tensor h_tm1(xs[num_inputs]->d, xs[num_inputs]->v, xs[num_inputs]->device, xs[num_inputs]->mem_pool);
    const Tensor *Wx = xs[num_inputs+1];
    const Tensor *Wh = xs[num_inputs+2];
//    const Tensor *b  = xs[num_inputs+3];

    unsigned hidden_dim = fx.d[0] / 4;
    unsigned input_dim = Wx->d[1];
    unsigned batch_size = xs[0]->d.bd;
    Eigen::DSizes<ptrdiff_t, 3> indices_mat_i(0, 0, 0);
    Eigen::DSizes<ptrdiff_t, 3> indices_mat_g(hidden_dim*3, 0, 0);
    Eigen::DSizes<ptrdiff_t, 3> indices_mat_i_inp(0, 0, 0);
    Eigen::DSizes<ptrdiff_t, 3> indices_mat_f_inp(input_dim*1, 0, 0);
    Eigen::DSizes<ptrdiff_t, 3> indices_mat_o_inp(input_dim*2, 0, 0);
    Eigen::DSizes<ptrdiff_t, 3> indices_mat_g_inp(input_dim*3, 0, 0);
    Eigen::DSizes<ptrdiff_t, 2> indices_i(0, 0);
    Eigen::DSizes<ptrdiff_t, 2> indices_f(hidden_dim,0);
    Eigen::DSizes<ptrdiff_t, 2> indices_o(hidden_dim*2,0);
    Eigen::DSizes<ptrdiff_t, 2> indices_g(hidden_dim*3,0);
    Eigen::DSizes<ptrdiff_t, 1> indices_i_nobatch(0);
    Eigen::DSizes<ptrdiff_t, 2> indices_mat_i_nobatch(0,0);
    Eigen::DSizes<ptrdiff_t, 1> indices_f_nobatch(hidden_dim);
    Eigen::DSizes<ptrdiff_t, 1> indices_o_nobatch(hidden_dim*2);
    Eigen::DSizes<ptrdiff_t, 1> indices_g_nobatch(hidden_dim*3);
    Eigen::DSizes<ptrdiff_t, 2> indices_mat_g_nobatch(hidden_dim*3, 0);
    Eigen::DSizes<ptrdiff_t, 3> sizes_mat_1(hidden_dim, 1, static_cast<ptrdiff_t>(fx.d.bd));
    Eigen::DSizes<ptrdiff_t, 3> sizes_mat_1_inp(input_dim, 1, static_cast<ptrdiff_t>(fx.d.bd));
    Eigen::DSizes<ptrdiff_t, 3> sizes_mat_3(hidden_dim*3, 1, static_cast<ptrdiff_t>(fx.d.bd));
    Eigen::array<ptrdiff_t, 1> vec_batch_axis = {1};

    AlignedMemoryPool* scratch_allocator = fx.device->pools[(int)DeviceMempool::SCS];

    // scratch memory to avoid striding for element-wise operations
    Tensor dEdf_ifo(Dim({hidden_dim*3, 1},batch_size), nullptr, fx.device, fx.mem_pool);
    dEdf_ifo.v = static_cast<float*>(scratch_allocator->allocate(dEdf_ifo.d.size() * sizeof(float)));
    Tensor dEdf_g(Dim({hidden_dim, 1},batch_size), nullptr, fx.device, fx.mem_pool);
    dEdf_g.v = static_cast<float*>(scratch_allocator->allocate(dEdf_g.d.size() * sizeof(float)));
    Tensor fx_ifo(Dim({hidden_dim*3, 1},batch_size), nullptr, fx.device, fx.mem_pool);
    fx_ifo.v = static_cast<float*>(scratch_allocator->allocate(fx_ifo.d.size() * sizeof(float)));
    Tensor fx_g(Dim({hidden_dim, 1},batch_size), nullptr, fx.device, fx.mem_pool);
    fx_g.v = static_cast<float*>(scratch_allocator->allocate(fx_g.d.size() * sizeof(float)));
    tb<2>(dEdf_ifo).device(*dev.edevice) = tb<2>(dEdf).slice(indices_mat_i, sizes_mat_3);
    tb<2>(dEdf_g).device(*dev.edevice)   = tb<2>(dEdf).slice(indices_mat_g, sizes_mat_1);
    tb<2>(fx_ifo).device(*dev.edevice) = tb<2>(fx).slice(indices_mat_i, sizes_mat_3);
    tb<2>(fx_g).device(*dev.edevice)   = tb<2>(fx).slice(indices_mat_g, sizes_mat_1);

    Tensor mask_x(Dim({input_dim}, batch_size), nullptr, fx.device, fx.mem_pool);
    Tensor mask_h(Dim({hidden_dim}, batch_size), nullptr, fx.device, fx.mem_pool);
    if(i < num_inputs+3 && dropout){
      if(xs[num_inputs+4]->d.bd == batch_size){
        mask_x.v = xs[num_inputs+4]->v;
      } else {
        mask_x.v = static_cast<float*>(scratch_allocator->allocate(mask_x.d.size() * sizeof(float)));
#ifdef __CUDACC__
        Eigen::array<ptrdiff_t, 2> bcast = {1, batch_size};
        tbvec(mask_x).device(*dev.edevice) = tbvec(*xs[num_inputs+4]).broadcast(bcast);
#else
        float *curr_ptr = mask_x.v, *end_ptr = curr_ptr + mask_x.d.size(), *in_ptr = xs[num_inputs+4]->v;
        do {
          memcpy(curr_ptr, in_ptr, sizeof(float)*xs[num_inputs+4]->d[0]);
          curr_ptr += xs[num_inputs+4]->d[0];
        } while(curr_ptr != end_ptr);
#endif
      }
      if(xs[num_inputs+5]->d.bd == batch_size){
        mask_h.v = xs[num_inputs+5]->v;
      } else {
        mask_h.v = static_cast<float*>(scratch_allocator->allocate(mask_h.d.size() * sizeof(float)));
#ifdef __CUDACC__
        Eigen::array<ptrdiff_t, 2> bcast = {1, batch_size};
        tbvec(mask_h).device(*dev.edevice) = tbvec(*xs[num_inputs+5]).broadcast(bcast);
#else
        float *curr_ptr = mask_h.v, *end_ptr = curr_ptr + mask_h.d.size(), *in_ptr = xs[num_inputs+5]->v;
        do {
          memcpy(curr_ptr, in_ptr, sizeof(float)*xs[num_inputs+5]->d[0]);
          curr_ptr += xs[num_inputs+5]->d[0];
        } while(curr_ptr != end_ptr);
#endif
      }
    }


    if(i < num_inputs){
      // dx_t = [Wx_i]^T   [di . i_t . (1-i_t)]
      //        [Wx_f]   * [df . f_t . (1-f_t)]
      //        [Wx_o]     [do . o_t . (1-o_t)]
      //        [Wx_g]     [dg . (1 - g_t^2)]
      //    where Wx is broadcasted over batches

      // when using multiple inputs: slice out the one we're backpropagating to
      Tensor Wx_slice(Dim({hidden_dim*4, xs[i]->d[0]},1), nullptr, fx.device, fx.mem_pool);
      if(num_inputs==1){
        Wx_slice.v = Wx->v;
      } else {
        unsigned offset=0;
        for(unsigned j=0; j<i; j++) offset += xs[j]->d[0];
        Eigen::DSizes<ptrdiff_t, 3> indices_Wx(0, offset, 0);
        Eigen::DSizes<ptrdiff_t, 3> sizes_Wx(hidden_dim*4, xs[i]->d[0], 1);
        Wx_slice.v = static_cast<float*>(scratch_allocator->allocate(Wx_slice.d.size() * sizeof(float)));
        tb<2>(Wx_slice).device(*dev.edevice) = tb<2>(*Wx).slice(indices_Wx, sizes_Wx);
      }

      // scratch memory for the matrix multiplication
      Tensor mult_r_ifo(Dim({hidden_dim*3, 1},batch_size), nullptr, fx.device, fx.mem_pool);
      mult_r_ifo.v = static_cast<float*>(scratch_allocator->allocate(mult_r_ifo.d.size() * sizeof(float)));
      Tensor mult_r_g(Dim({hidden_dim, 1},batch_size), nullptr, fx.device, fx.mem_pool);
      mult_r_g.v = static_cast<float*>(scratch_allocator->allocate(mult_r_g.d.size() * sizeof(float)));
      Tensor mult_r(Dim({hidden_dim*4, 1},batch_size), nullptr, fx.device, fx.mem_pool);
      mult_r.v = static_cast<float*>(scratch_allocator->allocate(mult_r.d.size() * sizeof(float)));

      // mult_r = [di . i_t . (1-i_t)]
      //          [df . f_t . (1-f_t)]
      //          [do . o_t . (1-o_t)]
      //          [dg . (1 - g_t^2)]
      tb<2>(mult_r_ifo).device(*dev.edevice) = tb<2>(dEdf_ifo) * tb<2>(fx_ifo) * (tb<2>(fx_ifo).constant(1) - tb<2>(fx_ifo));
      tb<2>(mult_r_g).device(*dev.edevice) = tb<2>(dEdf_g) * (tb<2>(fx_g).constant(1) - tb<2>(fx_g).square());
      tb<2>(mult_r).slice(indices_mat_i, sizes_mat_3).device(*dev.edevice) = tb<2>(mult_r_ifo);
      tb<2>(mult_r).slice(indices_mat_g, sizes_mat_1).device(*dev.edevice) = tb<2>(mult_r_g);
      // dx_t += mult_l^T * mult_r
      if(dropout){
        Tensor mult_y(Dim({xs[i]->d[0], 1},batch_size), nullptr, fx.device, fx.mem_pool);
        mult_y.v = static_cast<float*>(scratch_allocator->allocate(mult_y.d.size() * sizeof(float)));
        TensorTools::zero(mult_y);
        MatrixTranspMultiplyAcc(dev, Wx_slice, mult_r, mult_y);
        // when using multiple inputs: slice out the appropriate dropout mask
        Tensor dropout_mask(Dim({xs[i]->d[0]}, mask_x.d.bd), nullptr, fx.device, fx.mem_pool);
        if(num_inputs==1){
          dropout_mask.v = mask_x.v;
        } else {
          unsigned offset=0;
          for(unsigned j=0; j<i; j++) offset += xs[j]->d[0];
          Eigen::DSizes<ptrdiff_t, 2> indices_dropout(offset, 0);
          Eigen::DSizes<ptrdiff_t, 2> sizes_dropout(xs[i]->d[0], mask_x.d.bd);
          dropout_mask.v = static_cast<float*>(scratch_allocator->allocate(dropout_mask.d.size() * sizeof(float)));
          tbvec(dropout_mask).device(*dev.edevice) = tbvec(mask_x).slice(indices_dropout, sizes_dropout);
        }
        tvec(dEdxi).device(*dev.edevice) += tvec(mult_y) * tvec(dropout_mask);
      } else {
		    MatrixTranspMultiplyAcc(dev, Wx_slice, mult_r, dEdxi);
      }
    } else if(i==num_inputs){
      // dh_tm1 = [Wh_i]^T   [di . i_t . (1-i_t)]
      //          [Wh_f]   * [df . f_t . (1-f_t)]
      //          [Wh_o]     [do . o_t . (1-o_t)]
      //          [Wh_g]     [dg . (1 - g_t^2)]
      //    where Wh is broadcasted over batches

      // scratch memory for the matrix multiplication
      Tensor mult_r_ifo(Dim({hidden_dim*3, 1},batch_size), nullptr, fx.device, fx.mem_pool);
      mult_r_ifo.v = static_cast<float*>(scratch_allocator->allocate(mult_r_ifo.d.size() * sizeof(float)));
      Tensor mult_r_g(Dim({hidden_dim, 1},batch_size), nullptr, fx.device, fx.mem_pool);
      mult_r_g.v = static_cast<float*>(scratch_allocator->allocate(mult_r_g.d.size() * sizeof(float)));
      Tensor mult_r(Dim({hidden_dim*4, 1},batch_size), nullptr, fx.device, fx.mem_pool);
      mult_r.v = static_cast<float*>(scratch_allocator->allocate(mult_r.d.size() * sizeof(float)));

      // mult_r = [di . i_t . (1-i_t)]
      //          [df . f_t . (1-f_t)]
      //          [do . o_t . (1-o_t)]
      //          [dg . (1 - g_t^2)]
      tb<2>(mult_r_ifo).device(*dev.edevice) = tb<2>(dEdf_ifo) * tb<2>(fx_ifo) * (tb<2>(fx_ifo).constant(1) - tb<2>(fx_ifo));
      tb<2>(mult_r_g).device(*dev.edevice) = tb<2>(dEdf_g) * (tb<2>(fx_g).constant(1) - tb<2>(fx_g).square());
      tb<2>(mult_r).slice(indices_mat_i, sizes_mat_3).device(*dev.edevice) = tb<2>(mult_r_ifo);
      tb<2>(mult_r).slice(indices_mat_g, sizes_mat_1).device(*dev.edevice) = tb<2>(mult_r_g);

      // dx_t += mult_l * mult_r
      if(dropout){
        Tensor mult_y(Dim({hidden_dim, 1},batch_size), nullptr, fx.device, fx.mem_pool);
        mult_y.v = static_cast<float*>(scratch_allocator->allocate(mult_y.d.size() * sizeof(float)));
        TensorTools::zero(mult_y);
        MatrixTranspMultiplyAcc(dev, *Wh, mult_r, mult_y);
        tvec(dEdxi).device(*dev.edevice) += tvec(mult_y) * tvec(mask_h);
      } else {
        MatrixTranspMultiplyAcc(dev, *Wh, mult_r, dEdxi);
      }

    } else if(i==num_inputs+1){ // dWx
      // goal: dWx_i = [di . i_t . (1-i_t)] * x_t (here * is outer product), then sum over batches
      //       dWx_f = [di . f_t . (1-f_t)] * x_t (here * is outer product), then sum over batches
      //       dWx_o = [di . o_t . (1-o_t)] * x_t (here * is outer product), then sum over batches
      //       dWx_g = [dg . (1 - g_t^2)] * x_t (here * is outer product), then sum over batches

      // scratch memory for the matrix multiplication
      Tensor mult_l_ifo(Dim({hidden_dim*3, 1},batch_size), nullptr, fx.device, fx.mem_pool);
      mult_l_ifo.v = static_cast<float*>(scratch_allocator->allocate(mult_l_ifo.d.size() * sizeof(float)));
      Tensor mult_l_g(Dim({hidden_dim, 1},batch_size), nullptr, fx.device, fx.mem_pool);
      mult_l_g.v = static_cast<float*>(scratch_allocator->allocate(mult_l_g.d.size() * sizeof(float)));
      Tensor mult_l(Dim({hidden_dim*4, 1},batch_size), nullptr, fx.device, fx.mem_pool);
      mult_l.v = static_cast<float*>(scratch_allocator->allocate(mult_l.d.size() * sizeof(float)));

      Tensor x_t(Dim({input_dim}, batch_size), nullptr, fx.device, fx.mem_pool);
      if(num_inputs==1){
        x_t.v = xs[0]->v;
      } else {
        Tensor tmp(Dim({input_dim}, batch_size), nullptr, fx.device, fx.mem_pool);
        tmp.v = static_cast<float*>(scratch_allocator->allocate(tmp.d.size() * sizeof(float)));
        Eigen::DSizes<ptrdiff_t, 2> indices_tmp(0, 0);
        Eigen::DSizes<ptrdiff_t, 2> sizes_tmp(0, static_cast<ptrdiff_t>(fx.d.bd));
        for(unsigned i=0; i<num_inputs; i++){
          sizes_tmp[0] = xs[i]->d[0];
          tbvec(tmp).slice(indices_tmp, sizes_tmp).device(*dev.edevice) = tbvec(*xs[i]);
          indices_tmp[0] += xs[i]->d[0];
        }
        x_t.v = tmp.v;
      }


      // mult_l = [di . i_t . (1-i_t)]
      //          [df . f_t . (1-f_t)]
      //          [do . o_t . (1-o_t)]
      //          [dg . (1 - g_t^2)]
      tb<2>(mult_l_ifo).device(*dev.edevice) = tb<2>(dEdf_ifo) * tb<2>(fx_ifo) * (tb<2>(fx_ifo).constant(1) - tb<2>(fx_ifo));
      tb<2>(mult_l_g).device(*dev.edevice) = tb<2>(dEdf_g) * (tb<2>(fx_g).constant(1) - tb<2>(fx_g).square());
      tb<2>(mult_l).slice(indices_mat_i, sizes_mat_3).device(*dev.edevice) = tb<2>(mult_l_ifo);
      tb<2>(mult_l).slice(indices_mat_g, sizes_mat_1).device(*dev.edevice) = tb<2>(mult_l_g);

      if(dropout){
        Tensor x_t_dropped(Dim({input_dim}, batch_size), nullptr, fx.device, fx.mem_pool);
        x_t_dropped.v = static_cast<float*>(scratch_allocator->allocate(x_t_dropped.d.size() * sizeof(float)));
        tvec(x_t_dropped).device(*dev.edevice) = tvec(x_t) * tvec(mask_x);
        x_t.v = x_t_dropped.v;
      }

      // dWh += (mult_l * mult_r).sum_batches()
      MatrixMultiplyTranspAcc(dev, mult_l, x_t, dEdxi);

    } else if(i==num_inputs+2){ // dWh
      // goal: dWh_i = [di . i_t . (1-i_t)] * h_tm1 (here * is outer product), then sum over batches
      //       dWh_f = [df . f_t . (1-f_t)] * h_tm1 (here * is outer product), then sum over batches
      //       dWh_o = [do . o_t . (1-o_t)] * h_tm1 (here * is outer product), then sum over batches
      //       dWh_g = [dg . (1 - g_t^2)] * h_tm1 (here * is outer product), then sum over batches

      // scratch memory for the matrix multiplication
      Tensor mult_l_ifo(Dim({hidden_dim*3, 1},batch_size), nullptr, fx.device, fx.mem_pool);
      mult_l_ifo.v = static_cast<float*>(scratch_allocator->allocate(mult_l_ifo.d.size() * sizeof(float)));
      Tensor mult_l_g(Dim({hidden_dim, 1},batch_size), nullptr, fx.device, fx.mem_pool);
      mult_l_g.v = static_cast<float*>(scratch_allocator->allocate(mult_l_g.d.size() * sizeof(float)));
      Tensor mult_l(Dim({hidden_dim*4, 1},batch_size), nullptr, fx.device, fx.mem_pool);
      mult_l.v = static_cast<float*>(scratch_allocator->allocate(mult_l.d.size() * sizeof(float)));

      // mult_l = [di . i_t . (1-i_t)]
      //          [df . f_t . (1-f_t)]
      //          [do . o_t . (1-o_t)]
      //          [dg . (1 - g_t^2)]
      tb<2>(mult_l_ifo).device(*dev.edevice) = tb<2>(dEdf_ifo) * tb<2>(fx_ifo) * (tb<2>(fx_ifo).constant(1) - tb<2>(fx_ifo));
      tb<2>(mult_l_g).device(*dev.edevice) = tb<2>(dEdf_g) * (tb<2>(fx_g).constant(1) - tb<2>(fx_g).square());
      tb<2>(mult_l).slice(indices_mat_i, sizes_mat_3).device(*dev.edevice) = tb<2>(mult_l_ifo);
      tb<2>(mult_l).slice(indices_mat_g, sizes_mat_1).device(*dev.edevice) = tb<2>(mult_l_g);

      if(dropout){
        Tensor h_tm1_dropped(Dim({hidden_dim}, batch_size), nullptr, fx.device, fx.mem_pool);
        h_tm1_dropped.v = static_cast<float*>(scratch_allocator->allocate(h_tm1_dropped.d.size() * sizeof(float)));
        tvec(h_tm1_dropped).device(*dev.edevice) = tvec(h_tm1) * tvec(mask_h);
        h_tm1.v = h_tm1_dropped.v;
      }

      // dWh += (mult_l * mult_r).sum(batches)
      MatrixMultiplyTranspAcc(dev, mult_l, h_tm1, dEdxi);

    } else if(i==num_inputs+3){
      Eigen::DSizes<ptrdiff_t, 1> sizes_1_nobatch(hidden_dim);
      Eigen::DSizes<ptrdiff_t, 1> sizes_3_nobatch(hidden_dim*3);
      Eigen::DSizes<ptrdiff_t, 2> sizes_1(hidden_dim, static_cast<ptrdiff_t>(fx.d.bd));
      Eigen::DSizes<ptrdiff_t, 2> sizes_3(hidden_dim*3, static_cast<ptrdiff_t>(fx.d.bd));
      Tensor dEdxi_ifo(Dim({hidden_dim*3},1), nullptr, fx.device, fx.mem_pool);
      dEdxi_ifo.v = static_cast<float*>(scratch_allocator->allocate(dEdxi_ifo.d.size() * sizeof(float)));
      Tensor dEdxi_g(Dim({hidden_dim},1), nullptr, fx.device, fx.mem_pool);
      dEdxi_g.v = static_cast<float*>(scratch_allocator->allocate(dEdxi_g.d.size() * sizeof(float)));

      // db_i = di . i_t . (1-i_t), then sum over batches
      // db_f = df . f_t . (1-f_t), then sum over batches
      // db_o = do . f_t . (1-o_t), then sum over batches
      tvec(dEdxi_ifo).device(*dev.edevice) = (tbvec(dEdf_ifo) * tbvec(fx_ifo) * (tbvec(fx_ifo).constant(1) - tbvec(fx_ifo))).sum(vec_batch_axis);
      tvec(dEdxi).slice(indices_i_nobatch, sizes_3_nobatch).device(*dev.edevice) += tvec(dEdxi_ifo);

      // db_g = dg . (1 - g_t^2), then sum over batches
      tvec(dEdxi_g).device(*dev.edevice) = (tbvec(dEdf_g) * (tbvec(fx_g).constant(1) - tbvec(fx_g).square())).sum(vec_batch_axis);
      tvec(dEdxi).slice(indices_g_nobatch, sizes_1_nobatch).device(*dev.edevice) += tvec(dEdxi_g);
    }
    // no gradients for dropout masks computed

    scratch_allocator->free();


  }
  DYNET_NODE_INST_DEV_IMPL(VanillaLSTMGates)


  // ************* LSTM Cell *************

#ifndef __CUDACC__

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

  Dim VanillaLSTMC::dim_forward(const vector<Dim>& xs) const {
    DYNET_ARG_CHECK(xs.size() == 2, "Failed input count check in VanillaLSTMC");
    DYNET_ARG_CHECK(xs[0].ndims() == 1, "VanillaLSTMC: c_tm1 expected to be a vector");
    DYNET_ARG_CHECK(xs[1].ndims() == 1, "VanillaLSTMC: gates_t expected to be a vector");
    DYNET_ARG_CHECK(xs[0].size()*4 == xs[1].size(), "VanillaLSTMC: gates_t expected 4 times as big as c_t, but " << xs[0].size() << "*4 != " << xs[1].size());
    DYNET_ARG_CHECK(xs[0].bd == xs[1].bd, "VanillaLSTMC: gates_t and c_t expected to have equal batch size, but " << xs[0].bd << " != " << xs[1].bd);
    return xs[0];
  }

  int VanillaLSTMC::autobatch_sig(const ComputationGraph & cg, SigMap &sm) const {
    Sig s(nt::vanilla_lstm_h);
    s.add_dim(cg.nodes[args[0]]->dim);
    return sm.get_idx(s);
  }

  std::vector<int> VanillaLSTMC::autobatch_concat(const ComputationGraph & cg) const {
    return vector<int>(2, 1);
  }

#endif

  template<class MyDevice>
  void VanillaLSTMC::forward_dev_impl(const MyDevice & dev, const vector<const Tensor*>& xs, Tensor& fx) const {
    // computes cell state (elementwise multiplication)
    // c_t = i_t . g_t + f_t . c_tm1

    DYNET_ASSERT(xs.size() == 2, "Failed dimension check in VanillaLSTMC::forward");

    const Tensor *c_tm1 = xs[0];
    const Tensor *gates_t = xs[1];

    unsigned hidden_dim = c_tm1->d[0];
    unsigned batch_size = c_tm1->d.bd;

    Eigen::DSizes<ptrdiff_t, 2> indices_i(0, 0);
    Eigen::DSizes<ptrdiff_t, 2> indices_f(hidden_dim,0);
    Eigen::DSizes<ptrdiff_t, 2> indices_g(hidden_dim*3,0);
    Eigen::DSizes<ptrdiff_t, 2> sizes_1(hidden_dim, static_cast<ptrdiff_t>(fx.d.bd));

    AlignedMemoryPool* scratch_allocator = fx.device->pools[(int)DeviceMempool::SCS];
    Tensor f_t(Dim({hidden_dim,1},batch_size), nullptr, fx.device, fx.mem_pool);
    f_t.v = static_cast<float*>(scratch_allocator->allocate(f_t.d.size() * sizeof(float)));
    tbvec(f_t).device(*dev.edevice) = tbvec(*gates_t).slice(indices_f, sizes_1);

    tbvec(fx).device(*dev.edevice) = tbvec(*gates_t).slice(indices_i, sizes_1);
    tbvec(fx).device(*dev.edevice) = tbvec(fx) * tbvec(*gates_t).slice(indices_g, sizes_1) + tbvec(f_t) * tbvec(*c_tm1);
    scratch_allocator->free();

  }

  template<class MyDevice>
  void VanillaLSTMC::backward_dev_impl(const MyDevice & dev,
                               const vector<const Tensor*>& xs,
                               const Tensor& fx,
                               const Tensor& dEdf,
                               unsigned i,
                               Tensor& dEdxi) const {
    unsigned hidden_dim = fx.d[0];

    Eigen::DSizes<ptrdiff_t, 2> indices_i(0,0);
    Eigen::DSizes<ptrdiff_t, 2> indices_f(hidden_dim,0);
    Eigen::DSizes<ptrdiff_t, 2> indices_g(hidden_dim*3,0);
    Eigen::DSizes<ptrdiff_t, 2> sizes_1(hidden_dim, static_cast<ptrdiff_t>(fx.d.bd));

    if(i==0){ // dc_tm1 = dc_t . f_t
      tbvec(dEdxi).device(*dev.edevice) += tbvec(dEdf) * tbvec(*xs[1]).slice(indices_f, sizes_1);
    } else if(i==1){
      // di_t = dc_t . g_t
      tbvec(dEdxi).slice(indices_i, sizes_1).device(*dev.edevice) += tbvec(dEdf) * tbvec(*xs[1]).slice(indices_g, sizes_1);
      // df_t = dc_t . c_tm1
      tbvec(dEdxi).slice(indices_f, sizes_1).device(*dev.edevice) += tbvec(dEdf) * tbvec(*xs[0]);
      // dg_t = dc_t . i_t
      tbvec(dEdxi).slice(indices_g, sizes_1).device(*dev.edevice) += tbvec(dEdf) * tbvec(*xs[1]).slice(indices_i, sizes_1);
    }
  }

  DYNET_NODE_INST_DEV_IMPL(VanillaLSTMC)

  // ************* LSTM State *************

#ifndef __CUDACC__

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

  Dim VanillaLSTMH::dim_forward(const vector<Dim>& xs) const {
    DYNET_ARG_CHECK(xs.size() == 2, "Failed input count check in VanillaLSTMH");
    DYNET_ARG_CHECK(xs[0].ndims() == 1, "VanillaLSTMH: c_t expected to be a vector");
    DYNET_ARG_CHECK(xs[1].ndims() == 1, "VanillaLSTMH: gates_t expected to be a vector");
    DYNET_ARG_CHECK(xs[0].size()*4 == xs[1].size(), "VanillaLSTMH: gates_t expected 4 times as big as c_t, but " << xs[0].size() << "*4 != " << xs[1].size());
    DYNET_ARG_CHECK(xs[0].bd == xs[1].bd, "VanillaLSTMH: gates_t and c_t expected to have equal batch size, but " << xs[0].bd << " != " << xs[1].bd);
    return xs[0];
  }

  int VanillaLSTMH::autobatch_sig(const ComputationGraph & cg, SigMap &sm) const {
    Sig s(nt::vanilla_lstm_h);
    s.add_dim(cg.nodes[args[0]]->dim);
    return sm.get_idx(s);
  }

  std::vector<int> VanillaLSTMH::autobatch_concat(const ComputationGraph & cg) const {
    return vector<int>(2, 1);
  }

#endif

  template<class MyDevice>
  void VanillaLSTMH::forward_dev_impl(const MyDevice & dev, const vector<const Tensor*>& xs, Tensor& fx) const {
    // computes output state (elementwise multiplication)
    // h_t = o_t . tanh(c_t)

    DYNET_ASSERT(xs.size() == 2, "Failed dimension check in VanillaLSTMH::forward");

    const Tensor *c_t = xs[0];
    const Tensor *gates_t = xs[1];

    unsigned hidden_dim = c_t->d[0];
    unsigned batch_size = c_t->d.bd;

    Eigen::DSizes<ptrdiff_t, 3> indices_o(hidden_dim*2,0,0);
    Eigen::DSizes<ptrdiff_t, 3> sizes_1(hidden_dim, 1, static_cast<ptrdiff_t>(batch_size));

    tb<2>(fx).device(*dev.edevice) = tb<2>(*gates_t).slice(indices_o, sizes_1);
    tb<2>(fx).device(*dev.edevice) = tb<2>(fx) * tb<2>(*c_t).tanh();
  }

  template<class MyDevice>
  void VanillaLSTMH::backward_dev_impl(const MyDevice & dev,
                               const vector<const Tensor*>& xs,
                               const Tensor& fx,
                               const Tensor& dEdf,
                               unsigned i,
                               Tensor& dEdxi) const {
    unsigned hidden_dim = fx.d[0];
    unsigned batch_size = fx.d.bd;
    Eigen::DSizes<ptrdiff_t, 3> indices_o(hidden_dim*2,0,0);
    Eigen::DSizes<ptrdiff_t, 3> sizes_1(hidden_dim, 1, static_cast<ptrdiff_t>(batch_size));

    AlignedMemoryPool* scratch_allocator = fx.device->pools[(int)DeviceMempool::SCS];

    if(i==0){
      // dc_t = dh_t . o_t . (1 - tanh^2(c_t)))
      //      = dh_t . o_t . (1 - (h_t cdiv o_t)^2)

      Tensor o_t(Dim({hidden_dim,1},batch_size), nullptr, fx.device, fx.mem_pool);
      o_t.v = static_cast<float*>(scratch_allocator->allocate(o_t.d.size() * sizeof(float)));

      tb<2>(o_t).device(*dev.edevice) = tb<2>(*xs[1]).slice(indices_o, sizes_1);
      tb<2>(dEdxi).device(*dev.edevice) += tb<2>(dEdf)
                                            * tb<2>(o_t)
                                            * (tb<2>(*xs[0]).constant(1) - tb<2>(*xs[0]).tanh().square());
    } else if(i==1){
      Tensor dEdxi_o(Dim({hidden_dim,1},batch_size), nullptr, fx.device, fx.mem_pool);
      dEdxi_o.v = static_cast<float*>(scratch_allocator->allocate(dEdxi_o.d.size() * sizeof(float)));
      // do_t = dh_t . tanh(c_t)
      tb<2>(dEdxi_o).device(*dev.edevice) = tb<2>(dEdf) * tb<2>(*xs[0]).tanh();
      tb<2>(dEdxi).slice(indices_o, sizes_1).device(*dev.edevice) += tb<2>(dEdxi_o);
    }
    scratch_allocator->free();
  }

  DYNET_NODE_INST_DEV_IMPL(VanillaLSTMH)


}