Searched refs:orig_arg_num (Results 1 – 1 of 1) sorted by relevance
792 const xla::OpSharding& sharding, int orig_arg_num, DataType dtype, in CreateOrGetSplitNodesForInputSharding() argument798 ShardedInputIndex input_index{replica_id, orig_arg_num}; in CreateOrGetSplitNodesForInputSharding()2817 int64 orig_arg_num = in CreateOrGetPerHostVariableCopy() local2819 if (arg_shardings[orig_arg_num].type() != xla::OpSharding::OTHER) { in CreateOrGetPerHostVariableCopy()2856 int64 orig_arg_num = i + params_info.NumPerReplicaArgs() + in CreateOrGetPerHostVariableCopy() local2872 CreatePerHostDummyArgs(arg_shapes[orig_arg_num], host_cpu_device, in CreateOrGetPerHostVariableCopy()3115 int64 orig_arg_num = core_arg_nums[core][i]; in BuildExecuteNodes() local3117 << " orig_arg_num " << orig_arg_num; in BuildExecuteNodes()3118 if (params_info.IsPerReplicaArg(orig_arg_num) || in BuildExecuteNodes()3119 params_info.IsDistributedArg(orig_arg_num)) { in BuildExecuteNodes()[all …]