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, 3 insertions, 1 deletions

diff --git a/tensorflow/core/kernels/cast_op.cc b/tensorflow/core/kernels/cast_op.cc
index b4c97df38b..0478c93280 100644
--- a/tensorflow/core/kernels/cast_op.cc
+++ b/tensorflow/core/kernels/cast_op.cc
@@ -59,6 +59,8 @@ CastOpBase::CastOpBase(OpKernelConstruction* ctx) : OpKernel(ctx) {
 
   OP_REQUIRES_OK(ctx, ctx->GetAttr("DstT", &external_dst_dtype_));
 
+  OP_REQUIRES_OK(ctx, ctx->GetAttr("Truncate", &use_truncation_));
+
   // Quantized data types use the same underlying format as their non quantized
   // version so we use the non quantized implementation for casting.
   if (external_dst_dtype_ == DT_QUINT8) {
@@ -100,7 +102,7 @@ void CastOpBase::Compute(OpKernelContext* ctx) {
     Tensor* out = nullptr;
     OP_REQUIRES_OK(ctx, ctx->allocate_output(0, in.shape(), &out));
     out->set_dtype(dst_dtype_);
-    work_(ctx, in, out);
+    work_(ctx, in, out, use_truncation_);
     out->set_dtype(external_dst_dtype_);
   }
 }