Searched defs:arg_shardings (Results 1 – 3 of 3) sorted by relevance
| /external/tensorflow/tensorflow/compiler/tf2xla/ |
| D | xla_compiler.cc | 102 std::map<int, xla::OpSharding> arg_shardings; in ComputeArgAndRetvalShardings() local 169 const std::map<int, xla::OpSharding>& arg_shardings, in BuildComputation() 982 const std::map<int, xla::OpSharding>& arg_shardings, in BuildArguments() 1356 std::map<int, xla::OpSharding> arg_shardings; in CompileGraph() local
|
| /external/tensorflow/tensorflow/compiler/mlir/xla/ |
| D | mlir_hlo_to_hlo.cc | 538 llvm::SmallVectorImpl<std::optional<xla::OpSharding>>* arg_shardings, in ExtractShardingsFromFunction() 2271 llvm::SmallVector<std::optional<xla::OpSharding>, 4> arg_shardings; in RunOnFunction() local 2358 llvm::ArrayRef<std::optional<xla::OpSharding>> arg_shardings, in SetEntryTupleShardings() 2389 llvm::ArrayRef<std::optional<xla::OpSharding>> arg_shardings, in LowerBasicBlockAsFunction()
|
| /external/tensorflow/tensorflow/core/tpu/graph_rewrite/ |
| D | distributed_tpu_rewrite_pass.cc | 3155 const std::vector<xla::OpSharding>& arg_shardings, in CreateOrGetPerHostVariableCopy() 3245 const std::vector<xla::OpSharding>& arg_shardings, in BuildExecuteNodes()
|