Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Removed NDIMS template arg from BinaryElementWiseOp::Operate #36526

Merged
merged 2 commits into from
Feb 27, 2020
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
24 changes: 1 addition & 23 deletions tensorflow/core/framework/numeric_op.h
Original file line number Diff line number Diff line change
Expand Up @@ -82,29 +82,7 @@ class BinaryElementWiseOp : public BinaryOp<T> {
{0, 1}, 0, a.shape(), &output));

// Dispatch to the descendant's Operate() function.
switch (a.dims()) {
#define NDIM_CASE(NDIMS) \
case NDIMS: { \
static_cast<CHILD*>(this)->template Operate<NDIMS>(context, a, b, output); \
break; \
}

NDIM_CASE(0);
NDIM_CASE(1);
NDIM_CASE(2);
NDIM_CASE(3);
NDIM_CASE(4);
NDIM_CASE(5);
NDIM_CASE(6);
NDIM_CASE(7);
NDIM_CASE(8);
#undef NDIM_CASE

default:
context->SetStatus(errors::InvalidArgument(
"We only handle up to Tensor::dims() up to 8, not ", a.dims()));
break;
}
static_cast<CHILD*>(this)->Operate(context, a, b, output);
}
};

Expand Down
6 changes: 0 additions & 6 deletions tensorflow/core/kernels/fake_quant_ops.cc
Original file line number Diff line number Diff line change
Expand Up @@ -124,14 +124,8 @@ class FakeQuantWithMinMaxArgsGradientOp
quant_max_ = (1 << num_bits) - 1;
}

template <int NDIMS>
void Operate(OpKernelContext* context, const Tensor& gradient,
const Tensor& input, Tensor* output) {
OperateNoTemplate(context, gradient, input, output);
}

