void Compute()

in src/lstm_op.cc [264:305]
36 lines of code
2 McCabe index (conditional complexity)

  void Compute(OpKernelContext* ctx) override {

    const Tensor& c  = ctx->input(0);
    const Tensor& i  = ctx->input(1);
    const Tensor& u  = ctx->input(2);
    const Tensor& f  = ctx->input(3);
    const Tensor& o  = ctx->input(4);
    const Tensor& eh = ctx->input(5);

    int N = c.dim_size(0);
    int K = c.dim_size(1);

    Tensor* dc  = nullptr;
    Tensor* di  = nullptr;
    Tensor* df  = nullptr;
    Tensor* doo = nullptr;
    Tensor* du  = nullptr;
    OP_REQUIRES_OK(ctx, ctx->allocate_output(0, c.shape(), &dc));
    OP_REQUIRES_OK(ctx, ctx->allocate_output(1, c.shape(), &di));
    OP_REQUIRES_OK(ctx, ctx->allocate_output(2, c.shape(), &du));
    OP_REQUIRES_OK(ctx, ctx->allocate_output(3, c.shape(), &df));
    OP_REQUIRES_OK(ctx, ctx->allocate_output(4, c.shape(), &doo));

    VB1* dc_ptr = (VB1*)dc->flat<B>().data();
    VB1* di_ptr = (VB1*)di->flat<B>().data();
    VB1* df_ptr = (VB1*)df->flat<B>().data();
    VB1* do_ptr = (VB1*)doo->flat<B>().data();
    VB1* du_ptr = (VB1*)du->flat<B>().data();
    const VF1*  c_ptr = (const VF1*)c.flat<F>().data();
    const VF1*  i_ptr = (const VF1*)i.flat<F>().data();
    const VF1*  f_ptr = (const VF1*)f.flat<F>().data();
    const VF1*  o_ptr = (const VF1*)o.flat<F>().data();
    const VF1*  u_ptr = (const VF1*)u.flat<F>().data();
    const VB1* eh_ptr = (const VB1*)eh.flat<B>().data();
    const VB1* ec_ptr = nullptr;
    if (ctx->num_inputs() == 7)
      ec_ptr = (const VB1*)ctx->input(6).flat<B>().data();

    CUstream stream = ((CUDAStream*)ctx->op_device_context()->stream()->implementation())->cuda_stream();

    LSTM4_Gates_Backward<VB1,VF1,VB4,VF4>(stream, dc_ptr, di_ptr, df_ptr, do_ptr, du_ptr, ec_ptr, eh_ptr, c_ptr, i_ptr, f_ptr, o_ptr, u_ptr, N, K, forget_bias_);
  }