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_);
}