void Compute()

in src/cwise_linear_op.cc [133:189]


  void Compute(OpKernelContext* ctx) override {

    const Tensor& dy = ctx->input(0);
    OpInputList xy; ctx->input_list("xy", &xy);
    OpInputList a;  ctx->input_list("a",  &a);
    OpInputList b;  ctx->input_list("b",  &b);

    uint N = dy.dim_size(0);
    uint C = dy.dim_size(1);

    uint rank = dy.dims();
    uint DHW  = 1;
    for (uint r = 2; r < rank; r++)
      DHW *= dy.dim_size(r);

    V* dx_ptr = NULL;
    if (a.size() == 0 && !relu_)
      // no scale and no relu: just pass dy to dx
      ctx->set_output(0, dy);
    else
    {
      Tensor* dx; OP_REQUIRES_OK(ctx, ctx->allocate_output(0, dy.shape(), &dx));
      dx_ptr = (V*)dx->flat<T>().data();
    }

    float* da_ptr;
    if (a.size())
    {
      Tensor* da; OP_REQUIRES_OK(ctx, ctx->allocate_output(1, a[0].shape(), &da));
      da_ptr = da->flat<float>().data();
    }
    else
    {
      Tensor* da; OP_REQUIRES_OK(ctx, ctx->allocate_output(1, TensorShape(), &da));
      da_ptr = NULL;
    }
    float* db_ptr;
    if (b.size())
    {
      Tensor* db; OP_REQUIRES_OK(ctx, ctx->allocate_output(2, b[0].shape(), &db));
      db_ptr = db->flat<float>().data();
    }
    else
    {
      Tensor* db; OP_REQUIRES_OK(ctx, ctx->allocate_output(2, TensorShape(), &db));
      db_ptr = NULL;
    }

    const     V* dy_ptr =             (const V*)dy.flat<T>().data();
    const     V* xy_ptr = xy.size() ? (const V*)xy[0].flat<T>().data() : NULL;
    const float*  a_ptr = a.size()  ? a[0].flat<float>().data()        : NULL;
    const float*  b_ptr = b.size()  ? b[0].flat<float>().data()        : NULL;

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

    CWiseLinear_Backward<V>(stream, dx_ptr, da_ptr, db_ptr, dy_ptr, xy_ptr, a_ptr, b_ptr, N, C, DHW, relu_, swap_);
  }