experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Rename elem_bcast to then_else_bcast, and remove duplicate template specification

Browse Source 
based on review comment.

Signed-off-by: Yong Tang <yong.tang.github@outlook.com>
This commit is contained in:
Yong Tang 2019-05-02 21:43:46 +00:00
parent d9b98c4d16
commit e75409c2fe
2 changed files with 9 additions and 45 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

48
tensorflow/core/kernels/cwise_op_select.cc
View File

@ -162,15 +162,15 @@ class SelectV2Op : public OpKernel {
// 2-ary broadcast.
// Combine `then` and `else`.
BCast elem_bcast(BCast::FromShape(then->shape()),
BCast::FromShape(else_->shape()), false);
OP_REQUIRES(ctx, elem_bcast.IsValid(),
BCast then_else_bcast(BCast::FromShape(then->shape()),
BCast::FromShape(else_->shape()), false);
OP_REQUIRES(ctx, then_else_bcast.IsValid(),
errors::InvalidArgument(
"then ", then->shape().DebugString(), " and else ",
else_->shape().DebugString(), " must be broadcastable"));
// Combine `cond` with `then` and `else`.
BCast bcast(BCast::FromShape(cond->shape()),
BCast::FromShape(BCast::ToShape(elem_bcast.output_shape())),
BCast::FromShape(BCast::ToShape(then_else_bcast.output_shape())),
false);
OP_REQUIRES(ctx, bcast.IsValid(),
errors::InvalidArgument(
@ -443,46 +443,6 @@ struct BatchSelectFunctor<CPUDevice, T> {
d.parallelFor(batch, cost, work);
}
};
template <typename T, int NDIMS>
struct BCastSelectFunctor<CPUDevice, T, NDIMS> {
void operator()(const CPUDevice& d,
typename TTypes<T, NDIMS>::Tensor output_tensor,
typename TTypes<bool, NDIMS>::ConstTensor cond_tensor,
typename TTypes<T, NDIMS>::ConstTensor then_tensor,
typename TTypes<T, NDIMS>::ConstTensor else_tensor,
typename Eigen::array<Eigen::DenseIndex, NDIMS> cond_bcast,
typename Eigen::array<Eigen::DenseIndex, NDIMS> then_bcast,
typename Eigen::array<Eigen::DenseIndex, NDIMS> else_bcast) {
output_tensor.device(d) = cond_tensor.broadcast(cond_bcast)
.select(then_tensor.broadcast(then_bcast),
else_tensor.broadcast(else_bcast));
}
};
#ifdef TENSORFLOW_USE_SYCL
template <typename T>
struct BatchSelectFunctor<SYCLDevice, T>
: BatchSelectFunctorBase<SYCLDevice, T> {};
template <typename T, int NDIMS>
struct BCastSelectFunctor<SYCLDevice, T, NDIMS> {
void operator()(const SYCLDevice& d,
typename TTypes<T, NDIMS>::Tensor output_tensor,
typename TTypes<bool, NDIMS>::ConstTensor cond_tensor,
typename TTypes<T, NDIMS>::ConstTensor then_tensor,
typename TTypes<T, NDIMS>::ConstTensor else_tensor,
typename Eigen::array<Eigen::DenseIndex, NDIMS> cond_bcast,
typename Eigen::array<Eigen::DenseIndex, NDIMS> then_bcast,
typename Eigen::array<Eigen::DenseIndex, NDIMS> else_bcast) {
output_tensor.device(d) = cond_tensor.broadcast(cond_bcast)
.select(then_tensor.broadcast(then_bcast),
else_tensor.broadcast(else_bcast));
}
};
#endif // TENSORFLOW_USE_SYCL
} // namespace functor
} // namespace tensorflow

6
tensorflow/core/kernels/cwise_ops.h
View File

@ -1217,7 +1217,11 @@ struct BCastSelectFunctor {
typename TTypes<T, NDIMS>::ConstTensor else_tensor,
typename Eigen::array<Eigen::DenseIndex, NDIMS> cond_bcast,
typename Eigen::array<Eigen::DenseIndex, NDIMS> then_bcast,
typename Eigen::array<Eigen::DenseIndex, NDIMS> else_bcast);
typename Eigen::array<Eigen::DenseIndex, NDIMS> else_bcast) {
output_tensor.device(d) = cond_tensor.broadcast(cond_bcast)
.select(then_tensor.broadcast(then_bcast),
else_tensor.broadcast(else_bcast));
}
};
} // end namespace functor