Searched refs:input_sharding (Results 1 – 5 of 5) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | hlo_sharding_util_test.cc | 73 HloSharding input_sharding = in TEST() local
77 ReshapeSharding(input_shape, output_shape, input_sharding); in TEST()
85 HloSharding input_sharding = HloSharding::Tile(Array2D<int64_t>({{0}, {1}})); in TEST() local
89 ReshapeSharding(input_shape, output_shape, input_sharding); in TEST()
99 HloSharding input_sharding = HloSharding::Tile(tile); in TEST() local
103 ReshapeSharding(input_shape, output_shape, input_sharding); in TEST()
111 HloSharding input_sharding = in TEST() local
116 ReshapeSharding(input_shape, output_shape, input_sharding); in TEST()
137 HloSharding input_sharding = in TEST() local
142 ReshapeSharding(input_shape, output_shape, input_sharding); in TEST()
[all …]
|
| /external/tensorflow/tensorflow/dtensor/mlir/ |
| D | set_default_sharding.cc | 40 llvm::SmallVector<llvm::StringRef, 4> input_sharding(cluster.getNumOperands(), in SetDefaultSharding() local
46 builder->getStrArrayAttr(input_sharding)); in SetDefaultSharding()
|
| /external/tensorflow/tensorflow/dtensor/cc/ |
| D | dtensor_meta_ops.cc | 133 const std::vector<int32> input_sharding = input_layout.num_shards(); in __anon82f11ae70202() local
144 TF_RETURN_IF_ERROR(c->Multiply(dim, input_sharding[i], &out_dim)); in __anon82f11ae70202()
|
| /external/tensorflow/tensorflow/compiler/mlir/tensorflow/utils/ |
| D | xla_sharding_util.cc | 140 const mlir::Location& location, const xla::OpSharding& input_sharding, in HandleTileShardedInputs() argument
145 input_sharding.tile_assignment_devices_size()); in HandleTileShardedInputs()
151 const auto& dimension_splits = input_sharding.tile_assignment_dimensions(); in HandleTileShardedInputs()
189 tiled_inputs->reserve(input_sharding.tile_assignment_devices_size()); in HandleTileShardedInputs()
|
| /external/tensorflow/tensorflow/compiler/xla/service/spmd/ |
| D | spmd_partitioner.cc | 2471 const HloSharding& input_sharding = input->sharding(); in HandleSort() local
2473 input_sharding.tile_assignment().dim(sort_dim); in HandleSort()
2485 .Reshard(input_sharding) in HandleSort()
2499 input_sharding.GetTupleSharding(replicated_shape).ValueOrDie(); in HandleSort()
2515 slice_input->set_sharding(input_sharding); in HandleSort()
2525 slice_index->set_sharding(input_sharding); in HandleSort()
|