Searched refs:input_sharding (Results 1 – 3 of 3) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | hlo_sharding_util_test.cc | 58 HloSharding input_sharding = in TEST() local
62 ReshapeSharding(input_shape, output_shape, input_sharding); in TEST()
70 HloSharding input_sharding = HloSharding::Tile(Array2D<int64>({{0}, {1}})); in TEST() local
74 ReshapeSharding(input_shape, output_shape, input_sharding); in TEST()
84 HloSharding input_sharding = HloSharding::Tile(tile); in TEST() local
88 ReshapeSharding(input_shape, output_shape, input_sharding); in TEST()
96 HloSharding input_sharding = in TEST() local
101 ReshapeSharding(input_shape, output_shape, input_sharding); in TEST()
122 HloSharding input_sharding = in TEST() local
127 ReshapeSharding(input_shape, output_shape, input_sharding); in TEST()
[all …]
|
| /external/tensorflow/tensorflow/compiler/mlir/tensorflow/utils/ |
| D | xla_sharding_util.cc | 146 const mlir::Location& location, const xla::OpSharding& input_sharding, in HandleTileShardedInputs() argument
151 input_sharding.tile_assignment_devices_size()); in HandleTileShardedInputs()
157 const auto& dimension_splits = input_sharding.tile_assignment_dimensions(); in HandleTileShardedInputs()
195 tiled_inputs->reserve(input_sharding.tile_assignment_devices_size()); in HandleTileShardedInputs()
|
| /external/tensorflow/tensorflow/compiler/xla/service/spmd/ |
| D | spmd_partitioner.cc | 1661 const HloSharding& input_sharding = input->sharding(); in HandleSort() local
1663 input_sharding.tile_assignment().dim(sort_dim); in HandleSort()
1676 .Reshard(input_sharding) in HandleSort()
1690 input_sharding.GetTupleSharding(replicated_shape).ValueOrDie(); in HandleSort()
1706 slice_input->set_sharding(input_sharding); in HandleSort()
1716 slice_index->set_sharding(input_sharding); in HandleSort()
|