Home
last modified time | relevance | path

Searched refs:conv_spatial_dims (Results 1 – 6 of 6) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/service/
Ddot_as_convolution_util.h47 std::vector<DimNums> conv_spatial_dims; member
Ddot_as_convolution_util.cc98 dims.conv_spatial_dims.push_back({lhs, rhs, output, i}); in ParseConvolutionDimsInfo()
101 dims.conv_spatial_dims.push_back({lhs, rhs, output, i}); in ParseConvolutionDimsInfo()
Dsharding_propagation.cc507 instruction->operand(0), dot_dims.conv_spatial_dims, 0); in InferConvolutionShardingFromOperands()
509 instruction->operand(1), dot_dims.conv_spatial_dims, 1); in InferConvolutionShardingFromOperands()
510 if (dot_dims.conv_spatial_dims.empty() || in InferConvolutionShardingFromOperands()
1204 if (dot_dims.conv_spatial_dims.empty()) { in GetShardingFromUser()
/external/tensorflow/tensorflow/compiler/xla/service/spmd/
Dconvolution_handler.cc927 for (const auto& dim : dot_dnums.conv_spatial_dims) { in CreateShardedConvConvolution()
1018 for (const auto& dims : dims_info.conv_spatial_dims) { in HandleConvolution()
1019 mapping.conv_spatial_dims.emplace_back(); in HandleConvolution()
1020 mapping.conv_spatial_dims.back().lhs = dims.lhs; in HandleConvolution()
1021 mapping.conv_spatial_dims.back().rhs = dims.rhs; in HandleConvolution()
1022 mapping.conv_spatial_dims.back().output = dims.output; in HandleConvolution()
1023 mapping.conv_spatial_dims.back().spatial = dims.spatial_dim; in HandleConvolution()
1029 if (dims_info.conv_spatial_dims.empty()) { in HandleConvolution()
Ddot_handler.cc375 for (const auto& mapping : dims_mapping.conv_spatial_dims) { in ComputeDimensionIndexMapping()
2252 new_dims_mapping.conv_spatial_dims = dims_mapping.conv_spatial_dims; in ConvertDimsMappingWithFeatureGroupCount()
2275 new_dims_mapping.conv_spatial_dims = dims_mapping.conv_spatial_dims; in ConvertDimsMappingWithBatchGroupCount()
2491 lhs.sharding(), dims_mapping.conv_spatial_dims, 0); in PartitionDot()
2493 rhs.sharding(), dims_mapping.conv_spatial_dims, 1); in PartitionDot()
2495 output_sharding, dims_mapping.conv_spatial_dims, 2); in PartitionDot()
2508 if ((!dims_mapping.conv_spatial_dims.empty() && in PartitionDot()
Dspmd_partitioner.h397 std::vector<DimsMapping> conv_spatial_dims; member