This change started with an intention of adding an attribute to cast ops to decide

whether bfloat16 casts should use truncation or rounding.
This is a preparatory change before we switch the default float ==> bfloat16 cast
to use rounding instead of truncation. The attribute added can then be specified
on casts that rely on the truncation, e.g., the TensorFlow send/receive operations.
It later emerged that the choice of doing truncation is useful more generally.
Therefore, this change allows the new attribute to be used by all relevant casts
to use truncation instead of rounding.

PiperOrigin-RevId: 205996367

1 files changed, 2 insertions, 4 deletions

diff --git a/tensorflow/core/kernels/cast_op_impl_complex128.cc b/tensorflow/core/kernels/cast_op_impl_complex128.cc
index 52899d58cd..9a184e5954 100644
--- a/tensorflow/core/kernels/cast_op_impl_complex128.cc
+++ b/tensorflow/core/kernels/cast_op_impl_complex128.cc
@@ -20,15 +20,13 @@ namespace tensorflow {
 typedef Eigen::ThreadPoolDevice CPUDevice;
 typedef Eigen::GpuDevice GPUDevice;
 
-std::function<void(OpKernelContext*, const Tensor&, Tensor*)>
-GetCpuCastFromComplex128(DataType dst_dtype) {
+CastFunctorType GetCpuCastFromComplex128(DataType dst_dtype) {
   CURRY_TYPES3(CAST_CASE, CPUDevice, std::complex<double>);
   return nullptr;
 }
 
 #if GOOGLE_CUDA
-std::function<void(OpKernelContext*, const Tensor&, Tensor*)>
-GetGpuCastFromComplex128(DataType dst_dtype) {
+CastFunctorType GetGpuCastFromComplex128(DataType dst_dtype) {
   CURRY_TYPES3_NO_BF16(CAST_CASE, GPUDevice, std::complex<double>);
   return nullptr;
 }