Searched refs:sharded_lhs_hlo (Results 1 – 3 of 3) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | dot_as_convolution_util.cc | 111 HloInstruction* sharded_lhs_hlo, HloInstruction* sharded_rhs_hlo) { in CreateShardedConvForDotGeneralConvolution() argument
117 wd->set_size(sharded_lhs_hlo->shape().dimensions( in CreateShardedConvForDotGeneralConvolution()
127 wd->set_size(sharded_lhs_hlo->shape().dimensions( in CreateShardedConvForDotGeneralConvolution()
143 sharded_lhs_hlo->shape(), sharded_rhs_hlo->shape(), in CreateShardedConvForDotGeneralConvolution()
149 sharded_conv_shape, sharded_lhs_hlo, sharded_rhs_hlo, in CreateShardedConvForDotGeneralConvolution()
|
| D | dot_as_convolution_util.h | 63 HloInstruction* sharded_lhs_hlo, HloInstruction* sharded_rhs_hlo);
|
| /external/tensorflow/tensorflow/compiler/xla/service/spmd/ |
| D | convolution_handler.cc | 896 HloInstruction* sharded_lhs_hlo, HloInstruction* sharded_rhs_hlo, in CreateShardedConvConvolution() argument
903 wd->set_size(sharded_lhs_hlo->shape().dimensions( in CreateShardedConvConvolution()
913 wd->set_size(sharded_lhs_hlo->shape().dimensions( in CreateShardedConvConvolution()
937 feature_group_count = sharded_lhs_hlo->shape().dimensions( in CreateShardedConvConvolution()
946 sharded_lhs_hlo->shape().dimensions(conv_dnums.input_batch_dimension()); in CreateShardedConvConvolution()
952 sharded_lhs_hlo->shape(), sharded_rhs_hlo->shape(), in CreateShardedConvConvolution()
957 sharded_conv_shape, sharded_lhs_hlo, sharded_rhs_hlo, feature_group_count, in CreateShardedConvConvolution()
|