Searched refs:should_cast (Results 1 – 4 of 4) sorted by relevance
521 def _get_var_read_dtype(input_list, should_cast): argument523 if should_cast and input_list and input_list[0].dtype.is_floating:529 def autocast_context_manager(input_list, should_cast): argument544 var_read_dtype = _get_var_read_dtype(input_list, should_cast)
67 ShouldCastFunc should_cast = nullptr; member
197 const DataType cast_dtype = opts.should_cast ? opts.should_cast(edge) : dtype; in AddSend()260 if (opts.should_cast && !NeedSameDeviceSendRecv(edge, g_info)) { in AddRecv()261 cast_dtype = opts.should_cast(edge); in AddRecv()
1599 popts.should_cast = [enable_bfloat16_sendrecv](const Edge* e) { in BuildAndRegisterPartitions()