Searched refs:output_sharding_config (Results 1 – 3 of 3) sorted by relevance
| /external/tensorflow/tensorflow/compiler/mlir/tensorflow/utils/ |
| D | xla_sharding_util.h | 55 const int core_id, llvm::ArrayRef<xla::OpSharding> output_sharding_config,
63 llvm::ArrayRef<xla::OpSharding> output_sharding_config,
|
| D | xla_sharding_util.cc | 375 llvm::ArrayRef<xla::OpSharding> output_sharding_config, const int core_id, in MapClusterOutputIndexWithRegionOutputIndex() argument
380 const auto& sharding = output_sharding_config[output_index]; in MapClusterOutputIndexWithRegionOutputIndex()
501 const int core_id, llvm::ArrayRef<xla::OpSharding> output_sharding_config, in GetOutputTypesForLogicalDeviceComputation() argument
508 const auto& output_sharding = output_sharding_config[output_index]; in GetOutputTypesForLogicalDeviceComputation()
538 llvm::ArrayRef<xla::OpSharding> output_sharding_config, in RemapOutputsFromLogicalDevices() argument
545 const auto& output_sharding = output_sharding_config[output_index]; in RemapOutputsFromLogicalDevices()
606 output_sharding_config, logical_device_id, output_index); in RemapOutputsFromLogicalDevices()
|
| /external/tensorflow/tensorflow/compiler/mlir/tensorflow/transforms/ |
| D | tpu_rewrite_pass.cc | 450 const int core_id, llvm::ArrayRef<xla::OpSharding> output_sharding_config, in BuildExecuteOp() argument
457 core_id, output_sharding_config, cluster_func, &output_types); in BuildExecuteOp()
471 llvm::ArrayRef<xla::OpSharding> output_sharding_config, in BuildParallelExecuteOp() argument
488 core, output_sharding_config, cluster_func, &output_types); in BuildParallelExecuteOp()
522 result = BuildExecuteOp(core, output_sharding_config, execute_inputs, in BuildParallelExecuteOp()
|