Searched refs:cast_dtype (Results 1 – 2 of 2) sorted by relevance
201 const DataType cast_dtype = opts.should_cast ? opts.should_cast(edge) : dtype; in AddSend() local215 if (dtype != cast_dtype && !NeedSameDeviceSendRecv(edge, g_info)) { in AddSend()223 cast_builder.Attr("DstT", cast_dtype); in AddSend()225 if (cast_dtype == DT_BFLOAT16) { in AddSend()237 send_from.Reset(cast->name(), 0, cast_dtype); in AddSend()261 DataType cast_dtype = dtype; in AddRecv() local265 cast_dtype = opts.should_cast(edge); in AddRecv()305 .Attr("tensor_type", cast_dtype); in AddRecv()312 if (dtype != cast_dtype) { in AddRecv()318 .Input(recv->name(), 0, cast_dtype); in AddRecv()[all …]
723 cast_dtype = dtypes.as_dtype(tensor.op.get_attr("DstT"))724 return pre_cast.astype(cast_dtype.as_numpy_dtype)930 cast_dtype = dtypes.as_dtype(tensor.op.get_attr("DstT"))931 if cast_dtype not in (dtypes.int32, dtypes.int64):935 cast_dtype.as_numpy_dtype)