Home
last modified time | relevance | path

Searched refs:sharded_input (Results 1 – 5 of 5) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/service/spmd/
Dspmd_partitioner.cc759 auto sharded_input = in ReshardAsWindowedInput() local
764 sharded_input, shard_window, in ReshardAsWindowedInput()
1609 const Shape& operand_shape = reshard_operand->sharded_input->shape(); in HandleSlice()
1629 shard_shape, reshard_operand->sharded_input, start_indices, in HandleSlice()
1632 auto data = reshard_operand->sharded_input; in HandleSlice()
2080 reshard_operand->sharded_input->shape().dimensions(input_sharded_dim), in HandleReshape()
2085 output_shard_shape, reshard_operand->sharded_input)); in HandleReshape()
2135 reshard_output->sharded_input->shape().dimensions(output_sharded_dim), in HandleReshape()
2137 SetPartitionedHlo(hlo, [&] { return reshard_output->sharded_input; }); in HandleReshape()
2663 auto sharded_pad = reshard_operand->sharded_input; in HandlePad()
[all …]
Dspmd_partitioner_test.cc862 auto sharded_input = in TEST_F() local
865 op::CollectivePermute(op::Slice(sharded_input))); in TEST_F()
869 op::Pad(op::Concatenate(sharded_input, right_halo), op::Constant())), in TEST_F()
941 auto sharded_input = in TEST_F() local
944 op::CollectivePermute(op::Slice(sharded_input))); in TEST_F()
948 op::Pad(op::Concatenate(sharded_input, right_halo), op::Constant())), in TEST_F()
983 auto sharded_input = in TEST_F() local
986 op::CollectivePermute(op::Slice(sharded_input))); in TEST_F()
988 op::CollectivePermute(op::Slice(sharded_input))); in TEST_F()
993 op::Pad(op::Concatenate(left_halo, sharded_input, right_halo), in TEST_F()
[all …]
Dspmd_partitioner.h264 HloInstruction* sharded_input; member
Dconvolution_handler.cc800 create_sharded_conv(resharded_operand_and_window->sharded_input, in PartitionConvolutionTiledOutput()
Dspmd_partitioner_util.cc1106 return reshard_window->sharded_input; in HaloExchangeToPadOnLeft()