void OperateNoTemplate(OpKernelContext* context, const Tensor& gradient,
const Tensor& input, Tensor* output) {
OP_REQUIRES(context, input.IsSameSize(gradient),
InvalidArgument("gradient and input must be the same size"));
FakeQuantWithMinMaxArgsGradientFunctor<Device> functor;
Expand Down
96 changes: 20 additions & 76 deletions tensorflow/core/kernels/relu_op.h
Original file line number Diff line number Diff line change
Expand Up @@ -63,32 +63,21 @@ class ReluGradOp : public BinaryElementWiseOp<T, ReluGradOp<Device, T>> {
public:
using BinaryElementWiseOp<T, ReluGradOp<Device, T>>::BinaryElementWiseOp;

void OperateNoTemplate(OpKernelContext* context, const Tensor& g,
const Tensor& a, Tensor* output);

// INPUTS:
// g (gradients): backpropagated gradients
// a (inputs): either the inputs that were passed to ReluOp(), or its
// outputs (using either one yields the same result here).
// OUTPUT:
// gradients to backprop
template <int NDIMS>
void Operate(OpKernelContext* context, const Tensor& g, const Tensor& a,
Tensor* output) {
OperateNoTemplate(context, g, a, output);
if (!ReluHelpers::ValidateSameSize(context, g, a)) return;
functor::ReluGrad<Device, T> functor;
functor(context->eigen_device<Device>(), g.flat<T>(), a.flat<T>(),
output->flat<T>());
}
};

template <typename Device, typename T>
void ReluGradOp<Device, T>::OperateNoTemplate(OpKernelContext* context,
const Tensor& g, const Tensor& a,
Tensor* output) {
if (!ReluHelpers::ValidateSameSize(context, g, a)) return;
functor::ReluGrad<Device, T> functor;
functor(context->eigen_device<Device>(), g.flat<T>(), a.flat<T>(),
output->flat<T>());
}

template <typename Device, typename T>
class Relu6Op : public UnaryElementWiseOp<T, Relu6Op<Device, T>> {
public:
Expand All @@ -106,31 +95,20 @@ class Relu6GradOp : public BinaryElementWiseOp<T, Relu6GradOp<Device, T>> {
public:
using BinaryElementWiseOp<T, Relu6GradOp<Device, T>>::BinaryElementWiseOp;

void OperateNoTemplate(OpKernelContext* context, const Tensor& g,
const Tensor& a, Tensor* output);

// INPUTS:
// g (gradients): backpropagated gradients
// a (inputs): inputs that were passed to Relu6Op()
// OUTPUT:
// gradients to backprop
template <int NDIMS>
void Operate(OpKernelContext* context, const Tensor& g, const Tensor& a,
Tensor* output) {
OperateNoTemplate(context, g, a, output);
if (!ReluHelpers::ValidateSameSize(context, g, a)) return;
functor::Relu6Grad<Device, T> functor;
functor(context->eigen_device<Device>(), g.flat<T>(), a.flat<T>(),
output->flat<T>());
}
};

template <typename Device, typename T>
void Relu6GradOp<Device, T>::OperateNoTemplate(OpKernelContext* context,
const Tensor& g, const Tensor& a,
Tensor* output) {
if (!ReluHelpers::ValidateSameSize(context, g, a)) return;
functor::Relu6Grad<Device, T> functor;
functor(context->eigen_device<Device>(), g.flat<T>(), a.flat<T>(),
output->flat<T>());
}

template <typename Device, typename T>
class LeakyReluOp : public UnaryElementWiseOp<T, LeakyReluOp<Device, T>> {
public:
Expand Down Expand Up @@ -162,36 +140,24 @@ class LeakyReluGradOp
alpha_ = T(alpha_tmp);
}

void OperateNoTemplate(OpKernelContext* context, const Tensor& g,
const Tensor& a, T alpha, Tensor* output);

// INPUTS:
// g (gradients): backpropagated gradients
// a (inputs): either the inputs that were passed to LeakyReluOp(), or its
// outputs (using either one yields the same result here).
// OUTPUT:
// gradients to backprop
template <int NDIMS>
void Operate(OpKernelContext* context, const Tensor& g, const Tensor& a,
Tensor* output) {
OperateNoTemplate(context, g, a, alpha_, output);
if (!ReluHelpers::ValidateSameSize(context, g, a)) return;
functor::LeakyReluGrad<Device, T> functor;
functor(context->eigen_device<Device>(), g.flat<T>(), a.flat<T>(), alpha_,
output->flat<T>());
}

private:
T alpha_;
};

template <typename Device, typename T>
void LeakyReluGradOp<Device, T>::OperateNoTemplate(OpKernelContext* context,
const Tensor& g,
const Tensor& a, T alpha,
Tensor* output) {
if (!ReluHelpers::ValidateSameSize(context, g, a)) return;
functor::LeakyReluGrad<Device, T> functor;
functor(context->eigen_device<Device>(), g.flat<T>(), a.flat<T>(), alpha,
output->flat<T>());
};

template <typename Device, typename T>
class EluOp : public UnaryElementWiseOp<T, EluOp<Device, T>> {
public:
Expand All @@ -209,31 +175,20 @@ class EluGradOp : public BinaryElementWiseOp<T, EluGradOp<Device, T>> {
public:
using BinaryElementWiseOp<T, EluGradOp<Device, T>>::BinaryElementWiseOp;

void OperateNoTemplate(OpKernelContext* context, const Tensor& g,
const Tensor& a, Tensor* output);

// INPUTS:
// g (gradients): backpropagated gradients
// a (outputs): outputs of the EluOp()
// OUTPUT:
// gradients to backprop
template <int NDIMS>
void Operate(OpKernelContext* context, const Tensor& g, const Tensor& a,
Tensor* output) {
OperateNoTemplate(context, g, a, output);
if (!ReluHelpers::ValidateSameSize(context, g, a)) return;
functor::EluGrad<Device, T> functor;
functor(context->eigen_device<Device>(), g.flat<T>(), a.flat<T>(),
output->flat<T>());
}
};

template <typename Device, typename T>
void EluGradOp<Device, T>::OperateNoTemplate(OpKernelContext* context,
const Tensor& g, const Tensor& a,
Tensor* output) {
if (!ReluHelpers::ValidateSameSize(context, g, a)) return;
functor::EluGrad<Device, T> functor;
functor(context->eigen_device<Device>(), g.flat<T>(), a.flat<T>(),
output->flat<T>());
}

template <typename Device, typename T>
class SeluOp : public UnaryElementWiseOp<T, SeluOp<Device, T>> {
public:
Expand All @@ -251,31 +206,20 @@ class SeluGradOp : public BinaryElementWiseOp<T, SeluGradOp<Device, T>> {
public:
using BinaryElementWiseOp<T, SeluGradOp<Device, T>>::BinaryElementWiseOp;

void OperateNoTemplate(OpKernelContext* context, const Tensor& g,
const Tensor& a, Tensor* output);

// INPUTS:
// g (gradients): backpropagated gradients
// a (outputs): outputs of the SeluOp()
// OUTPUT:
// gradients to backprop
template <int NDIMS>
void Operate(OpKernelContext* context, const Tensor& g, const Tensor& a,
Tensor* output) {
OperateNoTemplate(context, g, a, output);
if (!ReluHelpers::ValidateSameSize(context, g, a)) return;
functor::SeluGrad<Device, T> functor;
functor(context->eigen_device<Device>(), g.flat<T>(), a.flat<T>(),
output->flat<T>());
}
};

template <typename Device, typename T>
void SeluGradOp<Device, T>::OperateNoTemplate(OpKernelContext* context,
const Tensor& g, const Tensor& a,
Tensor* output) {
if (!ReluHelpers::ValidateSameSize(context, g, a)) return;
functor::SeluGrad<Device, T> functor;
functor(context->eigen_device<Device>(), g.flat<T>(), a.flat<T>(),
output->flat<T>());
}

} // namespace tensorflow

#undef EIGEN_USE_THREADS
Expand Down
21 changes: 5 additions & 16 deletions tensorflow/core/kernels/softplus_op.cc
Original file line number Diff line number Diff line change
Expand Up @@ -50,31 +50,20 @@ class SoftplusGradOp
explicit SoftplusGradOp(OpKernelConstruction* context)
: BinaryElementWiseOp<T, SoftplusGradOp<Device, T>>(context) {}

void OperateNoTemplate(OpKernelContext* context, const Tensor& g,
const Tensor& a, Tensor* output);

// INPUTS:
// g (gradients): backpropagated gradients
// a (inputs): inputs that were passed to SoftplusOp()
// OUTPUT:
// gradients to backprop
template <int NDIMS>
void Operate(OpKernelContext* context, const Tensor& g, const Tensor& a,
Tensor* output) {
OperateNoTemplate(context, g, a, output);
OP_REQUIRES(context, a.IsSameSize(g),
errors::InvalidArgument("g and a must be the same size"));
functor::SoftplusGrad<Device, T> functor;
functor(context->eigen_device<Device>(), g.flat<T>(), a.flat<T>(),
output->flat<T>());
}
};
template <typename Device, typename T>
void SoftplusGradOp<Device, T>::OperateNoTemplate(OpKernelContext* context,
const Tensor& g,
const Tensor& a,
Tensor* output) {
OP_REQUIRES(context, a.IsSameSize(g),
errors::InvalidArgument("g and a must be the same size"));
functor::SoftplusGrad<Device, T> functor;
functor(context->eigen_device<Device>(), g.flat<T>(), a.flat<T>(),
output->flat<T>());
}

#define REGISTER_KERNELS(type) \
REGISTER_KERNEL_BUILDER( \
Expand Down
22 changes: 5 additions & 17 deletions tensorflow/core/kernels/softsign_op.cc
Original file line number Diff line number Diff line change
Expand Up @@ -50,33 +50,21 @@ class SoftsignGradOp
explicit SoftsignGradOp(OpKernelConstruction* context)
: BinaryElementWiseOp<T, SoftsignGradOp<Device, T>>(context) {}

void OperateNoTemplate(OpKernelContext* context, const Tensor& g,
const Tensor& a, Tensor* output);

// INPUTS:
// g (gradients): backpropagated gradients
// a (inputs): inputs that were passed to SoftsignOp()
// OUTPUT:
// gradients to backprop
template <int NDIMS>
void Operate(OpKernelContext* context, const Tensor& g, const Tensor& a,
Tensor* output) {
OperateNoTemplate(context, g, a, output);
OP_REQUIRES(context, a.IsSameSize(g),
errors::InvalidArgument("g and a must be the same size"));
functor::SoftsignGrad<Device, T> functor;
functor(context->eigen_device<Device>(), g.flat<T>(), a.flat<T>(),
output->flat<T>());
}
};

template <typename Device, typename T>
void SoftsignGradOp<Device, T>::OperateNoTemplate(OpKernelContext* context,
const Tensor& g,
const Tensor& a,
Tensor* output) {
OP_REQUIRES(context, a.IsSameSize(g),
errors::InvalidArgument("g and a must be the same size"));
functor::SoftsignGrad<Device, T> functor;
functor(context->eigen_device<Device>(), g.flat<T>(), a.flat<T>(),
output->flat<T>());
}

#define REGISTER_KERNELS(type) \
REGISTER_KERNEL_BUILDER( \
Name("Softsign").Device(DEVICE_CPU).TypeConstraint<type>("T"), \
Expand Down